Fused RMSNorm: multi-row programs and a one-pass backward for narrow norms

deepspeed/ops/triton_ops/fused_rms_norm.py

@@ -31,6 +31,10 @@
 _MAX_FORWARD_BLOCK = 2048
 _DWEIGHT_BLOCK_M = 16
 _DWEIGHT_BLOCK_N = 256
+# Norms at most this wide, such as the 128-wide query and key norms, give each program several rows: one row per
+# program leaves most of a program idle, and the backward reads each row once for both gradients.
+_NARROW_MAX_COLS = 128
+_NARROW_BLOCK_ELEMENTS = 4096
 _FUSED_INSTANCE_ATTRIBUTE = "_deepspeed_use_fused_rms_norm"
 # Class descriptors bind to the current receiver after copying or replication and are not stored in model pickles.
 _FUSED_CLASS_FORWARDS = {}
@@ -134,6 +138,71 @@ def _rms_norm_dweight_partial_kernel(
         partial = tl.sum(product, axis=0)
         tl.store(partial_ptr + row_block * n_cols + col_offsets, partial, mask=col_offsets < n_cols)
 
+    @triton.jit
+    def _rms_norm_narrow_forward_kernel(
+        hidden_ptr,
+        weight_ptr,
+        out_ptr,
+        rstd_ptr,
+        n_rows,
+        n_cols: tl.constexpr,
+        eps: tl.constexpr,
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+    ):
+        rows = tl.program_id(0).to(tl.int64) * BLOCK_M + tl.arange(0, BLOCK_M)
+        cols = tl.arange(0, BLOCK_N)
+        row_mask = rows < n_rows
+        mask = row_mask[:, None] & (cols[None, :] < n_cols)
+        offsets = rows[:, None] * n_cols + cols[None, :]
+
+        hidden = tl.load(hidden_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        variance = tl.sum(hidden * hidden, axis=1) / n_cols
+        rstd = tl.rsqrt(variance + eps)
+        tl.store(rstd_ptr + rows, rstd, mask=row_mask)
+
+        normalized = (hidden * rstd[:, None]).to(out_ptr.dtype.element_ty)
+        weight = tl.load(weight_ptr + cols, mask=cols < n_cols, other=0.0)
+        tl.store(out_ptr + offsets, normalized * weight[None, :], mask=mask)
+
+    @triton.jit
+    def _rms_norm_narrow_backward_kernel(
+        grad_out_ptr,
+        hidden_ptr,
+        weight_ptr,
+        rstd_ptr,
+        grad_hidden_ptr,
+        partial_ptr,
+        n_rows,
+        n_cols: tl.constexpr,
+        BLOCK_M: tl.constexpr,
+        BLOCK_N: tl.constexpr,
+    ):
+        row_block = tl.program_id(0).to(tl.int64)
+        rows = row_block * BLOCK_M + tl.arange(0, BLOCK_M)
+        cols = tl.arange(0, BLOCK_N)
+        row_mask = rows < n_rows
+        col_mask = cols < n_cols
+        mask = row_mask[:, None] & col_mask[None, :]
+        offsets = rows[:, None] * n_cols + cols[None, :]
+
+        grad_out = tl.load(grad_out_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        hidden = tl.load(hidden_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        weight = tl.load(weight_ptr + cols, mask=col_mask, other=0.0).to(tl.float32)
+        rstd = tl.load(rstd_ptr + rows, mask=row_mask, other=0.0).to(tl.float32)
+
+        # The same steps as the row kernels, so values round where the eager backward rounds them.
+        grad_norm = (grad_out * weight[None, :]).to(grad_hidden_ptr.dtype.element_ty).to(tl.float32)
+        grad_scaled = grad_norm * rstd[:, None]
+        grad_rstd = tl.sum(grad_norm * hidden, axis=1)
+        grad_variance = (-0.5 * grad_rstd) * (rstd * rstd * rstd)
+        grad_hidden = grad_scaled + (grad_variance / n_cols)[:, None] * (2.0 * hidden)
+        tl.store(grad_hidden_ptr + offsets, grad_hidden, mask=mask)
+
+        normalized = (hidden * rstd[:, None]).to(hidden_ptr.dtype.element_ty).to(tl.float32)
+        product = (grad_out * normalized).to(hidden_ptr.dtype.element_ty).to(tl.float32)
+        tl.store(partial_ptr + row_block * n_cols + cols, tl.sum(product, axis=0), mask=col_mask)
+
 
 def is_available() -> bool:
     """Whether this build can run the fused RMSNorm kernels."""
@@ -181,6 +250,13 @@ def _block_n(n_cols: int) -> int:
     return max(16, triton.next_power_of_2(n_cols))
 
 
+def _narrow_block_m(n_cols: int):
+    """Rows per program for the narrow kernels, or None where the row kernels apply."""
+    if n_cols > _NARROW_MAX_COLS:
+        return None
+    return _NARROW_BLOCK_ELEMENTS // _block_n(n_cols)
+
+
 class _FusedRMSNorm(torch.autograd.Function):
     """RMSNorm autograd with FP32 reductions and a rounded pre-gamma value."""
 
@@ -193,7 +269,20 @@ def forward(ctx, hidden, weight, eps):
         ctx.save_for_backward(hidden, weight, rstd)
         ctx.n_cols = n_cols
 
-        if n_rows > 0:
+        block_m = _narrow_block_m(n_cols)
+        if n_rows > 0 and block_m is not None:
+            _rms_norm_narrow_forward_kernel[(triton.cdiv(n_rows, block_m), )](
+                hidden,
+                weight,
+                out,
+                rstd,
+                n_rows,
+                n_cols,
+                eps,
+                BLOCK_M=block_m,
+                BLOCK_N=_block_n(n_cols),
+            )
+        elif n_rows > 0:
             _rms_norm_forward_kernel[(n_rows, )](
                 hidden,
                 weight,
@@ -218,6 +307,24 @@ def backward(ctx, grad_out):
         if n_rows == 0:
             return grad_hidden, torch.zeros_like(weight), None
 
+        block_m = _narrow_block_m(n_cols)
+        if block_m is not None:
+            n_row_blocks = triton.cdiv(n_rows, block_m)
+            partial = torch.empty((n_row_blocks, n_cols), dtype=torch.float32, device=hidden.device)
+            _rms_norm_narrow_backward_kernel[(n_row_blocks, )](
+                grad_out,
+                hidden,
+                weight,
+                rstd,
+                grad_hidden,
+                partial,
+                n_rows,
+                n_cols,
+                BLOCK_M=block_m,
+                BLOCK_N=_block_n(n_cols),
+            )
+            return grad_hidden, partial.sum(dim=0).to(weight.dtype), None
+
         _rms_norm_dx_kernel[(n_rows, )](
             grad_out,
             hidden,

docs/code-docs/source/autoep.rst

@@ -442,6 +442,16 @@ Where the kernels run, results are not bitwise identical to eager RMSNorm: the
 kernels' FP32 reductions sum in a different order, so a small fraction of
 outputs and gradients differ in their last bits.
 
+For norms at most 128 wide, the fused kernels schedule several rows per
+program and combine the input- and weight-gradient reads in backward. Wider
+norms retain the one-row kernels. The different FP32 reduction order still
+uses the existing ULP accuracy contract, not bitwise equality with eager.
+At width one, the pre-gamma input derivative in exact arithmetic is
+``eps / (x**2 + eps)**(3/2)``, which is nonzero when ``eps > 0``. With a
+very small epsilon and large ``abs(x)``, low-precision rounding can hide this
+gradient; width-one tests use a measurable positive epsilon and an FP64
+reference rather than assuming the derivative is zero.
+
 Requirements and limits:
 
 - The kernels need CUDA with Triton. On ROCm or without Triton, replaced

tests/unit/v1/ops/triton_ops/test_fused_rms_norm.py

@@ -126,6 +126,52 @@ def _check_fused_rms_norm_matches_hf_forward_and_backward(dtype, shape):
     _assert_ulp_close(fused_weight.grad, eager_weight.grad, max_ulp=8, min_frac_within_1=0.95, label="dgamma")
 
 
+def _check_fused_rms_norm_width_one_has_nonzero_input_gradient(dtype):
+    device = _device()
+    eps = 0.25
+    hidden = torch.linspace(0.25, 0.75, 257, device=device).to(dtype).reshape(-1, 1)
+    weight = torch.ones((1, ), dtype=dtype, device=device)
+    upstream = torch.ones_like(hidden)
+
+    eager_hidden = hidden.clone().requires_grad_(True)
+    eager_weight = weight.clone().requires_grad_(True)
+    eager_out = _hf_rms_norm(eager_hidden, eager_weight, eps)
+    eager_out.backward(upstream)
+
+    fused_hidden = hidden.clone().requires_grad_(True)
+    fused_weight = weight.clone().requires_grad_(True)
+    fused_out = _rms_norm().fused_rms_norm(fused_hidden, fused_weight, eps)
+    fused_out.backward(upstream)
+
+    # Before gamma, d(x / sqrt(x**2 + eps))/dx = eps / (x**2 + eps)**1.5.
+    reference_dx = (eps / (hidden.double().square() + eps).pow(1.5)).to(dtype)
+    assert torch.all(reference_dx > 0)
+    torch.testing.assert_close(fused_hidden.grad, reference_dx, rtol=2 * torch.finfo(dtype).eps, atol=0)
+    _assert_ulp_close(fused_out, eager_out, max_ulp=2, min_frac_within_1=0.99, label="width-one forward")
+    _assert_ulp_close(fused_hidden.grad, eager_hidden.grad, max_ulp=8, min_frac_within_1=0.95, label="width-one dx")
+    _assert_ulp_close(fused_weight.grad,
+                      eager_weight.grad,
+                      max_ulp=8,
+                      min_frac_within_1=0.95,
+                      label="width-one dgamma")
+
+
+def _check_fused_rms_norm_weight_gradient_is_deterministic(width):
+    device = _device()
+    generator = torch.Generator(device=device).manual_seed(7)
+    hidden = torch.randn((8192, width), device=device, dtype=torch.bfloat16, generator=generator)
+    weight = torch.randn((width, ), device=device, dtype=torch.bfloat16, generator=generator)
+    upstream = torch.randn((8192, width), device=device, dtype=torch.bfloat16, generator=generator)
+    grads = []
+    for _ in range(2):
+        leaf_hidden = hidden.clone().requires_grad_(True)
+        leaf_weight = weight.clone().requires_grad_(True)
+        _rms_norm().fused_rms_norm(leaf_hidden, leaf_weight, 1e-6).backward(upstream)
+        grads.append((leaf_hidden.grad, leaf_weight.grad))
+    assert torch.equal(grads[0][0], grads[1][0])
+    assert torch.equal(grads[0][1], grads[1][1])
+
+
 def _check_fused_rms_norm_handles_head_major_non_contiguous_input(dtype):
     device = _device()
     generator = torch.Generator(device=device).manual_seed(20260923)
@@ -634,11 +680,38 @@ class TestFusedRMSNormCUDA(DistributedTest):
     init_distributed = False
 
     @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
-    @pytest.mark.parametrize("shape", [(0, 2048), (7, 128), (5, 2048), (3, 130)])
+    @pytest.mark.parametrize(
+        "shape",
+        [
+            (0, 2048),
+            (7, 128),
+            (5, 2048),
+            (3, 130),
+            # First width beyond the narrow cutoff: it must retain the wide-path accuracy contract.
+            (3, 129),
+            # Narrow norms run several rows per program; these cover partial last blocks and masked columns.
+            (4099, 128),
+            (2, 300, 32, 128),
+            (1000, 96),
+            (257, 64),
+            (33, 16),
+            (257, 7),
+            (33, 15),
+        ])
     def test_fused_rms_norm_matches_hf_forward_and_backward(self, dtype, shape):
         _require_fused_engine()
         _check_fused_rms_norm_matches_hf_forward_and_backward(dtype, shape)
 
+    @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+    def test_fused_rms_norm_width_one_has_nonzero_input_gradient(self, dtype):
+        _require_fused_engine()
+        _check_fused_rms_norm_width_one_has_nonzero_input_gradient(dtype)
+
+    @pytest.mark.parametrize("width", [128, 2048])
+    def test_fused_rms_norm_weight_gradient_is_deterministic(self, width):
+        _require_fused_engine()
+        _check_fused_rms_norm_weight_gradient_is_deterministic(width)
+
     @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
     def test_fused_rms_norm_handles_head_major_non_contiguous_input(self, dtype):
         _require_fused_engine()