refactor(rms_norm): move RMS normalization logic to a new module for better organization and maintainability

torch-ext/activation/rms_norm.py

@@ -76,156 +76,5 @@ class FusedAddRMSNormFunction(torch.autograd.Function):
 
 
 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
+    from .rms_norm_meta import register_rms_norm_meta
+    register_rms_norm_meta()

torch-ext/activation/rms_norm_meta.py

@@ -0,0 +1,163 @@
+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.utils import (generate_redistribute_costs,
+                                                 register_op_strategy)
+from torch.distributed.tensor.placement_types import (Placement, Replicate,
+                                                      Shard)
+
+from ._ops import ops
+
+
+def register_rms_norm_meta():
+    """Dummy function to register the meta functions.
+    Registration happens at import time by the decorators below.
+    """
+    pass
+
+
+@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