From 52329febf68c0b3c1cc4e88c353f4492f74fa670 Mon Sep 17 00:00:00 2001 From: VikingDeng Date: Wed, 21 Oct 2026 16:25:03 -0500 Subject: [PATCH 1/3] [PyTorch] Add RMSNormGated fusible op Add a basic op that computes rmsnorm(x) * silu(gate) with the gate as an extra input, built on the existing rmsnorm_fwd/rmsnorm_bwd kernels. The normalized output and its gradient are rounded to the input dtype around the gate, matching Megatron-LM's unfused GDN path. Quantized output is not supported yet. First step of #3724. Co-Authored-By: Claude Opus 5.5 Signed-off-by: VikingDeng --- docs/api/pytorch.rst | 2 + tests/pytorch/test_fusible_ops.py | 137 +++++++++ .../pytorch/ops/basic/__init__.py | 1 + .../pytorch/ops/basic/rmsnorm_gated.py | 279 ++++++++++++++++++ 4 files changed, 419 insertions(+) create mode 100644 transformer_engine/pytorch/ops/basic/rmsnorm_gated.py diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 31d469cb77f..1c65b21bd66 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -259,6 +259,8 @@ Operation fuser .. autoapiclass:: transformer_engine.pytorch.ops.RMSNorm +.. autoapiclass:: transformer_engine.pytorch.ops.RMSNormGated + .. autoapiclass:: transformer_engine.pytorch.ops.SReGLU .. autoapiclass:: transformer_engine.pytorch.ops.SReLU diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index d4455c09d68..bbee82b86af 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2332,6 +2332,143 @@ def test_rmsnorm( torch.testing.assert_close(dx_test, x_ref.grad, **tols) torch.testing.assert_close(dw_test, w_ref.grad, **tols) + @pytest.mark.parametrize("in_shape", ((1, 37, 16, 128), (6, 16, 64))) + @pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16)) + @pytest.mark.parametrize("zero_centered_gamma", (False, True)) + @pytest.mark.parametrize("gate_layout", ("contiguous", "strided")) + def test_rmsnorm_gated( + self, + *, + in_shape: Iterable[int], + dtype: torch.dtype, + device: torch.device = "cuda", + eps: float = 1e-6, + zero_centered_gamma: bool, + gate_layout: str, + ) -> None: + """RMSNorm with SiLU gate""" + + # Skip unsupported BF16 + if dtype == torch.bfloat16 and not is_bf16_available(): + pytest.skip("BF16 requires sm_80 or higher") + + # Make input and weight shapes consistent + in_shape = list(in_shape) + weight_shape = (in_shape[-1],) + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + ) + w_ref, w_test = make_reference_and_test_tensors( + weight_shape, + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Gate tensor, optionally a non-contiguous view of a larger tensor + if gate_layout == "contiguous": + g_ref, g_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + ) + else: + g_base_ref, g_base_test = make_reference_and_test_tensors( + [*in_shape[:-1], 2 * in_shape[-1]], + test_dtype=dtype, + test_device=device, + ) + g_ref = g_base_ref[..., ::2] + g_test = g_base_test[..., ::2] + g_ref.retain_grad() + g_test.retain_grad() + assert not g_test.is_contiguous() + + # Plain PyTorch implementation + var_ref = x_ref.square().mean(dim=-1, keepdim=True) + gamma_ref = 1 + w_ref if zero_centered_gamma else w_ref + n_ref = x_ref * torch.rsqrt(eps + var_ref) * gamma_ref + # The op rounds the normalized output and its gradient to the input + # dtype around the gate (as the unfused Megatron path does), so the + # reference does too. + n_ref = n_ref + (n_ref.to(dtype).to(n_ref.dtype) - n_ref).detach() + n_ref.register_hook(lambda grad: grad.to(dtype).to(grad.dtype)) + y_ref = n_ref * torch.nn.functional.silu(g_ref) + y_ref.backward(dy_ref) + + # Implementation with fusible operation + op = te_ops.RMSNormGated( + weight_shape, + eps=eps, + device=device, + dtype=dtype, + zero_centered_gamma=zero_centered_gamma, + ) + with torch.no_grad(): + op.weight.copy_(w_test) + del w_test + y_test = op(x_test, g_test) + y_test.backward(dy_test) + + # Expected numerical error + tols = dtype_tols(dtype) + + # Check results + y_test = y_test.to(dtype=torch.float64, device="cpu") + dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") + dw_test = op.weight.grad.to(dtype=torch.float64, device="cpu") + dg_test = g_test.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(y_test, y_ref, **tols) + torch.testing.assert_close(dx_test, x_ref.grad, **tols) + torch.testing.assert_close(dw_test, w_ref.grad, **tols) + torch.testing.assert_close(dg_test, g_ref.grad, **tols) + + @pytest.mark.parametrize("quantization", _quantization_list[1:]) + def test_rmsnorm_gated_quantized_output( + self, + *, + quantization: str, + device: torch.device = "cuda", + in_shape: Iterable[int] = (16, 128), + dtype: torch.dtype = torch.bfloat16, + ) -> None: + """RMSNorm with SiLU gate rejects quantized output""" + + # Skip unsupported BF16 + if not is_bf16_available(): + pytest.skip("BF16 requires sm_80 or higher") + + # Random data + _, x_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + ) + _, g_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + ) + + # Implementation with fusible operation + op = te_ops.RMSNormGated(in_shape[-1], device=device, dtype=dtype) + forward = te_ops.Sequential( + op, + te_ops.Quantize(forward=True, backward=False), + ) + with te.autocast(enabled=True, recipe=make_recipe(quantization)): + with pytest.raises(NotImplementedError, match="quantized output"): + forward(x_test, g_test) + @pytest.mark.parametrize("in_shape", ((32,), (6, 16, 64), (32, 64))) @pytest.mark.parametrize("dtype", _dtypes) def test_l2normalization( diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 08299932a3d..26ab7efa7a4 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -36,6 +36,7 @@ from .reduce_scatter import ReduceScatter from .reshape import Reshape from .rmsnorm import RMSNorm +from .rmsnorm_gated import RMSNormGated from .swiglu import ( ClampedSwiGLU, ScaledClampedQGeGLU, diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm_gated.py b/transformer_engine/pytorch/ops/basic/rmsnorm_gated.py new file mode 100644 index 00000000000..65eeb2b8823 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/rmsnorm_gated.py @@ -0,0 +1,279 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusable operation for RMSNorm with SiLU gating.""" + +from __future__ import annotations +from collections.abc import Iterable, Sequence +import math +import os +from typing import Any, Optional + +import torch + +from transformer_engine_torch import rmsnorm_bwd, rmsnorm_fwd +from ...constants import TE_DType +from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload +from ...tensor import Quantizer +from ...utils import ( + canonicalize_device, + canonicalize_dtype, + clear_tensor_data, + devices_match, +) +from ..op import BasicOperation, OperationContext +from .._common import maybe_autocast_dtype, maybe_dequantize + + +class RMSNormGated(BasicOperation): + r"""Root Mean Square Layer Normalization with SiLU gate + + Applies Root Mean Square Layer Normalization over a mini-batch of + inputs and multiplies the result by the SiLU activation of a gate + tensor: + + .. math:: + y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * \gamma * \mathrm{SiLU}(g) + + :math:`\gamma` is a learnable affine transform parameter that + matches the inner-most dimensions of the input tensor. The gate + tensor :math:`g` is an extra operation input and must have the + same shape as the input tensor. + + Output quantization is not supported yet. The output is always + returned in FP32, FP16, or BF16. + + Parameters + ---------- + normalized_shape : int or iterable of int + Inner dimensions of input tensor + eps : float, default = 1e-5 + A value added to the denominator for numerical stability + device : torch.device, default = default CUDA device + Tensor device + dtype : torch.dtype, default = default dtype + Tensor datatype + zero_centered_gamma : bool, default = False + If ``True``, the :math:`\gamma` parameter is initialized to zero + and the calculation changes to + + .. math:: + y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) * \mathrm{SiLU}(g) + + sm_margin : int, default = 0 + Number of SMs to exclude when launching CUDA kernels. This + helps overlap with other kernels, e.g. communication kernels. + For more fine-grained control, provide a dict with the SM + margin at each compute stage ("forward", "backward", + "inference"). + + """ + + num_extra_inputs: int = 1 + + def __init__( + self, + normalized_shape: Iterable[int] | int, + *, + eps: float = 1e-5, + device: Optional[torch.device | str] = None, + dtype: Optional[torch.dtype] = None, + zero_centered_gamma: bool = False, + sm_margin: int = 0, + ) -> None: + super().__init__() + self.eps: float = eps + self.zero_centered_gamma: bool = zero_centered_gamma + + # Parameter shape + if not isinstance(normalized_shape, Iterable): + normalized_shape = (normalized_shape,) + else: + normalized_shape = tuple(normalized_shape) + + # Parameter device + defer_param_init = False + device = canonicalize_device(device) + if device.type == "meta": + defer_param_init = True + + # Initialize parameters if needed + weight = torch.empty( + normalized_shape, + device=device, + dtype=canonicalize_dtype(dtype), + ) + weight = torch.nn.Parameter(weight) + self.weight: torch.nn.Parameter + self.register_parameter("weight", weight) + if not defer_param_init: + self.reset_parameters() + + # Number of SMs to exclude when launching CUDA kernels + self._sm_margins: dict[str, int] + if isinstance(sm_margin, dict): + + def getenv(name: str) -> int: + return int(os.getenv(name, "0")) + + self._sm_margins = { + "forward": sm_margin.get("forward", getenv("NVTE_FWD_LAYERNORM_SM_MARGIN")), + "backward": sm_margin.get("backward", getenv("NVTE_BWD_LAYERNORM_SM_MARGIN")), + "inference": sm_margin.get("inference", getenv("NVTE_INF_LAYERNORM_SM_MARGIN")), + } + else: + + def getenv(name: str) -> int: + return int(os.getenv(name, str(sm_margin))) + + self._sm_margins = { + "forward": getenv("NVTE_FWD_LAYERNORM_SM_MARGIN"), + "backward": getenv("NVTE_BWD_LAYERNORM_SM_MARGIN"), + "inference": getenv("NVTE_INF_LAYERNORM_SM_MARGIN"), + } + + def reset_parameters(self) -> None: + """Initialize parameter buffers and values""" + + # Parameter device + weight = self.weight + device = weight.device + if device.type == "meta": + device = canonicalize_device(None) + + # Initialize param buffers + if not devices_match(weight.device, device): + weight = torch.empty_like(weight, device=device) + + # Initialize values + if self.zero_centered_gamma: + torch.nn.init.zeros_(weight) + else: + torch.nn.init.ones_(weight) + + # Save updated parameter + if not isinstance(weight, torch.nn.Parameter): + weight = torch.nn.Parameter(weight) + self.weight = weight + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: + del prev_op_grad_output_quantizer, basic_op_kwargs + + # Quantized output is not implemented in this fused operation + if next_op_input_quantizer is not None: + raise NotImplementedError( + "RMSNormGated does not support quantized output " + "(only FP32, FP16, or BF16). " + f"Got {type(next_op_input_quantizer).__name__} for the next operation's input." + ) + + weight = self.weight + weight_dims = tuple(weight.size()) + input_dims = tuple(input_.size()) + if len(input_dims) < len(weight_dims) or input_dims[-len(weight_dims) :] != weight_dims: + raise ValueError( + f"Input tensor (shape={input_dims}) " + f"and weight tensor (shape={weight_dims}) are not compatible" + ) + gate = basic_op_extra_inputs[0][0] + gate_dims = tuple(gate.size()) + if gate_dims != input_dims: + raise ValueError( + f"Input tensor (shape={input_dims}) " + f"and gate tensor (shape={gate_dims}) are not compatible" + ) + + inner_dim = math.prod(weight_dims) + dtype = maybe_autocast_dtype(default_dtype=weight.dtype) + x = maybe_dequantize(input_.contiguous(), dtype).view((-1, inner_dim)) + w = maybe_dequantize(self.weight, dtype).view((inner_dim,)) + g = maybe_dequantize(gate.contiguous()).view((-1, inner_dim)) + + # Compute RMSNorm without quantized output + n, _, rstdevs = rmsnorm_fwd( + x, + w, + self.eps, + None, + None, + TE_DType[dtype], + self._sm_margins["forward" if basic_op_ctxs[0].requires_grad else "inference"], + self.zero_centered_gamma, + ) + + # The SiLU gate is computed in FP32, like Megatron's fused implementation + y = (n.float() * torch.nn.functional.silu(g.float())).to(dtype) + + # Save state for backward pass + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + if is_cpu_offload_enabled(): + mark_activation_offload(x, n, g, rstdevs) + ctx.save_for_backward(x, n, g, rstdevs) + ctx.dtype = dtype + + out = y.view(input_dims) + return out, [()] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + del basic_op_grad_extra_outputs + + # Saved tensors from forward pass + ctx = basic_op_ctxs[0] + x, n, g, rstdevs = ctx.saved_tensors + + weight_dims = self.weight.size() + inner_dim = math.prod(weight_dims) + + dtype = ctx.dtype + dy = maybe_dequantize(grad_output.contiguous(), dtype).view(x.size()) + w = maybe_dequantize(self.weight, dtype).view((inner_dim,)) + + # Gradients of the SiLU gate, computed in FP32. Round dn to the input + # dtype so that the RMSNorm backward pass matches Megatron's implementation. + g_fp32 = g.float() + sigmoid = torch.sigmoid(g_fp32) + dy_fp32 = dy.float() + dn = (dy_fp32 * torch.nn.functional.silu(g_fp32)).to(dtype) + silu_grad = sigmoid * (1 + g_fp32 * (1 - sigmoid)) + dg = (dy_fp32 * n.float() * silu_grad).to(g.dtype) + + dx, dw = rmsnorm_bwd( + dn, + x, + rstdevs, + w, + self._sm_margins["backward"], + self.zero_centered_gamma, + ) + + # Clear saved tensors if possible + clear_tensor_data(x) + clear_tensor_data(n) + clear_tensor_data(g) + clear_tensor_data(rstdevs) + + grad_input = dx.view(grad_output.size()) + grad_gate = dg.view(grad_output.size()) + grad_weight = dw.view(weight_dims) + return grad_input, [(grad_weight,)], [(grad_gate,)] From 3df13f1b30e600d2e2e02b8d58a00284768b64d9 Mon Sep 17 00:00:00 2001 From: viking Date: Thu, 22 Oct 2026 03:02:14 -0500 Subject: [PATCH 2/3] [PyTorch] Replace RMSNormGated with a SiLUGate basic op Per review, keep the norm and the gate separate: gated RMSNorm is now Sequential(RMSNorm(...), SiLUGate()). A fused op can be added once there is a fused kernel. Co-Authored-By: Claude Opus 5.5 Signed-off-by: viking --- docs/api/pytorch.rst | 2 +- tests/pytorch/test_fusible_ops.py | 96 +++--- .../pytorch/ops/basic/__init__.py | 2 +- .../pytorch/ops/basic/rmsnorm_gated.py | 279 ------------------ .../pytorch/ops/basic/silu_gate.py | 110 +++++++ 5 files changed, 168 insertions(+), 321 deletions(-) delete mode 100644 transformer_engine/pytorch/ops/basic/rmsnorm_gated.py create mode 100644 transformer_engine/pytorch/ops/basic/silu_gate.py diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 1c65b21bd66..1a05a1e07f8 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -259,7 +259,7 @@ Operation fuser .. autoapiclass:: transformer_engine.pytorch.ops.RMSNorm -.. autoapiclass:: transformer_engine.pytorch.ops.RMSNormGated +.. autoapiclass:: transformer_engine.pytorch.ops.SiLUGate .. autoapiclass:: transformer_engine.pytorch.ops.SReGLU diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index bbee82b86af..9ee1c34df5e 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2332,11 +2332,63 @@ def test_rmsnorm( torch.testing.assert_close(dx_test, x_ref.grad, **tols) torch.testing.assert_close(dw_test, w_ref.grad, **tols) + @pytest.mark.parametrize("in_shape", ((1, 37, 16, 128), (6, 16, 64))) + @pytest.mark.parametrize("dtype", _dtypes) + def test_silu_gate( + self, + *, + in_shape: Iterable[int], + dtype: torch.dtype, + device: torch.device = "cuda", + ) -> None: + """SiLU gate""" + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + ) + g_ref, g_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch implementation. The op computes in FP32 and rounds + # the result to the input dtype, so the reference does the same. + y_ref = (x_ref.float() * torch.nn.functional.silu(g_ref.float())).to(dtype) + y_ref.backward(dy_ref.to(dtype)) + dx_ref = x_ref.grad.to(dtype) + dg_ref = g_ref.grad.to(dtype) + + # Implementation with fusible operation + op = te_ops.SiLUGate() + y_test = op(x_test, g_test) + y_test.backward(dy_test) + + # Expected numerical error + tols = dtype_tols(dtype) + + # Check results + y_test = y_test.to(dtype=torch.float64, device="cpu") + dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") + dg_test = g_test.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(y_test, y_ref.to(dtype=torch.float64, device="cpu"), **tols) + torch.testing.assert_close(dx_test, dx_ref.to(dtype=torch.float64, device="cpu"), **tols) + torch.testing.assert_close(dg_test, dg_ref.to(dtype=torch.float64, device="cpu"), **tols) + @pytest.mark.parametrize("in_shape", ((1, 37, 16, 128), (6, 16, 64))) @pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16)) @pytest.mark.parametrize("zero_centered_gamma", (False, True)) @pytest.mark.parametrize("gate_layout", ("contiguous", "strided")) - def test_rmsnorm_gated( + def test_rmsnorm_silu_gate( self, *, in_shape: Iterable[int], @@ -2406,7 +2458,7 @@ def test_rmsnorm_gated( y_ref.backward(dy_ref) # Implementation with fusible operation - op = te_ops.RMSNormGated( + op = te_ops.RMSNorm( weight_shape, eps=eps, device=device, @@ -2416,7 +2468,8 @@ def test_rmsnorm_gated( with torch.no_grad(): op.weight.copy_(w_test) del w_test - y_test = op(x_test, g_test) + forward = te_ops.Sequential(op, te_ops.SiLUGate()) + y_test = forward(x_test, g_test) y_test.backward(dy_test) # Expected numerical error @@ -2432,43 +2485,6 @@ def test_rmsnorm_gated( torch.testing.assert_close(dw_test, w_ref.grad, **tols) torch.testing.assert_close(dg_test, g_ref.grad, **tols) - @pytest.mark.parametrize("quantization", _quantization_list[1:]) - def test_rmsnorm_gated_quantized_output( - self, - *, - quantization: str, - device: torch.device = "cuda", - in_shape: Iterable[int] = (16, 128), - dtype: torch.dtype = torch.bfloat16, - ) -> None: - """RMSNorm with SiLU gate rejects quantized output""" - - # Skip unsupported BF16 - if not is_bf16_available(): - pytest.skip("BF16 requires sm_80 or higher") - - # Random data - _, x_test = make_reference_and_test_tensors( - in_shape, - test_dtype=dtype, - test_device=device, - ) - _, g_test = make_reference_and_test_tensors( - in_shape, - test_dtype=dtype, - test_device=device, - ) - - # Implementation with fusible operation - op = te_ops.RMSNormGated(in_shape[-1], device=device, dtype=dtype) - forward = te_ops.Sequential( - op, - te_ops.Quantize(forward=True, backward=False), - ) - with te.autocast(enabled=True, recipe=make_recipe(quantization)): - with pytest.raises(NotImplementedError, match="quantized output"): - forward(x_test, g_test) - @pytest.mark.parametrize("in_shape", ((32,), (6, 16, 64), (32, 64))) @pytest.mark.parametrize("dtype", _dtypes) def test_l2normalization( diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 26ab7efa7a4..285f6680b56 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -36,7 +36,7 @@ from .reduce_scatter import ReduceScatter from .reshape import Reshape from .rmsnorm import RMSNorm -from .rmsnorm_gated import RMSNormGated +from .silu_gate import SiLUGate from .swiglu import ( ClampedSwiGLU, ScaledClampedQGeGLU, diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm_gated.py b/transformer_engine/pytorch/ops/basic/rmsnorm_gated.py deleted file mode 100644 index 65eeb2b8823..00000000000 --- a/transformer_engine/pytorch/ops/basic/rmsnorm_gated.py +++ /dev/null @@ -1,279 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Fusable operation for RMSNorm with SiLU gating.""" - -from __future__ import annotations -from collections.abc import Iterable, Sequence -import math -import os -from typing import Any, Optional - -import torch - -from transformer_engine_torch import rmsnorm_bwd, rmsnorm_fwd -from ...constants import TE_DType -from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload -from ...tensor import Quantizer -from ...utils import ( - canonicalize_device, - canonicalize_dtype, - clear_tensor_data, - devices_match, -) -from ..op import BasicOperation, OperationContext -from .._common import maybe_autocast_dtype, maybe_dequantize - - -class RMSNormGated(BasicOperation): - r"""Root Mean Square Layer Normalization with SiLU gate - - Applies Root Mean Square Layer Normalization over a mini-batch of - inputs and multiplies the result by the SiLU activation of a gate - tensor: - - .. math:: - y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * \gamma * \mathrm{SiLU}(g) - - :math:`\gamma` is a learnable affine transform parameter that - matches the inner-most dimensions of the input tensor. The gate - tensor :math:`g` is an extra operation input and must have the - same shape as the input tensor. - - Output quantization is not supported yet. The output is always - returned in FP32, FP16, or BF16. - - Parameters - ---------- - normalized_shape : int or iterable of int - Inner dimensions of input tensor - eps : float, default = 1e-5 - A value added to the denominator for numerical stability - device : torch.device, default = default CUDA device - Tensor device - dtype : torch.dtype, default = default dtype - Tensor datatype - zero_centered_gamma : bool, default = False - If ``True``, the :math:`\gamma` parameter is initialized to zero - and the calculation changes to - - .. math:: - y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) * \mathrm{SiLU}(g) - - sm_margin : int, default = 0 - Number of SMs to exclude when launching CUDA kernels. This - helps overlap with other kernels, e.g. communication kernels. - For more fine-grained control, provide a dict with the SM - margin at each compute stage ("forward", "backward", - "inference"). - - """ - - num_extra_inputs: int = 1 - - def __init__( - self, - normalized_shape: Iterable[int] | int, - *, - eps: float = 1e-5, - device: Optional[torch.device | str] = None, - dtype: Optional[torch.dtype] = None, - zero_centered_gamma: bool = False, - sm_margin: int = 0, - ) -> None: - super().__init__() - self.eps: float = eps - self.zero_centered_gamma: bool = zero_centered_gamma - - # Parameter shape - if not isinstance(normalized_shape, Iterable): - normalized_shape = (normalized_shape,) - else: - normalized_shape = tuple(normalized_shape) - - # Parameter device - defer_param_init = False - device = canonicalize_device(device) - if device.type == "meta": - defer_param_init = True - - # Initialize parameters if needed - weight = torch.empty( - normalized_shape, - device=device, - dtype=canonicalize_dtype(dtype), - ) - weight = torch.nn.Parameter(weight) - self.weight: torch.nn.Parameter - self.register_parameter("weight", weight) - if not defer_param_init: - self.reset_parameters() - - # Number of SMs to exclude when launching CUDA kernels - self._sm_margins: dict[str, int] - if isinstance(sm_margin, dict): - - def getenv(name: str) -> int: - return int(os.getenv(name, "0")) - - self._sm_margins = { - "forward": sm_margin.get("forward", getenv("NVTE_FWD_LAYERNORM_SM_MARGIN")), - "backward": sm_margin.get("backward", getenv("NVTE_BWD_LAYERNORM_SM_MARGIN")), - "inference": sm_margin.get("inference", getenv("NVTE_INF_LAYERNORM_SM_MARGIN")), - } - else: - - def getenv(name: str) -> int: - return int(os.getenv(name, str(sm_margin))) - - self._sm_margins = { - "forward": getenv("NVTE_FWD_LAYERNORM_SM_MARGIN"), - "backward": getenv("NVTE_BWD_LAYERNORM_SM_MARGIN"), - "inference": getenv("NVTE_INF_LAYERNORM_SM_MARGIN"), - } - - def reset_parameters(self) -> None: - """Initialize parameter buffers and values""" - - # Parameter device - weight = self.weight - device = weight.device - if device.type == "meta": - device = canonicalize_device(None) - - # Initialize param buffers - if not devices_match(weight.device, device): - weight = torch.empty_like(weight, device=device) - - # Initialize values - if self.zero_centered_gamma: - torch.nn.init.zeros_(weight) - else: - torch.nn.init.ones_(weight) - - # Save updated parameter - if not isinstance(weight, torch.nn.Parameter): - weight = torch.nn.Parameter(weight) - self.weight = weight - - def fuser_forward( - self, - basic_op_ctxs: list[OperationContext], - input_: torch.Tensor, - *, - basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], - basic_op_kwargs: list[dict[str, Any]], - ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: - del prev_op_grad_output_quantizer, basic_op_kwargs - - # Quantized output is not implemented in this fused operation - if next_op_input_quantizer is not None: - raise NotImplementedError( - "RMSNormGated does not support quantized output " - "(only FP32, FP16, or BF16). " - f"Got {type(next_op_input_quantizer).__name__} for the next operation's input." - ) - - weight = self.weight - weight_dims = tuple(weight.size()) - input_dims = tuple(input_.size()) - if len(input_dims) < len(weight_dims) or input_dims[-len(weight_dims) :] != weight_dims: - raise ValueError( - f"Input tensor (shape={input_dims}) " - f"and weight tensor (shape={weight_dims}) are not compatible" - ) - gate = basic_op_extra_inputs[0][0] - gate_dims = tuple(gate.size()) - if gate_dims != input_dims: - raise ValueError( - f"Input tensor (shape={input_dims}) " - f"and gate tensor (shape={gate_dims}) are not compatible" - ) - - inner_dim = math.prod(weight_dims) - dtype = maybe_autocast_dtype(default_dtype=weight.dtype) - x = maybe_dequantize(input_.contiguous(), dtype).view((-1, inner_dim)) - w = maybe_dequantize(self.weight, dtype).view((inner_dim,)) - g = maybe_dequantize(gate.contiguous()).view((-1, inner_dim)) - - # Compute RMSNorm without quantized output - n, _, rstdevs = rmsnorm_fwd( - x, - w, - self.eps, - None, - None, - TE_DType[dtype], - self._sm_margins["forward" if basic_op_ctxs[0].requires_grad else "inference"], - self.zero_centered_gamma, - ) - - # The SiLU gate is computed in FP32, like Megatron's fused implementation - y = (n.float() * torch.nn.functional.silu(g.float())).to(dtype) - - # Save state for backward pass - ctx = basic_op_ctxs[0] - if ctx.requires_grad: - if is_cpu_offload_enabled(): - mark_activation_offload(x, n, g, rstdevs) - ctx.save_for_backward(x, n, g, rstdevs) - ctx.dtype = dtype - - out = y.view(input_dims) - return out, [()] - - def fuser_backward( - self, - basic_op_ctxs: list[OperationContext], - grad_output: torch.Tensor, - *, - basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], - ) -> tuple[ - torch.Tensor, - Iterable[Iterable[Optional[torch.Tensor]]], - Iterable[Iterable[Optional[torch.Tensor]]], - ]: - del basic_op_grad_extra_outputs - - # Saved tensors from forward pass - ctx = basic_op_ctxs[0] - x, n, g, rstdevs = ctx.saved_tensors - - weight_dims = self.weight.size() - inner_dim = math.prod(weight_dims) - - dtype = ctx.dtype - dy = maybe_dequantize(grad_output.contiguous(), dtype).view(x.size()) - w = maybe_dequantize(self.weight, dtype).view((inner_dim,)) - - # Gradients of the SiLU gate, computed in FP32. Round dn to the input - # dtype so that the RMSNorm backward pass matches Megatron's implementation. - g_fp32 = g.float() - sigmoid = torch.sigmoid(g_fp32) - dy_fp32 = dy.float() - dn = (dy_fp32 * torch.nn.functional.silu(g_fp32)).to(dtype) - silu_grad = sigmoid * (1 + g_fp32 * (1 - sigmoid)) - dg = (dy_fp32 * n.float() * silu_grad).to(g.dtype) - - dx, dw = rmsnorm_bwd( - dn, - x, - rstdevs, - w, - self._sm_margins["backward"], - self.zero_centered_gamma, - ) - - # Clear saved tensors if possible - clear_tensor_data(x) - clear_tensor_data(n) - clear_tensor_data(g) - clear_tensor_data(rstdevs) - - grad_input = dx.view(grad_output.size()) - grad_gate = dg.view(grad_output.size()) - grad_weight = dw.view(weight_dims) - return grad_input, [(grad_weight,)], [(grad_gate,)] diff --git a/transformer_engine/pytorch/ops/basic/silu_gate.py b/transformer_engine/pytorch/ops/basic/silu_gate.py new file mode 100644 index 00000000000..288124e34a6 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/silu_gate.py @@ -0,0 +1,110 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusable operation for SiLU gating.""" + +from __future__ import annotations +from collections.abc import Iterable, Sequence +from typing import Any, Optional + +import torch + +from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload +from ...tensor import Quantizer +from ..op import BasicOperation, OperationContext +from .._common import maybe_autocast_dtype, maybe_dequantize + + +class SiLUGate(BasicOperation): + r"""Multiply input by the SiLU activation of a gate tensor + + This operation takes the main input :math:`x` and an extra gate + tensor :math:`g` and computes + + .. math:: + y = x * \mathrm{SiLU}(g) = \frac{x g}{1 + \exp(-g)} + + The gate tensor is an extra operation input and must have the same + shape as the main input. Computation is performed in FP32 and the + output is returned in the input dtype. + + This operation is typically used after a normalization operation to + implement gated RMSNorm, as in Mamba2 and Qwen3-Next: + + .. code-block:: python + + op = Sequential(RMSNorm(hidden_size), SiLUGate()) + + Output quantization is not supported. + + """ + + num_extra_inputs: int = 1 + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: + # Output quantization is not supported. Return an unquantized + # tensor and leave quantization to the next operation. + del prev_op_grad_output_quantizer, next_op_input_quantizer, basic_op_kwargs + + # Check input tensors + x = maybe_dequantize(input_.contiguous()) + gate = maybe_dequantize(basic_op_extra_inputs[0][0].contiguous()) + if gate.shape != x.shape: + raise ValueError( + f"Input tensor (shape={tuple(x.shape)}) " + f"and gate tensor (shape={tuple(gate.shape)}) are not compatible" + ) + dtype = maybe_autocast_dtype(default_dtype=x.dtype) + x = maybe_dequantize(x, dtype) + + # Compute SiLU gate in FP32 + y = (x.float() * torch.nn.functional.silu(gate.float())).to(dtype) + + # Save state for backward pass + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + if is_cpu_offload_enabled(): + mark_activation_offload(x, gate) + ctx.save_for_backward(x, gate) + + return y, [()] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + del basic_op_grad_extra_outputs + + # Saved tensors from forward pass + ctx = basic_op_ctxs[0] + x, gate = ctx.saved_tensors + + dy = maybe_dequantize(grad_output.contiguous(), x.dtype) + + # Gradients of the SiLU gate, computed in FP32 + x_fp32 = x.float() + gate_fp32 = gate.float() + dy_fp32 = dy.float() + sigmoid = torch.sigmoid(gate_fp32) + silu_grad = sigmoid * (1 + gate_fp32 * (1 - sigmoid)) + dx = (dy_fp32 * torch.nn.functional.silu(gate_fp32)).to(x.dtype) + dgate = (dy_fp32 * x_fp32 * silu_grad).to(gate.dtype) + + return dx, [()], [(dgate,)] From 6b9da2d0ee1473c487baa74a0140c3fe3559fc41 Mon Sep 17 00:00:00 2001 From: viking Date: Sat, 24 Oct 2026 10:55:43 -0500 Subject: [PATCH 3/3] [PyTorch] Cover negative gates in SiLUGate tests Draw gates from [-6, 4) so both tests exercise the negative side of SiLU, including the zero of its derivative near -1.28. Co-Authored-By: Claude Opus 5.5 Signed-off-by: viking --- tests/pytorch/test_fusible_ops.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 9ee1c34df5e..050cef69eed 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2349,8 +2349,11 @@ def test_silu_gate( test_dtype=dtype, test_device=device, ) + # Signed gates cover SiLU's negative side, where its derivative changes sign g_ref, g_test = make_reference_and_test_tensors( in_shape, + min=-6.0, + max=4.0, test_dtype=dtype, test_device=device, ) @@ -2430,6 +2433,8 @@ def test_rmsnorm_silu_gate( if gate_layout == "contiguous": g_ref, g_test = make_reference_and_test_tensors( in_shape, + min=-6.0, + max=4.0, test_dtype=dtype, test_device=device, )