fixed case when ntasks_per_node is used instead (#1069)

EleutherAI · Nov 1, 2023 · 41f019e · 41f019e
1 parent a10f69c
commit 41f019e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/megatron/neox_arguments/arguments.py b/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")))