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

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

Author: anshul-si

torchtitan/distributed/parallel_dims.py

@@ -151,7 +151,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,13 +135,18 @@ 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)],
+            param_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_param],
+            reduce_dtype=TORCH_DTYPE_MAP[job_config.training.mixed_precision_reduce],
             enable_compile=model_compile_enabled,
             enable_compiled_autograd=job_config.parallelism.enable_compiled_autograd,
+            cpu_offload=job_config.training.enable_cpu_offload,
         )
 
     return model
@@ -317,17 +322,33 @@ def apply_fsdp(
 def apply_ddp(
     model: nn.Module,
     dp_mesh: DeviceMesh,
+    param_dtype: torch.dtype,
+    reduce_dtype: torch.dtype,
     enable_compile: bool,
     enable_compiled_autograd: bool,
+    cpu_offload: bool = False,
 ):
-    if enable_compile:
-        if enable_compiled_autograd:
-            torch._dynamo.config.optimize_ddp = (
-                "python_reducer_without_compiled_forward"
-            )
-        else:
-            torch._dynamo.config.optimize_ddp = "ddp_optimizer"
+    mp_policy = MixedPrecisionPolicy(param_dtype=param_dtype, reduce_dtype=reduce_dtype)
+    replicate_config = {"device_mesh": dp_mesh, "mp_policy": mp_policy}
+    if cpu_offload:
+        replicate_config["offload_policy"] = CPUOffloadPolicy()
 
-    replicate(model, device_mesh=dp_mesh, bucket_cap_mb=100)
+    if model.tok_embeddings is not None:
+        replicate(
+            model.tok_embeddings,
+            **replicate_config,
+        )
+    for layer_id, transformer_block in model.layers.items():
+        replicate(
+            transformer_block,
+            **replicate_config,
+        )
+
+    if model.norm is not None and model.output is not None:
+        replicate(
+            [model.norm, model.output],
+            **replicate_config,
+        )
+    replicate(model, **replicate_config)
 
     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