Support deepspeed with accelerate launch.

Author: qiuosier

ads/jobs/builders/infrastructure/dsc_job_runtime.py

@@ -1045,6 +1045,7 @@ class PyTorchDistributedRuntimeHandler(PythonRuntimeHandler):
     RUNTIME_CLASS = PyTorchDistributedRuntime
     CONST_WORKER_COUNT = "OCI__WORKER_COUNT"
     CONST_COMMAND = "OCI__LAUNCH_CMD"
+    CONST_DEEPSPEED = "OCI__DEEPSPEED"
 
     GIT_SPEC_MAPPINGS = {
         cluster_config_helper.OCI__RUNTIME_URI: GitPythonRuntime.CONST_GIT_URL,
@@ -1054,14 +1055,19 @@ class PyTorchDistributedRuntimeHandler(PythonRuntimeHandler):
     }
 
     SPEC_MAPPINGS = PythonRuntimeHandler.SPEC_MAPPINGS
-    SPEC_MAPPINGS.update({PyTorchDistributedRuntime.CONST_COMMAND: CONST_COMMAND})
+    SPEC_MAPPINGS.update(
+        {
+            PyTorchDistributedRuntime.CONST_COMMAND: CONST_COMMAND,
+        }
+    )
 
     def _translate_artifact(self, runtime: PyTorchDistributedRuntime):
         return PyTorchDistributedArtifact(runtime.source_uri, runtime)
 
     def _translate_env(self, runtime: PyTorchDistributedRuntime) -> dict:
         envs = super()._translate_env(runtime)
         replica = runtime.replica if runtime.replica else 1
+        # WORKER_COUNT = REPLICA - 1
         envs[self.CONST_WORKER_COUNT] = str(replica - 1)
         envs[self.CONST_JOB_ENTRYPOINT] = PyTorchDistributedArtifact.CONST_DRIVER_SCRIPT
         if runtime.inputs:
@@ -1080,6 +1086,8 @@ def _translate_env(self, runtime: PyTorchDistributedRuntime) -> dict:
                 envs[driver_utils.CONST_ENV_PIP_REQ] = runtime.dependencies[
                     PyTorchDistributedRuntime.CONST_PIP_REQ
                 ]
+        if runtime.use_deepspeed:
+            envs[self.CONST_DEEPSPEED] = "1"
         return envs
 
     def _extract_envs(self, dsc_job) -> dict:
@@ -1117,6 +1125,8 @@ def _extract_envs(self, dsc_job) -> dict:
             )
         if dep:
             spec[PyTorchDistributedRuntime.CONST_DEP] = dep
+        if envs.pop(self.CONST_DEEPSPEED, None):
+            spec[PyTorchDistributedRuntime.CONST_DEEPSPEED] = True
         # Envs
         if envs:
             spec[PythonRuntime.CONST_ENV_VAR] = envs

ads/jobs/builders/runtimes/pytorch_runtime.py

@@ -14,6 +14,7 @@ class PyTorchDistributedRuntime(PythonRuntime):
     CONST_PIP_REQ = "pipRequirements"
     CONST_PIP_PKG = "pipPackages"
     CONST_COMMAND = "command"
+    CONST_DEEPSPEED = "deepspeed"
 
     def with_git(
         self, url: str, branch: str = None, commit: str = None, secret_ocid: str = None
@@ -77,13 +78,21 @@ def with_dependency(self, pip_req=None, pip_pkg=None):
     def dependencies(self) -> dict:
         return self.get_spec(self.CONST_DEP)
 
-    def with_command(self, command: str):
+    def with_command(self, command: str, use_deepspeed=False):
+        if use_deepspeed:
+            self.set_spec(self.CONST_DEEPSPEED, True)
         return self.set_spec(self.CONST_COMMAND, command)
 
     @property
     def command(self):
         return self.get_spec(self.CONST_COMMAND)
 
+    @property
+    def use_deepspeed(self):
+        if self.get_spec(self.CONST_DEEPSPEED):
+            return True
+        return False
+
     def run(self, dsc_job, **kwargs):
         replicas = self.replica if self.replica else 1
         main_run = None

ads/jobs/templates/driver_pytorch.py

@@ -43,7 +43,7 @@
 CONST_ENV_HOST_JOB_RUN_OCID = "MAIN_JOB_RUN_OCID"
 CONST_ENV_LD_PRELOAD = "LD_PRELOAD"
 CONST_ENV_LAUNCH_CMD = "OCI__LAUNCH_CMD"
-CONST_ENV_LAUNCH_ARGS = "OCI__LAUNCH_ARGS"
+CONST_ENV_DEEPSPEED = "OCI__DEEPSPEED"
 LOG_PREFIX_HOST_IP = "Distributed Training HOST IP: "
 LOG_PREFIX_NODE_IP = "Node IP: "
 LOG_PREFIX_PUBLIC_KEY = "HOST PUBLIC KEY: "
@@ -246,6 +246,8 @@ def get_entrypoint_with_args(self, prefix=""):
         return cmd
 
     def prepare_cmd(self, launch_args, prefix=None):
+        if not launch_args:
+            launch_args = []
         # Append launch cmd args specified by the user.
         if self.launch_cmd:
             launch_args.append(self.launch_cmd[len(self.LAUNCHER) + 1 :])
@@ -486,67 +488,90 @@ def save_deepspeed_env(self):
         logger.debug("Environment variables saved to %s", self.ENV_FILE)
         self.run_command(f"cat {self.ENV_FILE}")
 
-    def run(self):
-        if self.is_host:
-            self.generate_key_pair().generate_hostfile()
-            self.save_deepspeed_env()
-            # Wait for nodes to be ready
-            for run in self.node_runs:
-                self.wait_for_log(run, LOG_PREFIX_PUBLIC_KEY)
+    def run_deepspeed_host(self, launch_args=None):
+        """Prepares the host and launch the deepspeed training.
 
-            for node_ip in self.node_ip_list:
-                self.run_command(
-                    f"ssh-keyscan -H {node_ip} >> {SSH_DIR}/known_hosts",
-                    level=logging.DEBUG,
-                    check=True,
-                )
-            # For DeepSpeed, we only need to run the cmd on the host
-            launch_args = [f"--hostfile={self.HOST_FILE}"]
+        Parameters
+        ----------
+        launch_args : str, optional
+            Additional command line arguments, by default None.
+            The deepspeed host file should be specified in the launch args.
+            For "deepspeed": --hostfile
+            For "accelerate launch": --deepspeed_hostfile
+        """
+        self.generate_key_pair().generate_hostfile()
+        self.save_deepspeed_env()
+        # Wait for nodes to be ready
+        for run in self.node_runs:
+            self.wait_for_log(run, LOG_PREFIX_PUBLIC_KEY)
+
+        for node_ip in self.node_ip_list:
+            self.run_command(
+                f"ssh-keyscan -H {node_ip} >> {SSH_DIR}/known_hosts",
+                level=logging.DEBUG,
+                check=True,
+            )
 
-            cmd = self.prepare_cmd(launch_args)
+        cmd = self.prepare_cmd(launch_args)
+        # For DeepSpeed, we only need to run the cmd on the host
+        try:
+            self.time_cmd(cmd)
+        except:
+            # Caution: file will not be generated if job run is killed from the console.
+            self.touch_file(self.ERROR_FILE)
+            raise
+        # Signal stop
+        self.touch_file(self.STOP_FILE)
+
+    def run_deepspeed_worker(self):
+        self.fetch_host_public_key()
+        # Keep the job run alive until host job run is finished.
+        while not os.path.exists(self.STOP_FILE):
+            time.sleep(60)
+            # Stop the node if the host touched the error file.
+            if os.path.exists(self.ERROR_FILE):
+                logger.error("There is an error in the host job run.")
+                sys.exit(1)
+            # Stop the node if the host job run is CANCELLED or in unexpected state.
+            self.host_job_run.sync()
+            if self.host_job_run.status not in [
+                "ACCEPTED",
+                "IN_PROGRESS",
+                "SUCCEEDED",
+            ]:
+                logger.info(
+                    "Host job run status is %s. Stopping job run...",
+                    self.host_job_run.status,
+                )
+                sys.exit(2)
+        logger.info("Job finished successfully. Stopping job run...")
 
-            try:
-                self.time_cmd(cmd)
-            except:
-                # Caution: file will not be generated if job run is killed from the console.
-                self.touch_file(self.ERROR_FILE)
-                raise
-            # Signal stop
-            self.touch_file(self.STOP_FILE)
+    def run(self):
+        if self.is_host:
+            launch_args = [f"--hostfile={self.HOST_FILE}"]
+            self.run_deepspeed_host(launch_args)
         else:
-            self.fetch_host_public_key()
-            # Keep the job run alive until host job run is finished.
-            while not os.path.exists(self.STOP_FILE):
-                time.sleep(60)
-                # Stop the node if the host touched the error file.
-                if os.path.exists(self.ERROR_FILE):
-                    logger.error("There is an error in the host job run.")
-                    sys.exit(1)
-                # Stop the node if the host job run is CANCELLED or in unexpected state.
-                self.host_job_run.sync()
-                if self.host_job_run.status not in [
-                    "ACCEPTED",
-                    "IN_PROGRESS",
-                    "SUCCEEDED",
-                ]:
-                    logger.info(
-                        "Host job run status is %s. Stopping job run...",
-                        self.host_job_run.status,
-                    )
-                    sys.exit(2)
-            logger.info("Job finished successfully. Stopping job run...")
+            self.run_deepspeed_worker()
 
 
 class AccelerateRunner(TorchRunner, DeepSpeedRunner):
-    DEFAULT_ARGS = ["num_processes", "num_machines", "machine_rank"]
-    TORCHRUN_ARGS = ["main_process_ip", "main_process_port"]
+    # accelerate launch will add main_process_port for deepspeed cmd even if it is not needed.
+    # https://github.com/huggingface/accelerate/blob/70920895e80f78d96d8f91e0beeb3ebdb8e5e5d6/src/accelerate/utils/launch.py#L233
+    DEFAULT_ARGS = [
+        "num_processes",
+        "num_machines",
+        "machine_rank",
+        "main_process_ip",
+        "main_process_port",
+    ]
+    TORCHRUN_ARGS = []
     LAUNCHER = "accelerate launch"
 
     def __init__(self, code_dir: str = driver_utils.DEFAULT_CODE_DIR) -> None:
         super().__init__(code_dir)
         # For "accelerate launch", only one of the following options can be used at one time
         # `--cpu`, `--multi_gpu`, `--tpu`, `--use_deepspeed`, `--use_fsdp`.
-        # When a config file is not provided, 
+        # When a config file is not provided,
         # --multi_gpu will be set automatically if there is more than 1 GPU
         # self.multi_gpu = bool(self.node_count > 1 or self.gpu_count > 1)
         self.num_machines = self.node_count
@@ -560,7 +585,9 @@ def __init__(self, code_dir: str = driver_utils.DEFAULT_CODE_DIR) -> None:
         self.main_process_ip = None
 
     def use_deepspeed(self):
-        return self.launch_cmd_contains("use_deepspeed")
+        return os.environ.get(CONST_ENV_DEEPSPEED) or self.launch_cmd_contains(
+            "use_deepspeed"
+        )
 
     def accelerate_args(self):
         args = []
@@ -584,13 +611,18 @@ def run_with_torchrun(self):
         self.time_cmd(cmd=cmd)
 
     def run_with_deepspeed(self):
-        raise NotImplementedError
+        if self.is_host:
+            launch_args = self.accelerate_args()
+            launch_args.append(f"--deepspeed_hostfile={self.HOST_FILE}")
+            self.run_deepspeed_host(launch_args)
+        else:
+            self.run_deepspeed_worker()
 
     def run(self):
         # Check if any default argument is provided by the user
         for arg in self.DEFAULT_ARGS:
             if self.launch_cmd_contains(arg):
-                logger.debug("%s found in launch args.", arg)
+                logger.debug("%s found in command.", arg)
                 setattr(self, arg, None)
         if self.use_deepspeed():
             self.run_with_deepspeed()