Update on "[cp] fix the device mesh access issue when CP is not used …

…with DP" [ghstack-poisoned]
pytorch · Nov 20, 2024 · 9736890 · 9736890
1 parent ac2c7e1
commit 9736890
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchtitan/parallelisms/parallel_dims.py b/torchtitan/parallelisms/parallel_dims.py
@@ -78,7 +78,7 @@ def build_mesh(self, device_type):
         if self.cp > 1:
             if self.dp_replicate > 1 and self.dp_shard > 1:  # HSDP
                 mesh["dp_replicate", "dp_shard", "cp"]._flatten(mesh_dim_name="dp_cp")
-            elif self.dp_replicate > 1 or self.dp_shard > 1:  # FSDP or DDP
+            elif self.dp_shard > 1:  # FSDP
                 mesh["dp", "cp"]._flatten(mesh_dim_name="dp_cp")
 
         return mesh