feat: support sequence parallel with rms_norm

tests/test_rms_norm_sequence_parallel.py

@@ -0,0 +1,137 @@
+import random
+import sys
+from collections.abc import Sequence
+
+import pytest
+import torch
+import torch.distributed as dist
+from packaging import version
+from torch.distributed.tensor.placement_types import (Partial, Placement,
+                                                      Replicate, Shard)
+
+import activation
+
+from .utils import assert_close, opcheck
+
+DTYPES = [torch.float32]
+NUM_TOKENS = [512]  # Arbitrary values for testing
+SEQUENCE_DIMS = [0, 1]  # 0 is for [T, D] (packed), 1 is for [B, S, D]
+D = [16]  # Arbitrary values for testing
+SEEDS = [0]
+
+from torch.distributed._tensor import DTensor
+from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
+from torch.distributed.tensor.parallel import (SequenceParallel,
+                                               parallelize_module)
+
+
+@pytest.fixture(scope="session", autouse=True)
+def init_dist(request):
+    if version.parse(torch.__version__) < version.parse("2.8"):
+        pytest.skip("torch>=2.8.0 is required for sequence parallel")
+        return
+
+    try:
+        dist.init_process_group(backend="nccl")
+        torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())
+    except Exception as e:
+        print(f"Failed to initialize torch.distributed: {e}")
+        pytest.skip("Failed to initialize torch.distributed")
+
+    if dist.get_world_size() < 2:
+        pytest.skip("Need at least 2 processes in dist group. "
+                    "You can run with `torchrun --nproc-per-node=2 "
+                    "--local-ranks-filter 0 -m pytest "
+                    "test_rms_norm_sequence_parallel.py`")
+
+    yield
+    dist.destroy_process_group()
+
+
+class Model(torch.nn.Module):
+
+    def __init__(self, num_tokens, d) -> None:
+        super().__init__()
+        self.rms_norm = activation.layers.RMSNorm(d)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.rms_norm(x)
+
+
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
+@pytest.mark.parametrize("d", D)
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("seed", SEEDS)
+@pytest.mark.parametrize("sequence_dim", SEQUENCE_DIMS)
+def test_rms_norm(
+    num_tokens: int,
+    d: int,
+    dtype: torch.dtype,
+    seed: int,
+    sequence_dim: int,
+) -> None:
+    if num_tokens % dist.get_world_size() != 0:
+        # It hangs at `y.full_tensor()` if not divisible
+        pytest.skip("num_tokens must be divisible by world_size for sharding")
+
+    random.seed(seed)
+    torch.manual_seed(seed)
+
+    num_ranks = dist.get_world_size()
+    rank = dist.get_rank()
+    mesh = init_device_mesh("cuda", (num_ranks, ), mesh_dim_names=("shard", ))
+
+    match sequence_dim:
+        case 0:
+            x_shape = (num_tokens, d)
+        case 1:
+            BATCH_SIZE = 2
+            x_shape = (BATCH_SIZE, num_tokens, d)
+        case _:
+            raise ValueError(f"Invalid sequence_dim: {sequence_dim}")
+
+    x = torch.randn(x_shape, dtype=dtype, requires_grad=True).cuda()
+    weight = torch.ones(d, dtype=dtype, requires_grad=True).cuda()
+    eps = 1e-05
+
+    x.retain_grad()
+    weight.retain_grad()
+
+    # Copy x, weight for reference
+    x_ref = x.detach().clone().requires_grad_(True)
+    weight_ref = weight.detach().clone().requires_grad_(True)
+
+    model_sharded = Model(num_tokens, d).to(dtype=dtype).cuda()
+    model_sharded.rms_norm.weight = torch.nn.Parameter(weight)
+    parallelize_module(
+        model_sharded, mesh,
+        {"rms_norm": SequenceParallel(sequence_dim=sequence_dim)})
+    x_sharded = DTensor.from_local(
+        x.chunk(num_ranks, dim=sequence_dim)[rank].contiguous(),
+        placements=(Shard(sequence_dim), ),
+        device_mesh=mesh,
+    )
+    y = model_sharded(x_sharded)
+    y_from_sharded = y.full_tensor()
+
+    model_unsharded = Model(num_tokens, d).to(dtype=dtype).cuda()
+    model_unsharded.rms_norm.weight = torch.nn.Parameter(weight_ref)
+
+    y_from_unsharded = model_unsharded(x_ref)
+
+    assert_close(y_from_sharded, y_from_unsharded)
+
+    # Backward
+    y_grad = torch.randn_like(y_from_unsharded)
+    y_from_sharded.backward(y_grad)
+    y_from_unsharded.backward(y_grad)
+
+    weight_grad_from_sharded = model_sharded.rms_norm.weight.grad._local_tensor
+    weight_grad_from_unsharded = model_unsharded.rms_norm.weight.grad
+
+    torch.distributed.all_reduce(x.grad, op=torch.distributed.ReduceOp.SUM)
+    torch.distributed.all_reduce(weight_grad_from_sharded,
+                                 op=torch.distributed.ReduceOp.SUM)
+
+    assert_close(x.grad, x_ref.grad)
+    assert_close(weight_grad_from_sharded, weight_grad_from_unsharded)

torch-ext/activation/rms_norm.py

@@ -1,4 +1,7 @@
+from collections.abc import Sequence
+
 import torch
+from packaging import version
 
 from ._ops import ops
 
@@ -70,3 +73,159 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
         residual_grad = grad if need_res else None
 
         return input_grad, residual_grad, weight_grad, None
+
+
+if version.parse(torch.__version__) >= version.parse("2.8"):
+    from torch.distributed.tensor._dtensor_spec import DTensorSpec
+    from torch.distributed.tensor._op_schema import (OpSchema, OpSpec,
+                                                     OpStrategy,
+                                                     RuntimeSchemaInfo)
+    from torch.distributed.tensor._ops.utils import (
+        generate_redistribute_costs, register_op_strategy)
+    from torch.distributed.tensor.placement_types import (Placement, Replicate,
+                                                          Shard)
+
+    @torch.library.register_fake(ops.rms_norm.default)
+    def rms_norm_abstract(x, weight, eps):
+        return torch.empty_like(x)
+
+    @torch.library.register_fake(ops.rms_norm_backward.default)
+    def rms_norm_backward_abstract(output_grad, x, weight, eps):
+        return torch.empty_like(x), torch.empty_like(weight)
+
+    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.rms_norm.default,
+                          schema_info=RuntimeSchemaInfo(1))
+    def rms_norm_strategy(op_schema: OpSchema) -> OpStrategy:
+        mesh = op_schema.get_mesh_from_args()
+
+        assert len(op_schema.args_schema) == 3
+        (
+            input_strategy,
+            weight_strategy,
+            _,  # eps
+        ) = op_schema.args_schema
+
+        assert isinstance(input_strategy, OpStrategy)
+        assert isinstance(weight_strategy, OpStrategy)
+
+        assert len(input_strategy.strategies) == len(
+            weight_strategy.strategies)
+
+        last_dim = input_strategy.ndim - 1
+        strategy = OpStrategy([])
+        for idx in range(len(input_strategy.strategies)):
+            input_src = input_strategy.strategies[idx].output_spec
+            weight_src = weight_strategy.strategies[idx].output_spec
+
+            assert isinstance(input_src, DTensorSpec)
+            assert isinstance(weight_src, DTensorSpec)
+
+            redistribute_costs = []
+
+            # Input can be sharded in any dim except the last dim.
+            input_tgt = DTensorSpec(
+                mesh=mesh,
+                placements=_replicate_dims_start_at(input_src.placements,
+                                                    last_dim),
+                tensor_meta=input_src.tensor_meta,
+            )
+            redistribute_costs.append(
+                generate_redistribute_costs(input_strategy, input_tgt))
+
+            # Weight cannot be sharded, so always replicate it.
+            weight_tgt = DTensorSpec(
+                mesh=mesh,
+                placements=(Replicate(), ),
+                tensor_meta=weight_src.tensor_meta,
+            )
+            redistribute_costs.append(
+                generate_redistribute_costs(weight_strategy, weight_tgt))
+
+            strategy.strategies.append(
+                OpSpec(
+                    output_specs=input_tgt,
+                    input_specs=[input_tgt, weight_tgt],
+                    redistribute_cost=redistribute_costs,
+                ))
+        return strategy
+
+    @register_op_strategy(ops.rms_norm_backward.default,
+                          schema_info=RuntimeSchemaInfo(1))
+    def rms_norm_backward_strategy(op_schema: OpSchema) -> OpStrategy:
+        mesh = op_schema.get_mesh_from_args()
+
+        assert len(op_schema.args_schema) == 4
+        (
+            output_grad_strategy,
+            input_strategy,
+            weight_strategy,
+            _,  # eps
+        ) = op_schema.args_schema
+
+        assert isinstance(output_grad_strategy, OpStrategy)
+        assert isinstance(input_strategy, OpStrategy)
+        assert isinstance(weight_strategy, OpStrategy)
+
+        assert len(input_strategy.strategies) == len(
+            weight_strategy.strategies)
+        assert len(input_strategy.strategies) == len(
+            output_grad_strategy.strategies)
+
+        last_dim = input_strategy.ndim - 1
+        strategy = OpStrategy([])
+        for idx in range(len(input_strategy.strategies)):
+            output_grad_src = output_grad_strategy.strategies[idx].output_spec
+            input_src = input_strategy.strategies[idx].output_spec
+            weight_src = weight_strategy.strategies[idx].output_spec
+
+            assert isinstance(output_grad_src, DTensorSpec)
+            assert isinstance(input_src, DTensorSpec)
+            assert isinstance(weight_src, DTensorSpec)
+
+            redistribute_costs = []
+
+            # Output grad and input can be sharded in any dim except the last dim.
+            output_grad_tgt = DTensorSpec(
+                mesh=mesh,
+                placements=_replicate_dims_start_at(output_grad_src.placements,
+                                                    last_dim),
+                tensor_meta=output_grad_src.tensor_meta,
+            )
+            redistribute_costs.append(
+                generate_redistribute_costs(output_grad_strategy,
+                                            output_grad_tgt))
+            input_tgt = DTensorSpec(
+                mesh=mesh,
+                placements=_replicate_dims_start_at(input_src.placements,
+                                                    last_dim),
+                tensor_meta=input_src.tensor_meta,
+            )
+            redistribute_costs.append(
+                generate_redistribute_costs(input_strategy, input_tgt))
+
+            # Weight cannot be sharded, so always replicate it.
+            weight_tgt = DTensorSpec(
+                mesh=mesh,
+                placements=(Replicate(), ),
+                tensor_meta=weight_src.tensor_meta,
+            )
+            redistribute_costs.append(
+                generate_redistribute_costs(weight_strategy, weight_tgt))
+
+            strategy.strategies.append(
+                OpSpec(
+                    output_specs=[input_tgt, weight_tgt],
+                    input_specs=[output_grad_tgt, input_tgt, weight_tgt],
+                    redistribute_cost=redistribute_costs,
+                ))
+        return strategy