fixed case when ntasks_per_node is used instead (EleutherAI#1069)

megatron/neox_arguments/arguments.py

@@ -732,7 +732,8 @@ def configure_distributed_args(self):
  if self.deepspeed_slurm:
  os.environ["LOCAL_RANK"] = os.environ["SLURM_LOCALID"]
  os.environ["RANK"] = os.environ["SLURM_PROCID"]
- os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"]
+ os.environ["WORLD_SIZE"] = os.environ["SLURM_NTASKS"] if os.environ.get("SLURM_NTASKS") is not None \
+ else str(int(os.environ["SLURM_NNODES"]) * int(os.environ["SLURM_NTASKS_PER_NODE"]))
 
  self.update_value("local_rank", int(os.getenv("LOCAL_RANK", "0")))
  self.update_value("rank", int(os.getenv("RANK", "0")))