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)) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index a974d41ef9..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,11 +414,7 @@ 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." - ) + self._maybe_warn_activation_recompute_in_mlp() extra_input = basic_op_extra_inputs[0][0] @@ -445,12 +453,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) @@ -481,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(