remove ring optimization when not enable

PrimeIntellect-ai · Nov 18, 2024 · cc36033 · cc36033
1 parent 89e7eeb
commit cc36033
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/src/zeroband/comms.py b/src/zeroband/comms.py
@@ -76,12 +76,13 @@ def __init__(
         self.cpu_local_mesh = init_device_mesh("cpu", mesh_shape=(self.local_pg.size(),))
 
         # Logging
-        self._optimize_ring_ranks()
-        if self.live_recovery_rank_src is not None:
-            self.live_recovery.ask_for_live_ckpt(self.live_recovery_rank_src)
-        self.global_pg.barrier().wait()
+        if self.enable:
+            self._optimize_ring_ranks()
+            if self.live_recovery_rank_src is not None:
+                self.live_recovery.ask_for_live_ckpt(self.live_recovery_rank_src)
+            self.global_pg.barrier().wait()
 
-        self._logger.info(f"global_pg size : {self.global_pg.size()}, local_pg size: {self.local_pg.size()}")
+            self._logger.info(f"global_pg size : {self.global_pg.size()}, local_pg size: {self.local_pg.size()}")
 
     def __del__(self):
         self._stop_heartbeat()