Add jsrun launcher based on neox 2.0

EleutherAI · Mar 18, 2023 · 3782c7a · 3782c7a
1 parent e2d8a24
commit 3782c7a
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/configs/neox_arguments.md b/configs/neox_arguments.md
@@ -1809,7 +1809,7 @@ Args for deepspeed runner (deepspeed.launcher.runner).
 
 
 
-- **launcher**: typing.Literal['pdsh', 'openmpi', 'mvapich', 'slurm']
+- **launcher**: typing.Literal['pdsh', 'openmpi', 'mvapich', 'slurm', 'jsrun']
 
  Default = pdsh
 

diff --git a/megatron/neox_arguments/deepspeed_args.py b/megatron/neox_arguments/deepspeed_args.py
@@ -192,7 +192,7 @@ class NeoXArgsDeepspeedRunner(NeoXArgsTemplate):
  IP address of node 0, will be inferred via 'hostname -I' if not specified.
  """
 
- launcher: Literal["pdsh", "openmpi", "mvapich", "slurm"] = "pdsh"
+ launcher: Literal["pdsh", "openmpi", "mvapich", "slurm", "jsrun"] = "pdsh"
  """
  Launcher backend for multi-node training. Options currently include PDSH, OpenMPI, MVAPICH.
  """

diff --git a/megatron/neox_arguments/neox_args.py b/megatron/neox_arguments/neox_args.py
@@ -612,6 +612,11 @@ class NeoXArgsOther(NeoXArgsTemplate):
  Run via SLURM, this will attempt to discover the necessary variables to initialize torch distributed from the SLURM environment
  """
 
+ deepspeed_jsrun: bool = False
+ """
+ Run via JSRUN, this will attempt to discover the necessary variables to initialize torch distributed from the IBM LSF environment
+ """
+
  user_script: str = None
  """
  user script to be run