fix: fix fused add rms norm sharding strategy

tests/test_fused_add_rms_norm_sequence_parallel.py

@@ -55,7 +55,7 @@ class Model(torch.nn.Module):
         self.fused_add_rms_norm = activation.layers.FusedAddRMSNorm(d)
 
     def forward(self, x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
-        return self.fused_add_rms_norm(x, residual=residual)
+        return self.fused_add_rms_norm(x, residual)
 
 
 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@@ -122,18 +122,18 @@ def test_fused_add_rms_norm_sequence_parallel(
         ResidualSequenceParallel(sequence_dim=sequence_dim)
     })
 
-    x_sharded = DTensor.from_local(
-        x.chunk(num_ranks, dim=sequence_dim)[rank].contiguous(),
-        placements=(Shard(sequence_dim), ),
+    x_replicate = DTensor.from_local(
+        x,
+        placements=(Replicate(), ),
         device_mesh=mesh,
     )
-    residual_sharded = DTensor.from_local(
-        residual.chunk(num_ranks, dim=sequence_dim)[rank].contiguous(),
-        placements=(Shard(sequence_dim), ),
+    residual_replicate = DTensor.from_local(
+        residual,
+        placements=(Replicate(), ),
         device_mesh=mesh,
     )
 
-    y, add_output = model_sharded(x_sharded, residual_sharded)
+    y, add_output = model_sharded(x_replicate, residual_replicate)
 
     y_from_sharded = y.full_tensor()
     add_output_from_sharded = add_output.full_tensor()
@@ -156,21 +156,16 @@ def test_fused_add_rms_norm_sequence_parallel(
     (y_grad * y_from_unsharded +
      add_output_grad * add_output_from_unsharded).sum().backward()
 
-    weight_grad_from_sharded = model_sharded.fused_add_rms_norm.weight.grad._local_tensor
+    weight_grad_from_sharded = model_sharded.fused_add_rms_norm.weight.grad.full_tensor(
+    )
     weight_grad_from_unsharded = model_unsharded.fused_add_rms_norm.weight.grad
 
     assert (x.grad is None) ^ x_requires_grad
     assert (residual.grad is None) ^ residual_requires_grad
 
-    torch.distributed.all_reduce(weight_grad_from_sharded,
-                                 op=torch.distributed.ReduceOp.SUM)
-
-    if x.grad is not None:
-        torch.distributed.all_reduce(x.grad, op=torch.distributed.ReduceOp.SUM)
+    if x_requires_grad:
         assert_close(x.grad, x_ref.grad)
-    if residual.grad is not None:
-        torch.distributed.all_reduce(residual.grad,
-                                     op=torch.distributed.ReduceOp.SUM)
+    if residual_requires_grad:
         assert_close(residual.grad, residual_ref.grad)
 
     assert_close(weight_grad_from_sharded, weight_grad_from_unsharded)

torch-ext/activation/fused_add_rms_norm_meta.py

@@ -4,6 +4,9 @@ import torch
 from torch.distributed.tensor._dtensor_spec import DTensorSpec
 from torch.distributed.tensor._op_schema import (OpSchema, OpSpec, OpStrategy,
                                                  RuntimeSchemaInfo)
+from torch.distributed.tensor._ops._math_ops import (
+    _infer_reduce_dims_map, _replicate_dims_start_at,
+    map_placements_after_reduction)
 from torch.distributed.tensor._ops.utils import (generate_redistribute_costs,
                                                  register_op_strategy)
 from torch.distributed.tensor.placement_types import (Placement, Replicate,
@@ -19,17 +22,6 @@ def register_fused_add_rms_norm_meta():
     pass
 
 
-def _replicate_dims_start_at(placements: Sequence[Placement],
-                             start_dim: int = 0) -> tuple[Placement, ...]:
-    new_placements: list[Placement] = []
-    for p in placements:
-        if p.is_partial() or (isinstance(p, Shard) and p.dim >= start_dim):
-            new_placements.append(Replicate())  # make it replicate
-        else:
-            new_placements.append(p)  # keep the placement
-    return tuple(new_placements)
-
-
 @register_op_strategy(ops.fused_add_rms_norm.default,
                       schema_info=RuntimeSchemaInfo(1))
 def fused_add_rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
@@ -89,7 +81,7 @@ def fused_add_rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
         # Weight cannot be sharded, so always replicate it.
         weight_tgt = DTensorSpec(
             mesh=mesh,
-            placements=(Replicate(), ),
+            placements=_replicate_dims_start_at(weight_src.placements),
             tensor_meta=weight_src.tensor_meta,
         )
         redistribute_costs.append(
@@ -141,6 +133,8 @@ def fused_add_rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
     )
 
     last_dim = output_grad_strategy.ndim - 1
+    outer_dims = list(range(last_dim))
+
     strategy = OpStrategy([])
     for output_grad, add_output_grad, add_output, weight in zipped:
         output_grad_src = output_grad.output_spec
@@ -179,16 +173,35 @@ def fused_add_rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
         # Weight cannot be sharded, so always replicate it.
         weight_tgt = DTensorSpec(
             mesh=mesh,
-            placements=(Replicate(), ),
+            placements=_replicate_dims_start_at(weight_src.placements),
             tensor_meta=weight_src.tensor_meta,
         )
         redistribute_costs.append(
             generate_redistribute_costs(weight_strategy, weight_tgt))
 
+        # from torch/distributed/tensor/_ops/_math_ops.py::layer_norm_bwd_strategy()
+
+        # Weight cannot be sharded, so always replicate it.
+        # TODO: now d_weight spec follows input spec w/ a reduction.
+        # we may need to change to a pointwise rule over grad_out and
+        # input, then apply a reduction.
+        inp_placements = _replicate_dims_start_at(output_grad_src.placements,
+                                                  last_dim)
+        reduce_dims_map = _infer_reduce_dims_map(outer_dims,
+                                                 output_grad_src.ndim, False)
+        out_placements = map_placements_after_reduction(
+            inp_placements, outer_dims, reduce_dims_map, "sum")
+        weight_grad_tgt = DTensorSpec(
+            mesh=mesh,
+            placements=out_placements,
+            tensor_meta=weight_src.tensor_meta,
+        )
+
         strategy.strategies.append(
             OpSpec(
                 output_specs=[
-                    output_grad_tgt if need_input_grad else None, weight_tgt
+                    output_grad_tgt if need_input_grad else None,
+                    weight_grad_tgt
                 ],
                 input_specs=[
                     output_grad_tgt, add_output_grad_tgt, add_output_tgt,