Update main.py

deepmodeling · Oct 23, 2024 · f9b43e8 · f9b43e8
1 parent 0b0d943
commit f9b43e8
Showing 1 changed file with 1 addition and 12 deletions.
diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py
@@ -105,18 +105,7 @@ def get_trainer(
     local_rank = os.environ.get("LOCAL_RANK")
     if local_rank is not None:
         local_rank = int(local_rank)
-        nccl_available = dist.is_nccl_available()
-        gloo_available = dist.is_gloo_available()
-        # nccl first
-        if nccl_available:
-            backend = "nccl"
-        elif gloo_available:
-            backend = "gloo"
-        else:
-            raise RuntimeError(
-                "No suitable backend found. Neither NCCL nor Gloo is available."
-            )
-        dist.init_process_group(backend=backend)
+        dist.init_process_group(backend="cuda:nccl,cpu:gloo")
 
     def prepare_trainer_input_single(
         model_params_single, data_dict_single, rank=0, seed=None