diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 31d469cb77f..1a05a1e07f8 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.SiLUGate + .. 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..050cef69eed 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2332,6 +2332,164 @@ 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, + ) + # 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, + ) + 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_silu_gate( + 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, + min=-6.0, + max=4.0, + 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.RMSNorm( + 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 + forward = te_ops.Sequential(op, te_ops.SiLUGate()) + y_test = forward(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("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..285f6680b56 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 .silu_gate import SiLUGate from .swiglu import ( ClampedSwiGLU, ScaledClampedQGeGLU, 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,)]