Use cuDNN's deterministic dprob in the fused grouped MLP

tests/pytorch/test_grouped_mlp.py

@@ -2772,6 +2772,181 @@ def train_step(
                 assert_close(graph_grad, param.grad, **tols)
 
 
+class TestGroupedMLPDeterminism:
+    """Determinism coverage for the CuTe DSL fused grouped MLP.
+
+    cuDNN's grouped-GEMM dactivation epilogue accumulates the scale gradient (``dprob``)
+    with cross-CTA atomic adds, so its summation order follows the tile scheduler. The
+    dSReLU wrapper takes a ``deterministic`` argument from cuDNN frontend 1.28.0 on that
+    replaces those atomics with per-subtile slots reduced in a canonical order. The dGLU
+    wrapper has no equivalent, so a determinism request it cannot honor is an error.
+    """
+
+    @pytest.fixture
+    def _restore_torch_determinism(self):
+        """``use_deterministic_algorithms`` is process-global, so put it back."""
+        previous = torch.are_deterministic_algorithms_enabled()
+        yield
+        torch.use_deterministic_algorithms(previous)
+
+    @pytest.mark.parametrize(
+        "allow_nondeterministic,torch_flag,expected",
+        (
+            (None, False, False),  # default: non-deterministic algorithms are allowed
+            ("1", False, False),
+            ("0", False, True),  # the TE variable alone
+            (None, True, True),  # the torch flag alone, which TE must not ignore
+            ("1", True, True),  # ... including when the TE variable says otherwise
+            ("0", True, True),
+        ),
+    )
+    def test_either_knob_requests_determinism(
+        self,
+        monkeypatch,
+        _restore_torch_determinism,
+        *,
+        allow_nondeterministic: Optional[str],
+        torch_flag: bool,
+        expected: bool,
+    ) -> None:
+        """The same union DotProductAttention uses, and uncached so both stay live.
+
+        ``NVTE_ALLOW_NONDETERMINISTIC_ALGO=1`` is not a request for non-determinism, only
+        the absence of one, so the torch flag still wins that row.
+        """
+        if allow_nondeterministic is None:
+            monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False)
+        else:
+            monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", allow_nondeterministic)
+        torch.use_deterministic_algorithms(torch_flag)
+        assert grouped_mlp_module._deterministic_algorithms_required() is expected
+
+    def test_only_the_srelu_path_can_be_deterministic(self) -> None:
+        """The capability is a property of the wrapper, not of the environment.
+
+        Needs no GPU, and no cuDNN either: the SReLU probe answers False when the import
+        fails, so both assertions hold anywhere.
+        """
+        glu = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU
+        unary = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary
+        assert glu.grouped_gemm_dactivation_is_deterministic() is False
+        assert isinstance(unary.grouped_gemm_dactivation_is_deterministic(), bool)
+
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    @pytest.mark.parametrize("activation", ("scaled_swiglu", "scaled_srelu"))
+    def test_determinism_either_runs_or_refuses(self, monkeypatch, *, activation: str) -> None:
+        """A request TE cannot honor must fail loudly, not train on silently.
+
+        ``group_size > 1`` here, so the unit-activation-scale shortcut is off and the cuDNN
+        epilogue really does produce ``dprob``. When the request *can* be honored, the
+        wrapped test checks numerics as usual.
+        """
+        if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
+            pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
+
+        monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0")
+        fused_cls = (
+            grouped_mlp_module.GroupedMLP_CuTeGEMMUnary
+            if activation == "scaled_srelu"
+            else grouped_mlp_module.GroupedMLP_CuTeGEMMGLU
+        )
+
+        def _run() -> None:
+            TestGroupedMLPFusedOp().test_grouped_mlp(
+                bias=False,
+                hidden_size=128,
+                quantization="mxfp8",
+                single_grouped_weight=False,
+                activation=activation,
+            )
+
+        if fused_cls.grouped_gemm_dactivation_is_deterministic():
+            _run()
+        else:
+            with pytest.raises(RuntimeError, match="dprob"):
+                _run()
+
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None:
+        """The actual claim: identical inputs produce an identical ``dprob``.
+
+        The end-to-end test above cannot see this. Reordering the same atomic adds moves
+        the result by about an ulp, which every tolerance in this file accepts, so a run
+        that is silently not reproducible passes it. Only comparing two runs exactly can
+        tell the difference.
+        """
+        fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary
+        if not fused_cls.is_supported():
+            pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
+        if not fused_cls.grouped_gemm_dactivation_is_deterministic():
+            pytest.skip("dSReLU determinism needs cuDNN frontend 1.28.0 or later")
+
+        monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0")
+
+        device = torch.device("cuda")
+        dtype = torch.bfloat16
+        group_size = 4
+        # Wide enough that the dprob reduction spans several 256-wide N-tiles. With a
+        # single tile there is one writer per token, nothing to reorder, and the test is
+        # vacuous.
+        hidden_size = 1024
+        split_sizes = torch.tensor([256] * group_size, dtype=torch.int, device=device)
+        num_tokens = int(split_sizes.sum().item())
+
+        recipe = make_recipe("mxfp8")
+        tensor_kwargs = {
+            "min": -0.25,
+            "max": 0.25,
+            "quantization": "mxfp8",
+            "test_dtype": dtype,
+            "test_device": device,
+        }
+        _, x = make_reference_and_test_tensors((num_tokens, hidden_size), **tensor_kwargs)
+        _, dy = make_reference_and_test_tensors(
+            (num_tokens, hidden_size), requires_grad=False, **tensor_kwargs
+        )
+        _, probs = make_reference_and_test_tensors(
+            (num_tokens,), test_dtype=dtype, test_device=device
+        )
+
+        # No bias: with an FC2 scale_bias the scale gradient is finished by the Triton
+        # grouped-dbias kernel instead of arriving straight from the cuDNN epilogue, and
+        # probs.grad would no longer be the dprob under test.
+        with te.quantized_model_init(enabled=True, recipe=recipe):
+            module = te.ops.Sequential(
+                te.ops.GroupedLinear(
+                    group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype
+                ),
+                te.ops.ScaledSReLU(),
+                te.ops.GroupedLinear(
+                    group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype
+                ),
+            )
+
+        def _run() -> torch.Tensor:
+            x.grad = None
+            probs.grad = None
+            with te.autocast(enabled=True, recipe=recipe):
+                y = module(x, split_sizes, probs, split_sizes)
+            y.backward(dy)
+            return probs.grad.detach().clone()
+
+        first = _run()
+        # The fusion has to have happened, or dprob never came from the cuDNN epilogue and
+        # the comparison below proves nothing.
+        forward_ops = module._module_groups[0]._forward_ops
+        assert len(forward_ops) == 1
+        assert isinstance(forward_ops[0][0], fused_cls)
+        second = _run()
+
+        # Exact, not assert_close. Weight gradients are deliberately not compared: the CuTe
+        # DSL wgrad kernel has its own K-split atomics that this change does not address.
+        assert torch.equal(first, second), (
+            "dprob differed between two identical runs under determinism; max |delta| ="
+            f" {(first.float() - second.float()).abs().max().item()}"
+        )
+
+
 def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None:
     if not mxfp8_available:
         pytest.skip(reason_for_no_mxfp8)

transformer_engine/pytorch/ops/fused/grouped_mlp.py

@@ -137,6 +137,24 @@ def _cudnn_frontend_supports_single_group_runtime_offsets(
     )
 
 
+def _deterministic_algorithms_required() -> bool:
+    """Whether the user has asked for bit-exact reproducibility, by either route.
+
+    Same union as ``DotProductAttention``: the two knobs answer different questions.
+    ``NVTE_ALLOW_NONDETERMINISTIC_ALGO`` is set once in a job launcher, applies uniformly
+    across ranks and is the only one TE's C++ layer can read; ``use_deterministic_algorithms``
+    is the framework standard, is togglable at runtime, and is what a user who wants
+    reproducibility usually reaches for. Honoring only the first would leave the second
+    silently unenforced.
+
+    Deliberately uncached: both inputs can change during the process.
+    """
+    return (
+        not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")))
+        or torch.are_deterministic_algorithms_enabled()
+    )
+
+
 def _wrap_single_quantized_as_grouped(
     tensor: torch.Tensor,
     quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage,
@@ -917,6 +935,16 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable:
         """Fused kernel for grouped GEMM, activation backward, and scale grad."""
         raise NotImplementedError
 
+    @classmethod
+    def grouped_gemm_dactivation_is_deterministic(cls) -> bool:
+        """Whether this op's dactivation kernel can be asked for a bit-exact ``dprob``.
+
+        Reported per subclass rather than per environment variable: the argument exists on
+        the dSReLU wrapper only, so a GLU activation stays non-deterministic no matter how
+        new the installed cuDNN front-end is.
+        """
+        return False
+
     @classmethod
     @functools.lru_cache(maxsize=None)
     def grouped_gemm_quant_kernel(cls) -> Callable:
@@ -2028,10 +2056,26 @@ def fuser_backward(
         scales_f32 = None
         scales_tensor = None
         dscales_tensor = None
+        deterministic_dactivation = False
         if not unit_activation_scale:
             scales_f32 = scales.detach().to(dtype=torch.float32)
             scales_tensor = scales_f32.reshape(-1, 1, 1)
             dscales_tensor = torch.zeros_like(scales_tensor)
+            # Checked here rather than up front because this is where dprob is actually
+            # produced: with a unit activation scale the cuDNN epilogue never runs its
+            # atomic dprob accumulation, so there is nothing to make deterministic.
+            if _deterministic_algorithms_required():
+                deterministic_dactivation = self.grouped_gemm_dactivation_is_deterministic()
+                if not deterministic_dactivation:
+                    raise RuntimeError(
+                        "Deterministic execution was requested"
+                        " (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or"
+                        " torch.use_deterministic_algorithms), but the cuDNN grouped-GEMM"
+                        " dactivation backward accumulates the scale gradient (dprob) with"
+                        " cross-CTA atomic adds whose summation order is set by the tile"
+                        " scheduler. A bit-exact dprob needs the scaled-SReLU activation and"
+                        " nvidia-cudnn-frontend 1.28.0 or later."
+                    )
 
         fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn
         if use_nvfp4:
@@ -2079,6 +2123,11 @@ def fuser_backward(
             "use_dynamic_sched": True,
         }
         dactivation_kernel = self.grouped_gemm_dactivation_kernel()
+        if deterministic_dactivation:
+            # Only reachable when the installed wrapper takes the argument: the check above
+            # raises otherwise, and passing it to a wrapper that does not accept it -- the
+            # dGLU one, or a front-end older than 1.28.0 -- is a TypeError.
+            fc2_dactivation_kwargs["deterministic"] = True
         if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)):
             fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
         if self._cudnn_dact_func is not None:
@@ -2607,6 +2656,24 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable:
 
         return grouped_gemm_dsrelu_wrapper_sm100
 
+    @classmethod
+    @functools.lru_cache(maxsize=None)
+    def grouped_gemm_dactivation_is_deterministic(cls) -> bool:
+        """Whether the installed dSReLU wrapper takes ``deterministic`` (cuDNN FE 1.28.0+).
+
+        Asked of the signature rather than of the package version, because the version
+        cannot answer it: NVIDIA/cudnn-frontend#521 merged after ``v1.27.0`` was tagged, but
+        ``develop`` had already called itself ``1.28.0`` for the eleven days before that, so
+        a build from that window passes a version check and then raises ``TypeError: ...
+        unexpected keyword argument 'deterministic'``. The same coarseness bit
+        ``use_single_group_runtime_offsets`` on a cuDNN reporting 1.27.0 without it.
+        """
+        try:
+            kernel = cls.grouped_gemm_dactivation_kernel()
+        except ImportError:
+            return False
+        return "deterministic" in inspect.signature(kernel).parameters
+
 
 def fuse_ops(
     ops: list[FusibleOperation],