fix(sae): do not init device mesh in single device mode

OpenMOSS · Jul 31, 2024 · 63f0a13 · 63f0a13
1 parent 7785d67
commit 63f0a13
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/lm_saes/sae.py b/src/lm_saes/sae.py
@@ -60,9 +60,10 @@ def __init__(self, cfg: SAEConfig):
         )
         torch.nn.init.kaiming_uniform_(self.encoder.weight)
         torch.nn.init.zeros_(self.encoder.bias)
-        self.device_mesh = init_device_mesh(
-            "cuda", (cfg.ddp_size, cfg.tp_size), mesh_dim_names=("ddp", "tp")
-        )
+        if cfg.tp_size > 1 or cfg.ddp_size > 1:
+            self.device_mesh = init_device_mesh(
+                "cuda", (cfg.ddp_size, cfg.tp_size), mesh_dim_names=("ddp", "tp")
+            )
 
         if cfg.use_glu_encoder: