[SimpleFSDP] Add support for ddp+tp (#1250)

As titled, this PR adds support for DDP+TP under SimpleFSDP's
`replicate` mode.

1. Profile trace for DDP. As seen, the DDP bwd communication is
`all-reduce`.

<img width="1109" alt="Screenshot 2025-06-01 at 1 10 07 PM"
src="https://github.com/user-attachments/assets/91ca56f4-c116-433d-98bf-96869a72de0c"
/>


2. Numerical convergence: As seen, the loss convergence discrepancy is
in 1e-3 for [ddp:2, tp:2] and [fsdp:2, tp:2] (with mixed-precision
training)


<img width="1568" alt="Screenshot 2025-06-01 at 11 39 49 PM"
src="https://github.com/user-attachments/assets/ef429276-da2b-41cd-bed3-fa880cd1efa6"
/>

The loss convergence is the same for [ddp:2, tp:2] and [fsdp:2, tp:2]
(without mixed-precision training)
<img width="1541" alt="Screenshot 2025-06-02 at 11 59 09 AM"
src="https://github.com/user-attachments/assets/2a18ef51-3ebf-4d5f-a27f-70fd15ee59d6"
/>

Author: ruisizhang123

torchtitan/experiments/simple_fsdp/simple_fsdp.py

@@ -77,16 +77,26 @@ def _distribute_dtensor(
     assert inner_mesh.mesh_dim_names is not None
     submesh_names = outer_mesh.mesh_dim_names + inner_mesh.mesh_dim_names
     spanned_mesh = outer_global_mesh[submesh_names]
-    shard_dim = placements[0].dim
-    split_factor = inner_spec.num_shards_map[shard_dim]
-    tensor_placement = (
-        (
-            _StridedShard(shard_dim, split_factor=split_factor)
-            if split_factor > 1
-            else placements[0]
-        ),
-        inner_spec.placements[0],
-    )
+
+    if placements[0].is_shard():
+        # for FSDP + TP dtensor placement
+        shard_dim = placements[0].dim
+        split_factor = inner_spec.num_shards_map[shard_dim]
+        tensor_placement = (
+            (
+                _StridedShard(shard_dim, split_factor=split_factor)
+                if split_factor > 1
+                else placements[0]
+            ),
+            inner_spec.placements[0],
+        )
+    elif placements[0].is_replicate():
+        # for DDP + TP dtensor placement
+        tensor_placement = (placements[0], inner_spec.placements[0])
+    else:
+        raise ValueError(
+            f"Unsupported placement {placements[0]} for distributing DTensor {tensor}"
+        )
 
     current_spec = DTensorSpec(
         mesh=outer_mesh,
@@ -154,10 +164,8 @@ def replicate_compute(self, x):
         # the gradients are partial tensors that needs to perform reduction
         # (i.e. DDP: allreduce, FSDP: reduce_scatter, HSDP: mix of both)
 
-        # NOTE: specifying mixed precision is only available in pytorch_intern24
-        #       https://github.com/tianyu-l/pytorch_intern24/pull/20
-        # support for FSDP + TP (assuming TP shards the inner-most dim)
-        if self.mode == "fully_shard" and x._spec.mesh.ndim == 2:
+        # support for FSDP/DDP + TP (assuming TP shards the inner-most dim)
+        if x._spec.mesh.mesh_dim_names[-1] == "tp":
             dp_placement, tp_placement = x._spec.placements
             # TODO: remove tp_mesh as an input arg to data_parallel API and use x._spec.mesh["tp"]
             #       after DeviceMesh supports slicing a non-root mesh
@@ -170,7 +178,8 @@ def replicate_compute(self, x):
                 sharded_local_tensor, dp_mesh, self.param_sharding
             )
 
-            # the actuall FSDP all-gather on dp_mesh
+            # the actual FSDP's fwd all-gather & bwd reduce-scatter
+            # DDP's bwd all-reduce on dp_mesh
             replicated_dtensor = sharded_dtensor.redistribute(
                 placements=self.compute_placements,
                 forward_dtype=self.param_dtype,

torchtitan/experiments/simple_fsdp/tests/integration_tests.py

@@ -131,7 +131,7 @@ def build_test_list():
             "hsdp",
             ngpu=4,
         ),
-        # TODO: Adds back after HSDP+TP & DDP+TP is supported by SimpleFSDP
+        # TODO: Adds back after HSDP+TP is supported by SimpleFSDP
         # OverrideDefinitions(
         #     [
         #         [
@@ -144,17 +144,17 @@ def build_test_list():
         #     "hsdp+tp",
         #     ngpu=8,
         # ),
-        # OverrideDefinitions(
-        #     [
-        #         [
-        #             "--parallelism.data_parallel_replicate_degree=2",
-        #             "--parallelism.tensor_parallel_degree=2",
-        #         ]
-        #     ],
-        #     "DDP+TP",
-        #     "ddp+tp",
-        #     ngpu=4,
-        # ),
+        OverrideDefinitions(
+            [
+                [
+                    "--parallelism.data_parallel_replicate_degree=2",
+                    "--parallelism.tensor_parallel_degree=2",
+                ]
+            ],
+            "DDP+TP",
+            "ddp+tp",
+            ngpu=4,
+        ),
         OverrideDefinitions(
             [
                 [