[NPU] load EXPORT_ENV based on different accelerators to support mult…

…i-node training on other devices (#4830)

Different hardwares may require different environment variables. To
support multi-node training feature for NPU and other devices that rely
on different env vars, I add a method `export_envs()` to each
accelerator and load them in runner.py

For works about NPU, see #4567

Co-authored-by: Olatunji Ruwase <[email protected]>

accelerator/abstract_accelerator.py

@@ -258,3 +258,7 @@ def get_op_builder(self, class_name):
     @abc.abstractmethod
     def build_extension(self):
         ...
+
+    @abc.abstractmethod
+    def export_envs(self):
+        ...

accelerator/cpu_accelerator.py

@@ -280,3 +280,6 @@ def get_op_builder(self, class_name):
     def build_extension(self):
         from torch.utils.cpp_extension import BuildExtension
         return BuildExtension
+
+    def export_envs(self):
+        return []

accelerator/cuda_accelerator.py

@@ -322,3 +322,6 @@ def get_op_builder(self, class_name):
     def build_extension(self):
         from torch.utils.cpp_extension import BuildExtension
         return BuildExtension
+
+    def export_envs(self):
+        return ['NCCL']

accelerator/mps_accelerator.py

@@ -235,3 +235,6 @@ def build_extension(self):
         from torch.utils.cpp_extension import BuildExtension
 
         return BuildExtension
+
+    def export_envs(self):
+        return []

accelerator/npu_accelerator.py

@@ -255,3 +255,6 @@ def get_op_builder(self, class_name):
     def build_extension(self):
         from torch.utils.cpp_extension import BuildExtension
         return BuildExtension
+
+    def export_envs(self):
+        return ['ASCEND', 'HCCL', 'LD_LIBRARY', 'PATH']

deepspeed/launcher/runner.py

@@ -31,7 +31,7 @@
 from deepspeed.accelerator import get_accelerator
 
 DLTS_HOSTFILE = "/job/hostfile"
-EXPORT_ENVS = ['MLFLOW', 'NCCL', 'PYTHON', 'MV2', 'UCX']
+EXPORT_ENVS = ['MLFLOW', 'PYTHON', 'MV2', 'UCX']
 EXPORT_ENVS += NEBULA_EXPORT_ENVS
 DEEPSPEED_ENVIRONMENT_NAME = os.getenv("DS_ENV_FILE", ".deepspeed_env")
 DEEPSPEED_ENVIRONMENT_PATHS = [os.path.expanduser("~"), '.']
@@ -544,9 +544,10 @@ def main(args=None):
                 # key exists in launcher env -> var list should be used
                 excluded_vars += var_list
 
-        exports = ""
+        # load envs from accelerator
+        exports = EXPORT_ENVS + get_accelerator().export_envs()
         for var in env.keys():
-            if any([var.startswith(name) for name in EXPORT_ENVS]):
+            if any([var.startswith(name) for name in exports]):
                 if not any([var == name for name in excluded_vars]):
                     runner.add_export(var, env[var])