set backend to hccl only if hccl initialization sets a local rank

huggingface · Feb 5, 2025 · f66c5df · f66c5df
1 parent 7f72745
commit f66c5df
Showing 1 changed file with 4 additions and 10 deletions.
diff --git a/src/accelerate/state.py b/src/accelerate/state.py
@@ -730,18 +730,12 @@ def _prepare_backend(
             backend = "xla"
             distributed_type = DistributedType.XLA
         elif is_hpu_available():
-            from habana_frameworks.torch.distributed.hccl import initialize_distributed_hpu  # noqa: F401
+            import habana_frameworks.torch.distributed.hccl
 
-            print("after importing habana_frameworks in prepare_backend")
-
-            print(os.environ.get("LOCAL_RANK", -1))
-            print(os.environ.get("WORLD_SIZE", -1))
-            print(os.environ.get("RANK", -1))
-            print(os.environ.get("MASTER_ADDR", -1))
-            print(os.environ.get("MASTER_PORT", -1))
+            if int(os.environ.get("LOCAL_RANK", -1)) != -1:
+                backend = "hccl"
+                distributed_type = DistributedType.MULTI_HPU
 
-            backend = "hccl"
-            distributed_type = DistributedType.MULTI_HPU
         elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
             if is_mlu_available():
                 backend = "cncl"