Add jsrun launcher based on neox 2.0

Author: Quentin-Anthony

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
 

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.
     """

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