Add jsrun launcher based on neox 2.0

Quentin-Anthony · Quentin-Anthony · commit 3782c7ae60f8 · 2023-03-18T14:33:43.000-04:00
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
