feat(pt): support CPU parallel training with PT

iProzd · Oct 16, 2024 · 6acfd3b · 6acfd3b
1 parent 5050f61
commit 6acfd3b
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py
@@ -105,8 +105,16 @@ def get_trainer(
     local_rank = os.environ.get("LOCAL_RANK")
     if local_rank is not None:
         local_rank = int(local_rank)
-        assert dist.is_nccl_available()
-        dist.init_process_group(backend="nccl")
+        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)
 
     def prepare_trainer_input_single(
         model_params_single, data_dict_single, rank=0, seed=None