[torchtitan][replicate] experimenting new replicate integration with torchtitan

ghstack-source-id: ea5b964
Pull Request resolved: #1714

Author: anshul-si

torchtitan/distributed/parallel_dims.py

@@ -100,7 +100,13 @@ def _build_mesh_with_ep(self) -> DeviceMesh:
         ):
             # dp_shard_mod_ep is needed even if it's 1, whose FSDP wrapping
             # helps the MoE layers do mixed precision training
-            if d > 1 or name == "dp_shard_mod_ep":
+            # dp_shard_in_ep is included even if it equals 1 when replicate > 1
+            # to make device_mesh compatible with replicate function
+            if (
+                d > 1
+                or name == "dp_shard_mod_ep"
+                or (name == "dp_shard_in_ep" and self.dp_replicate > 1)
+            ):
                 dims.append(d)
                 names.append(name)
 
@@ -151,7 +157,9 @@ def _build_mesh_without_ep(self) -> DeviceMesh:
             [self.pp, self.dp_replicate, self.dp_shard, self.cp, self.tp],
             ["pp", "dp_replicate", "dp_shard", "cp", "tp"],
         ):
-            if d > 1:
+            # Include dp_shard dimension even if it equals 1 when replicate > 1
+            # to make device_mesh compatible with replicate function
+            if d > 1 or (name == "dp_shard" and self.dp_replicate > 1):
                 dims.append(d)
                 names.append(name)
 

torchtitan/models/llama3/infra/parallelize.py

@@ -9,7 +9,7 @@
 
 import torch
 import torch.nn as nn
-from torch.distributed._composable.replicate import replicate
+from torch.distributed._composable.replicate_with_fsdp import replicate
 
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard, MixedPrecisionPolicy
@@ -135,11 +135,13 @@ def parallelize_llama(
         if job_config.training.enable_cpu_offload:
             logger.info("Applied CPU Offloading to the model")
     elif parallel_dims.dp_replicate_enabled:
-        if world_mesh.ndim > 1:
-            raise RuntimeError("DDP has not supported > 1D parallelism")
+        # if world_mesh.ndim > 1:
+        #    raise RuntimeError("DDP has not supported > 1D parallelism")
+
+        dp_mesh_dim_names = ("dp_replicate", "dp_shard")
         apply_ddp(
             model,
-            world_mesh,
+            world_mesh[tuple(dp_mesh_dim_names)],
             enable_compile=model_compile_enabled,
             enable_compiled_autograd=job_config.parallelism.enable_compiled_autograd,
         )
@@ -328,6 +330,6 @@ def apply_ddp(
         else:
             torch._dynamo.config.optimize_ddp = "ddp_optimizer"
 
-    replicate(model, device_mesh=dp_mesh, bucket_cap_mb=100)
+    replicate(model, device_mesh=dp_mesh)
 
     logger.info("Applied DDP to the model")

torchtitan/models/llama3/train_configs/debug_model.toml

@@ -43,7 +43,7 @@ steps = 10
 dataset = "c4_test"  # supported datasets: c4_test (2K), c4 (177M)
 
 [parallelism]
-data_parallel_replicate_degree = 1
+data_parallel_replicate_degree = 8
 data_parallel_shard_degree = -1
 fsdp_reshard_after_forward = "default" # default / never / always
 tensor_parallel_degree = 1