diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 0083fd9c0f..6cd1fc3065 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -3101,6 +3101,101 @@ def test_scaled_srelu( assert_close_grads(x_test, x_ref, **tols) assert_close_grads(scales_test, scales_ref, **tols) + @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128))) + @pytest.mark.parametrize("input_requires_grad", (False, True)) + @pytest.mark.parametrize("scales_requires_grad", (False, True)) + @pytest.mark.parametrize("tanh_clamp_scale", (0.5, 2.0)) + def test_scaled_tanh_srelu( + self, + *, + in_shape: Iterable[int], + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + input_requires_grad: bool, + scales_requires_grad: bool, + tanh_clamp_scale: float, + ) -> None: + """Tanh soft-clamped SReLU with post-scale. + + Covers the unfused path specifically: the fused grouped-MLP op goes straight + to the cuDNN srelu_tanh epilogue and never runs this code. Small clamp scales + are used so tanh genuinely saturates -- with a large scale the result is + numerically indistinguishable from plain ScaledSReLU. + """ + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + requires_grad=input_requires_grad, + ) + scales_ref, scales_test = make_reference_and_test_tensors( + in_shape[:-1], + test_dtype=dtype, + test_device=device, + requires_grad=scales_requires_grad, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch implementation. Autograd supplies the reference gradients, so + # the op's hand-written backward is checked against a derivative it played no + # part in computing. + y = ( + tanh_clamp_scale * torch.tanh(torch.nn.functional.relu(x_ref) / tanh_clamp_scale) + ).square() + y_ref = scales_ref.unsqueeze(-1) * y + if input_requires_grad or scales_requires_grad: + y_ref.backward(dy_ref) + + # Implementation with fusible operation + op = te_ops.ScaledTanhSReLU(tanh_clamp_scale=tanh_clamp_scale) + y_test = op(x_test, scales_test) + if input_requires_grad or scales_requires_grad: + y_test.backward(dy_test) + + # Check results + tols = dtype_tols(dtype) + y_test = y_test.to(dtype=torch.float64, device="cpu") + assert_close(y_test, y_ref, **tols) + if input_requires_grad: + assert_close_grads(x_test, x_ref, **tols) + if scales_requires_grad: + assert_close_grads(scales_test, scales_ref, **tols) + + def test_scaled_tanh_srelu_saturates(self) -> None: + """Large inputs pin the output at tanh_clamp_scale**2, unlike plain SReLU.""" + s = 2.0 + x = torch.full((4, 8), 1.0e3, device="cuda", dtype=torch.float32) + scales = torch.ones((4,), device="cuda", dtype=torch.float32) + + y = te_ops.ScaledTanhSReLU(tanh_clamp_scale=s)(x, scales) + torch.testing.assert_close(y, torch.full_like(y, s * s)) + + # Plain SReLU on the same input is ~250000x larger, so this cannot pass by + # accident if the clamp were silently dropped. + y_unclamped = te_ops.ScaledSReLU()(x, scales) + assert y_unclamped.min().item() > 1.0e5 + + @pytest.mark.parametrize("tanh_clamp_scale", (0.0, -1.0, float("inf"), float("nan"))) + def test_scaled_tanh_srelu_rejects_bad_clamp_scale(self, tanh_clamp_scale) -> None: + """The clamp scale must be finite and positive.""" + with pytest.raises(ValueError, match="tanh_clamp_scale"): + te_ops.ScaledTanhSReLU(tanh_clamp_scale=tanh_clamp_scale) + + def test_scaled_tanh_srelu_activation_recompute_in_mlp_config(self) -> None: + """Tanh SReLU exposes the same activation recompute knob as ScaledSReLU.""" + op = te_ops.ScaledTanhSReLU(tanh_clamp_scale=2.0) + assert op.activation_recompute_in_mlp is False + assert te_ops.ScaledTanhSReLU( + tanh_clamp_scale=2.0, activation_recompute_in_mlp=True + ).activation_recompute_in_mlp + def test_interleaved_scaled_swiglu(self): """SwiGLU with post-scale and block interleaved input format""" self.test_scaled_swiglu( diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 5049f58ce4..d48f7afae6 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -57,6 +57,12 @@ mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +# Soft-clamp scale for ScaledTanhSReLU coverage. Deliberately small relative to the +# FC1 outputs these tests produce, so tanh actually saturates -- a large scale would +# make the activation numerically indistinguishable from plain ScaledSReLU and the +# test would pass even if the clamp were dropped. +_TANH_SRELU_CLAMP_SCALE: float = 2.0 + # Supported data types _dtypes: list[torch.dtype] = [torch.float32, torch.float16] if is_bf16_available(): # bf16 requires sm_80 or higher @@ -1072,6 +1078,7 @@ class TestGroupedMLPFusedOp: "scaled_clamped_qgeglu", "scaled_clamped_qgeglu_custom", "scaled_srelu", + "scaled_tanh_srelu", ), ) def test_grouped_mlp( @@ -1140,12 +1147,25 @@ def test_grouped_mlp( pytest.skip("Unary activations do not use GLU interleaving") if quantization == "nvfp4_4over6": pytest.skip("NVFP4 4over6 grouped quantization is not supported") - if activation == "scaled_srelu" and quantization in ("nvfp4", "nvfp4_rht") and bias: + if ( + activation in ("scaled_srelu", "scaled_tanh_srelu") + and quantization in ("nvfp4", "nvfp4_rht") + and bias + ): pytest.skip("NVFP4 SReLU grouped MLP coverage is limited to no-bias") if quantization == "nvfp4_rht": if activation == "scaled_swiglu" and (bias or glu_interleave_size != 32): pytest.skip("NVFP4 RHT SwiGLU grouped MLP coverage is limited to no-bias") - if activation not in ("scaled_swiglu", "scaled_situglu", "scaled_srelu"): + # tanh-SReLU is included deliberately: NVFP4 is the only path that exercises + # the dsrelu-family fc2 alpha (full product rather than sqrt), so excluding + # it here would leave that branch untested. It runs without the hadamard + # sub-kernel, which the fused op handles via the generic quantize fallback. + if activation not in ( + "scaled_swiglu", + "scaled_situglu", + "scaled_srelu", + "scaled_tanh_srelu", + ): pytest.skip( "NVFP4 RHT grouped MLP coverage is limited to SwiGLU, SiTU-GLU, and SReLU" ) @@ -1271,6 +1291,11 @@ def _apply_activation(x: torch.Tensor) -> torch.Tensor: return (x2c + geglu_offset) * (x1c * torch.sigmoid(geglu_alpha * x1c)) if activation == "scaled_srelu": return torch.nn.functional.relu(x).square() + if activation == "scaled_tanh_srelu": + s = _TANH_SRELU_CLAMP_SCALE + return ( + (s * torch.tanh(torch.nn.functional.relu(x.float()) / s)).square().to(x.dtype) + ) raise ValueError(f"Unexpected grouped MLP activation ({activation})") # Reference implementation @@ -1314,6 +1339,8 @@ def _make_scaled_act(): return te.ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave_size) if activation == "scaled_srelu": return te.ops.ScaledSReLU() + if activation == "scaled_tanh_srelu": + return te.ops.ScaledTanhSReLU(tanh_clamp_scale=_TANH_SRELU_CLAMP_SCALE) raise ValueError(f"Unexpected grouped MLP activation ({activation})") def _make_module(): @@ -1400,15 +1427,20 @@ def _make_module(): fc2.backward_dw() # Check for expected fusions - cudnn_frontend_supports_grouped_mlp = ( - grouped_mlp_module._cudnn_frontend_supports_grouped_gemm_situglu() - if activation == "scaled_situglu" - else ( + if activation == "scaled_situglu": + cudnn_frontend_supports_grouped_mlp = ( + grouped_mlp_module._cudnn_frontend_supports_grouped_gemm_situglu() + ) + elif activation == "scaled_srelu": + cudnn_frontend_supports_grouped_mlp = _cudnn_frontend_supports_grouped_gemm_srelu() + elif activation == "scaled_tanh_srelu": + # Needs both the base srelu kernels and the tanh_clamp_scale parameter. + cudnn_frontend_supports_grouped_mlp = ( _cudnn_frontend_supports_grouped_gemm_srelu() - if activation == "scaled_srelu" - else _cudnn_frontend_version_supported() + and grouped_mlp_module._cudnn_frontend_supports_grouped_gemm_srelu_tanh() ) - ) + else: + cudnn_frontend_supports_grouped_mlp = _cudnn_frontend_version_supported() expected_grouped_mlp_fusion = cudnn_frontend_supports_grouped_mlp and ( ( quantization == "mxfp8" @@ -1908,6 +1940,7 @@ def _run(grad): "scaled_clamped_qgeglu", "scaled_clamped_qgeglu_custom", "scaled_srelu", + "scaled_tanh_srelu", ), ) def test_grouped_mlp_fp16( diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 4eb2796b2c..dae15330c1 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -14,6 +14,7 @@ ReGLU, SReLU, ScaledSReLU, + ScaledTanhSReLU, SReGLU, SiLU, ) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index a974d41ef9..5c33c08b44 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -6,6 +6,7 @@ from __future__ import annotations import abc +import math from collections.abc import Iterable, Sequence from typing import Any, Optional @@ -29,6 +30,7 @@ "ReGLU", "SReLU", "ScaledSReLU", + "ScaledTanhSReLU", "SReGLU", "SiLU", ] @@ -508,6 +510,88 @@ def _scaled_unary_backward( ) +class ScaledTanhSReLU(_ScaledUnary): + r"""Tanh soft-clamped squared ReLU with per-row post-scaling. + + Identical to :class:`ScaledSReLU` except that the ReLU output is soft-clamped + by a tanh before squaring, which bounds the activation by + ``tanh_clamp_scale ** 2``: + + .. math:: + y = \left( s \cdot \tanh\left( \frac{\mathrm{relu}(x)}{s} \right) \right)^2 \cdot \mathrm{scales} + + "Tanh" rather than "Clamped" is deliberate: in these ops ``Clamped`` already + means a hard min/max clamp (see :class:`ScaledClampedQGeGLU`), whereas this is + a smooth saturating clamp. + + Parameters + ---------- + tanh_clamp_scale : float + The soft-clamp scale ``s``; must be finite and positive. Required, since + an unclamped tanh-SReLU is just :class:`ScaledSReLU`. + activation_recompute_in_mlp : bool, default = ``False`` + Enable fused grouped MLP kernels to recompute activation outputs + during backward when supported instead of saving them. + """ + + def __init__( + self, + *, + tanh_clamp_scale: float, + activation_recompute_in_mlp: bool = False, + ) -> None: + super().__init__(activation_recompute_in_mlp=activation_recompute_in_mlp) + self.tanh_clamp_scale: float = float(tanh_clamp_scale) + if not math.isfinite(self.tanh_clamp_scale) or self.tanh_clamp_scale <= 0.0: + raise ValueError( + f"tanh_clamp_scale must be finite and positive, got {tanh_clamp_scale}" + ) + + def _tanh_srelu_terms(self, input_: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``b = s*tanh(relu(x)/s)`` and ``t``, in fp32. + + Matches the fused kernels, which evaluate the whole epilogue in fp32 and + round once at the store rather than at each step. + """ + s = self.tanh_clamp_scale + t = torch.tanh(torch.relu(input_.float()) / s) + return s * t, t + + def _scaled_unary_forward( + self, + input_: torch.Tensor, + scales: torch.Tensor, + ) -> torch.Tensor: + # No fused CUDA kernel for this activation yet, so the unfused path is + # expressed in torch. It only serves non-SM100 / non-fused configurations; + # the fused grouped-MLP path goes straight to the cuDNN srelu_tanh epilogue. + b, _ = self._tanh_srelu_terms(input_) + out = b.square() * scales.float().unsqueeze(-1) + return out.to(input_.dtype) + + def _scaled_unary_backward( + self, + grad_output: torch.Tensor, + input_: torch.Tensor, + scales: torch.Tensor, + *, + compute_scale_grad: bool, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + # d/dx (b^2) = 2*b*(1 - t^2); the scale multiplies it, exactly as the + # scaled-unary CUDA kernel does for the plain squared ReLU. + b, t = self._tanh_srelu_terms(input_) + grad_output_f32 = grad_output.float() + grad_input = grad_output_f32 * scales.float().unsqueeze(-1) * (2.0 * b * (1.0 - t * t)) + + grad_scales = None + if compute_scale_grad: + # Gradient wrt the per-row scale is sum over the last dim of the + # unscaled activation times the incoming gradient. + grad_scales = (b.square() * grad_output_f32).sum(dim=-1).to(scales.dtype) + + return grad_input.to(input_.dtype), grad_scales + + class SReGLU(_ActivationOperation): r"""Squared Rectified Gated Linear Unit diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index dca1119800..61f80b9d9f 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -47,6 +47,7 @@ ScaledSiTUGLU, ScaledSReLU, ScaledSwiGLU, + ScaledTanhSReLU, ) from ..fuser import register_forward_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext @@ -98,6 +99,34 @@ def _cudnn_frontend_supports_grouped_gemm_srelu_hadamard() -> bool: return _cudnn_frontend_version_at_least("1.26.0") +@functools.lru_cache(maxsize=None) +def _cudnn_frontend_supports_grouped_gemm_srelu_tanh() -> bool: + """Feature-detect complete cuDNN frontend grouped tanh-SReLU support. + + Both directions are required: a frontend with only the forward clamp would + train against an unclamped backward. Detected by signature rather than + version so this can be developed against an editable cuDNN FE checkout; a + min-version constant can replace it once the feature is in a release. + """ + try: + from cudnn import ( # pylint: disable=import-outside-toplevel + grouped_gemm_dsrelu_wrapper_sm100, + grouped_gemm_srelu_wrapper_sm100, + ) + except ImportError: + return False + try: + wrappers = ( + grouped_gemm_srelu_wrapper_sm100, + grouped_gemm_dsrelu_wrapper_sm100, + ) + return all( + "tanh_clamp_scale" in inspect.signature(wrapper).parameters for wrapper in wrappers + ) + except (TypeError, ValueError): + return False + + @functools.lru_cache(maxsize=None) def _cudnn_frontend_supports_grouped_gemm_situglu() -> bool: """Feature-detect complete cuDNN frontend grouped SiTU-GLU support.""" @@ -132,9 +161,13 @@ def _cudnn_frontend_supports_single_group_runtime_offsets( activation_type: type[FusibleOperation], ) -> bool: """Check cuDNN FE support for single-group runtime offsets.""" - return not issubclass(activation_type, ScaledSReLU) and _cudnn_frontend_version_at_least( - "1.27.0" - ) + # The srelu/dsrelu wrappers take no use_single_group_runtime_offsets argument, + # so every activation in that family has to be excluded here, not just + # ScaledSReLU -- passing it through would raise TypeError on the single-group + # (shared expert) path. + return not issubclass( + activation_type, (ScaledSReLU, ScaledTanhSReLU) + ) and _cudnn_frontend_version_at_least("1.27.0") def _wrap_single_quantized_as_grouped( @@ -794,7 +827,7 @@ def validate_grouped_mlp_dims(fc1, activation_op, fc2) -> None: ) if is_glu_activation(activation_op): expected_fc1_out_features = 2 * fc2.in_features - elif isinstance(activation_op, ScaledSReLU): + elif isinstance(activation_op, (ScaledSReLU, ScaledTanhSReLU)): expected_fc1_out_features = fc2.in_features else: raise TypeError(f"Unsupported grouped MLP activation ({activation_op.__class__.__name__}).") @@ -1008,6 +1041,21 @@ def __init__( self._cudnn_situ_beta1: float = activation.beta1 self._cudnn_situ_beta2: float = activation.beta2 + # Set unconditionally: the forward/backward paths read this attribute for + # every activation, so leaving it undefined would break plain ScaledSReLU. + self._pass_srelu_tanh_params: bool = isinstance(activation, ScaledTanhSReLU) + if self._pass_srelu_tanh_params: + # Fail at construction rather than silently running unclamped, which + # would train a different model than the config asks for. + if not _cudnn_frontend_supports_grouped_gemm_srelu_tanh(): + raise RuntimeError( + "ScaledTanhSReLU requires a cuDNN frontend whose " + "grouped_gemm_srelu_wrapper_sm100 and grouped_gemm_dsrelu_wrapper_sm100 " + "accept tanh_clamp_scale. The installed frontend does not, and running " + "without it would silently apply an unclamped squared ReLU." + ) + self._cudnn_tanh_clamp_scale: float = activation.tanh_clamp_scale + def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -1386,6 +1434,10 @@ def fuser_forward( and fc2_input_quantizer.with_rht and fc2_input_quantizer.with_post_rht_amax ) + # Deliberately ScaledSReLU only: the hadamard kernel this selects is the GLU + # one, which has no soft-clamp support, so ScaledTanhSReLU must not reach it. + # The cost is that NVFP4 RHT gives up hadamard fusion for tanh-SReLU -- a + # performance limitation, not a correctness one. activation_is_srelu = isinstance(activation_op, ScaledSReLU) activation_supports_hadamard = self._cudnn_act_func in ("swiglu", "situglu") or ( activation_is_srelu and _cudnn_frontend_supports_grouped_gemm_srelu_hadamard() @@ -1434,6 +1486,8 @@ def fuser_forward( situ_beta1=self._cudnn_situ_beta1, situ_beta2=self._cudnn_situ_beta2, ) + if self._pass_srelu_tanh_params: + fc1_activation_kwargs.update(tanh_clamp_scale=self._cudnn_tanh_clamp_scale) if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM. @@ -1749,6 +1803,10 @@ def fuser_forward( mark_grouped_tensor(saved_fc1_x, activation_in, scales, grouped_fc2_x) activation_op = self.basic_ops[1] cpu_offloading = is_cpu_offload_enabled() + # Deliberately ScaledSReLU only for now: ScaledTanhSReLU falls back to + # saving fc2_x, which costs memory but stays correct. The cuDNN dsrelu + # d_srelu regeneration does honour the clamp, so enabling recompute here + # is a viable follow-up rather than a blocker. activation_is_srelu = isinstance(activation_op, ScaledSReLU) activation_recompute_in_mlp = bool( getattr(activation_op, "activation_recompute_in_mlp", False) @@ -1835,7 +1893,9 @@ def fuser_backward( # Get basic operations fc1_op, activation_op, fc2_op = self.basic_ops - activation_is_srelu = isinstance(activation_op, ScaledSReLU) + # Selects how the NVFP4 fc2 alpha is folded below: the whole dsrelu family + # applies alpha once, unlike the gated kernels which need sqrt(product). + activation_is_srelu = isinstance(activation_op, (ScaledSReLU, ScaledTanhSReLU)) fc1_ctx, _activation_ctx, fc2_ctx = basic_op_ctxs # Tensor properties @@ -2098,6 +2158,8 @@ def fuser_backward( situ_beta1=self._cudnn_situ_beta1, situ_beta2=self._cudnn_situ_beta2, ) + if self._pass_srelu_tanh_params: + fc2_dactivation_kwargs.update(tanh_clamp_scale=self._cudnn_tanh_clamp_scale) fc2_leader = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 if is_distributed_weight(fc2_leader): @@ -2629,13 +2691,21 @@ def fuse_srelu_ops( recipe: Optional[Recipe] = None, **unused, # pylint: disable=unused-argument ) -> list[FusibleOperation]: - """Apply joint GroupedLinear + ScaledSReLU + GroupedLinear fusion.""" + """Apply joint GroupedLinear + scaled unary activation + GroupedLinear fusion.""" + + # ScaledTanhSReLU joins only when the installed cuDNN frontend can actually + # clamp. Listing it unconditionally would let the op fuse and then raise from + # _GroupedMLP_CuTeGEMMBase.__init__; leaving it out simply declines the fusion + # and runs the correct unfused activation instead. + activation_op_types: tuple[type[FusibleOperation], ...] = (ScaledSReLU,) + if _cudnn_frontend_supports_grouped_gemm_srelu_tanh(): + activation_op_types += (ScaledTanhSReLU,) return fuse_grouped_mlp_ops( ops, recipe=recipe, fused_op_cls=GroupedMLP_CuTeGEMMUnary, - activation_op_types=(ScaledSReLU,), + activation_op_types=activation_op_types, )