Merge remote-tracking branch 'determined/eleuther_dai' into igor

Author: marcmac

deepspeed/launcher/launch.py

@@ -2,7 +2,6 @@
 """
 DeepSpeed launcher, this is similar to torch.distributed.launch but supports
 additional features such as abitrary gpu exclusion.
-
 deepspeed.launcher.launch is intended to be run on a single worker node and
 will spawn several worker sub-processes depending on how many devices/ranks
 are on the worker.
@@ -53,6 +52,22 @@ def parse_args():
     parser.add_argument("--detect_nvlink_pairs", action="store_true",
                         help="autodetects nvlink pairs and remaps CUDA_VISIBLE_DEVICES along the fastest connections")
 
+    parser.add_argument("--module",
+                        action="store_true",
+                        help="Change each process to interpret the launch "
+                        "script as a Python module, executing with the same "
+                        "behavior as 'python -m'.")
+
+    parser.add_argument("--no_python",
+                        action="store_true",
+                        help="Skip prepending the training script with "
+                        "'python' - just execute it directly.")
+
+    parser.add_argument("--no_local_rank",
+                        action="store_true",
+                        help="Do not pass local_rank as an argument when calling "
+                        "the user's training script.")
+
     # positional
     parser.add_argument("training_script",
                         type=str,
@@ -117,6 +132,9 @@ def main():
     current_env["MASTER_ADDR"] = args.master_addr
     current_env["MASTER_PORT"] = str(args.master_port)
     current_env["WORLD_SIZE"] = str(dist_world_size)
+    current_env["CROSS_RANK"] = str(args.node_rank)
+    current_env["CROSS_SIZE"] = str(args.nnodes)
+    current_env["LOCAL_SIZE"] = str(num_local_procs)
 
     processes = []
     for local_rank in range(0, num_local_procs):
@@ -126,12 +144,20 @@ def main():
         current_env["LOCAL_RANK"] = str(local_rank)
 
         # spawn the processes
-        cmd = [
-            sys.executable,
-            "-u",
-            args.training_script,
-            "--local_rank={}".format(local_rank)
-        ] + args.training_script_args
+        cmd = []
+        if not args.no_python:
+            cmd = [sys.executable, "-u"]
+            if args.module:
+                cmd.append("-m")
+        else:
+            if args.module:
+                raise ValueError("Don't use both the '--no_python' flag"
+                                 " and the '--module' flag at the same time.")
+        cmd.append(args.training_script)
+        # A user may not want to pass local_rank as a keyword arg so we make this optional.
+        if not args.no_local_rank:
+            cmd.append(f"--local_rank={local_rank}")
+        cmd += args.training_script_args
 
         sig_names = {2: "SIGINT", 15: "SIGTERM"}
         last_return_code = None

deepspeed/launcher/multinode_runner.py

@@ -4,6 +4,7 @@
 import subprocess
 import warnings
 from abc import ABC, abstractmethod
+from shlex import quote
 
 from ..utils import logger
 from .constants import PDSH_MAX_FAN_OUT, MVAPICH_TMP_HOSTFILE
@@ -56,7 +57,7 @@ def get_cmd(self, environment, active_resources):
 
         exports = ""
         for key, val in self.exports.items():
-            exports += "export {}={}; ".format(key, val)
+            exports += f"export {key}={quote(val)}; "
 
         deepspeed_launch = [
             exports,
@@ -71,7 +72,13 @@ def get_cmd(self, environment, active_resources):
             "--master_port={}".format(self.args.master_port)
         ]
         if self.args.detect_nvlink_pairs:
-            deepspeed_launch += ["--detect_nvlink_pairs"]
+            deepspeed_launch.append("--detect_nvlink_pairs")
+        if self.args.no_python:
+            deepspeed_launch.append("--no_python")
+        if self.args.module:
+            deepspeed_launch.append("--module")
+        if self.args.no_local_rank:
+            deepspeed_launch.append("--no_local_rank")
 
         return pdsh_cmd_args + deepspeed_launch + [self.user_script
                                                    ] + self.user_arguments

deepspeed/launcher/runner.py

@@ -24,7 +24,7 @@
 from ..utils import logger
 
 DLTS_HOSTFILE = "/job/hostfile"
-EXPORT_ENVS = ["NCCL", "PYTHON", "MV2", 'UCX']
+EXPORT_ENVS = ["NCCL", "PYTHON", "MV2", "UCX"]
 DEEPSPEED_ENVIRONMENT_NAME = ".deepspeed_env"
 DEEPSPEED_ENVIRONMENT_PATHS = [os.path.expanduser("~"), '.']
 PDSH_MAX_FAN_OUT = 1024
@@ -104,6 +104,22 @@ def parse_args(args=None):
                         help="(optional) pass launcher specific arguments as a "
                         "single quoted argument.")
 
+    parser.add_argument("--module",
+                        action="store_true",
+                        help="Change each process to interpret the launch "
+                        "script as a Python module, executing with the same "
+                        "behavior as 'python -m'.")
+
+    parser.add_argument("--no_python",
+                        action="store_true",
+                        help="Skip prepending the training script with "
+                        "'python' - just execute it directly.")
+
+    parser.add_argument("--no_local_rank",
+                        action="store_true",
+                        help="Do not pass local_rank as an argument when calling "
+                        "the user's training script.")
+
     parser.add_argument("--force_multi",
                         action="store_true",
                         help="Force multi-node launcher mode, helps in cases where user "
@@ -154,11 +170,9 @@ def fetch_hostfile(hostfile_path):
 
 def parse_resource_filter(host_info, include_str="", exclude_str=""):
     '''Parse an inclusion or exclusion string and filter a hostfile dictionary.
-
     String format is NODE_SPEC[@NODE_SPEC ...], where
         NODE_SPEC = NAME[:SLOT[,SLOT ...]].
     If :SLOT is omitted, include/exclude all slots on that host.
-
     Examples:
         include_str="worker-0@worker-1:0,2" will use all slots on worker-0 and
           slots [0, 2] on worker-1.
@@ -326,7 +340,13 @@ def main(args=None):
             "--master_port={}".format(args.master_port)
         ]
         if args.detect_nvlink_pairs:
-            deepspeed_launch += ["--detect_nvlink_pairs"]
+            deepspeed_launch.append("--detect_nvlink_pairs")
+        if args.no_python:
+            deepspeed_launch.append("--no_python")
+        if args.module:
+            deepspeed_launch.append("--module")
+        if args.no_local_rank:
+            deepspeed_launch.append("--no_local_rank")
         cmd = deepspeed_launch + [args.user_script] + args.user_args
     else:
         args.launcher = args.launcher.lower()
@@ -358,7 +378,7 @@ def main(args=None):
             if os.path.isfile(environ_file):
                 with open(environ_file, 'r') as fd:
                     for var in fd.readlines():
-                        key, val = var.split('=')
+                        key, val = var.split('=', maxsplit=1)
                         runner.add_export(key, val)
 
         cmd = runner.get_cmd(env, active_resources)