fix cpu local mehs

PrimeIntellect-ai · Nov 20, 2024 · f0245f1 · f0245f1
1 parent 84bf9b3
commit f0245f1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/zeroband/diloco.py b/src/zeroband/diloco.py
@@ -9,6 +9,7 @@
 from zeroband.utils.logging import get_logger
 from torch.distributed._tensor.api import DTensor
 from functools import lru_cache
+from torch.distributed.device_mesh import init_device_mesh
 
 
 class DilocoConfig(BaseConfig):
@@ -67,6 +68,7 @@ def __init__(
             # just force compilation
 
         self.pccl_communicator = pccl_communicator
+        self.cpu_local_mesh = init_device_mesh("cpu", mesh_shape=(self.local_pg.size(),))
 
         self._logger = get_logger()
         self.world_info = get_world_info()