Fix bugs for PyTorchDistributedRuntime (#271)

Author: qiuosier

ads/jobs/builders/runtimes/pytorch_runtime.py

@@ -206,7 +206,7 @@ def run(self, dsc_job, **kwargs):
                     envs = {}
                 # Huggingface accelerate requires machine rank
                 envs["RANK"] = str(i)
-                envs["WORLD_SIZE"] = replicas
+                envs["WORLD_SIZE"] = str(replicas)
                 if main_run:
                     envs["MAIN_JOB_RUN_OCID"] = main_run.id
                 name = replica_kwargs.get("display_name")

ads/jobs/templates/driver_pytorch.py

@@ -270,6 +270,7 @@ def find_self_ip(self):
                 ip = snics[0].address
                 if ipaddress.ip_address(ip) in ipaddress.ip_network(cidr):
                     logger.info("Node IP address: %s", ip)
+                    # Specify the network interface for NCCL/GLOO
                     os.environ["GLOO_SOCKET_IFNAME"] = interface
                     os.environ["NCCL_SOCKET_IFNAME"] = interface
                     return ip
@@ -579,13 +580,22 @@ def save_deepspeed_env(self):
                 # as the .deepspeed_env file is parsed line by line.
                 if not v or "\n" in v:
                     continue
+                # Ignore variables that are node specific
+                # The network interface name for each job run is a unique string, e.g. ens300f0v1604
+                if k in ["NCCL_SOCKET_IFNAME", "GLOO_SOCKET_IFNAME", "JOB_RUN_OCID"]:
+                    continue
                 # Quote the value if it contains space
                 # Environment variable containing space may not be exported correctly when using pdsh
                 # https://github.com/microsoft/DeepSpeed/blob/v0.9.2/deepspeed/launcher/multinode_runner.py#L79
                 if " " in v:
                     v = shlex.quote(v)
 
                 f.write(f"{k}={v}\n")
+            # The following are required for specifying the network interface to be used by NCCL/GLOO
+            # The value should be the prefix of the expected network interface name
+            # https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-socket-ifname
+            f.write("NCCL_SOCKET_IFNAME=ens\n")
+            f.write("GLOO_SOCKET_IFNAME=ens\n")
         logger.debug("Environment variables saved to %s", self.ENV_FILE)
         self.run_command(f"cat {self.ENV_FILE}")
 

tests/unitary/default_setup/jobs/test_jobs_pytorch_ddp.py

@@ -137,13 +137,13 @@ def test_create_job_runs(self, patched_run, *args):
             [
                 {
                     "display_name": "None-0",
-                    "environment_variables": {"RANK": "0", "WORLD_SIZE": 2},
+                    "environment_variables": {"RANK": "0", "WORLD_SIZE": "2"},
                 },
                 {
                     "display_name": "None-1",
                     "environment_variables": {
                         "RANK": "1",
-                        "WORLD_SIZE": 2,
+                        "WORLD_SIZE": "2",
                         "MAIN_JOB_RUN_OCID": test_ocid,
                     },
                 },