From a4d79a71f870f8ab9bef08cf27701b4de68a8157 Mon Sep 17 00:00:00 2001 From: Ravi Ghadia Date: Thu, 27 Aug 2026 17:19:00 -0700 Subject: [PATCH 1/3] Remove redundant runtime checks for activation recompute in MLP from _ScaledUnary class in activation.py Signed-off-by: Ravi Ghadia --- transformer_engine/pytorch/ops/basic/activation.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index a974d41ef9..7530825c40 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -402,12 +402,6 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: - if self.activation_recompute_in_mlp: - raise RuntimeError( - f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " - "fused grouped MLP path." - ) - extra_input = basic_op_extra_inputs[0][0] if torch.is_autocast_enabled(): @@ -445,12 +439,6 @@ def fuser_backward( ]: del basic_op_grad_extra_outputs - if self.activation_recompute_in_mlp: - raise RuntimeError( - f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " - "fused grouped MLP path." - ) - ctx = basic_op_ctxs[0] x, scales = ctx.saved_tensors x = maybe_dequantize(x.contiguous(), ctx.dtype) From d77372b8a128cd73947dae6a6c0009d2e7b3c5bb Mon Sep 17 00:00:00 2001 From: Ravi Ghadia Date: Mon, 31 Aug 2026 21:55:16 -0700 Subject: [PATCH 2/3] Add warning for activation recompute in MLP outside fused path in _ScaledUnary class Signed-off-by: Ravi Ghadia --- .../pytorch/ops/basic/activation.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 7530825c40..91d5d1339d 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -8,6 +8,7 @@ import abc from collections.abc import Iterable, Sequence from typing import Any, Optional +import warnings import torch @@ -356,6 +357,17 @@ class _ScaledUnary(BasicOperation, metaclass=abc.ABCMeta): def __init__(self, *, activation_recompute_in_mlp: bool = False) -> None: super().__init__() self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp + self._warned_activation_recompute_in_mlp: bool = False + + def _maybe_warn_activation_recompute_in_mlp(self) -> None: + """Warn if activation recompute is requested outside the fused grouped MLP.""" + if not self.activation_recompute_in_mlp or self._warned_activation_recompute_in_mlp: + return + self._warned_activation_recompute_in_mlp = True + warnings.warn( + f"{self.__class__.__name__}(activation_recompute_in_mlp=True) is only supported " + "in the fused grouped MLP path." + ) @abc.abstractmethod def _scaled_unary_forward( @@ -402,6 +414,8 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]: + self._maybe_warn_activation_recompute_in_mlp() + extra_input = basic_op_extra_inputs[0][0] if torch.is_autocast_enabled(): @@ -469,7 +483,9 @@ class ScaledSReLU(_ScaledUnary): ---------- activation_recompute_in_mlp : bool, default = ``False`` Enable fused grouped MLP kernels to recompute activation outputs - during backward when supported instead of saving them. + during backward when supported instead of saving them. Outside the + fused grouped MLP path this option has no effect and a warning is + emitted. """ def _scaled_unary_forward( From b3ae0cbb38221c4a10e90c95f6ffcbb02664d9c8 Mon Sep 17 00:00:00 2001 From: Ravi Ghadia Date: Tue, 1 Sep 2026 11:18:59 -0700 Subject: [PATCH 3/3] Add test for Scaled SReLU activation recompute warning outside fused MLP path Signed-off-by: Ravi Ghadia --- tests/pytorch/test_fusible_ops.py | 62 +++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index c20eb7afa7..15db957ef1 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -3128,6 +3128,68 @@ def test_scaled_srelu_activation_recompute_in_mlp_config(self) -> None: assert te_ops.ScaledSReLU().activation_recompute_in_mlp is False assert te_ops.ScaledSReLU(activation_recompute_in_mlp=True).activation_recompute_in_mlp + def test_scaled_srelu_activation_recompute_in_mlp_warns_outside_fused_mlp( + self, + *, + in_shape: Iterable[int] = (71, 192), + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + ) -> None: + """Scaled SReLU warns, but still runs, when recompute is unavailable""" + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + requires_grad=True, + ) + scales_ref, scales_test = make_reference_and_test_tensors( + in_shape[:-1], + test_dtype=dtype, + test_device=device, + requires_grad=True, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch implementation + y_ref = scales_ref.unsqueeze(-1) * torch.nn.functional.relu(x_ref).square() + y_ref.backward(dy_ref) + + # Activation recompute is only honored within the fused grouped MLP, + # so the standalone op falls back to saving the activation input. + op = te_ops.ScaledSReLU(activation_recompute_in_mlp=True) + with pytest.warns(UserWarning, match="fused grouped MLP path") as warning_log: + y_test = op(x_test, scales_test) + y_test.backward(dy_test) + + # Gradients accumulate, so snapshot them before the second pass + dx_test = x_test.grad.detach().clone() + dscales_test = scales_test.grad.detach().clone() + + # Second pass through the same op + op(x_test, scales_test).backward(dy_test) + warnings_seen = [ + warning + for warning in warning_log + if "activation_recompute_in_mlp" in str(warning.message) + ] + + # Warning is emitted at most once per op + assert len(warnings_seen) == 1 + + # Check results + tols = dtype_tols(dtype) + y_test = y_test.to(dtype=torch.float64, device="cpu") + assert_close(y_test, y_ref, **tols) + assert_close(dx_test, x_ref.grad, **tols) + assert_close(dscales_test, scales_ref.grad, **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))