Relax runtime checks for activation recompute into Warnings

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))

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(