diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py index 5049f58ce4..4623bcb98a 100644 --- a/tests/pytorch/test_grouped_mlp.py +++ b/tests/pytorch/test_grouped_mlp.py @@ -5,6 +5,7 @@ from __future__ import annotations from collections.abc import Iterable +import contextlib import functools import os import math @@ -2772,6 +2773,194 @@ def train_step( assert_close(graph_grad, param.grad, **tols) +class TestGroupedMLPDeterminism: + """Determinism coverage for the CuTe DSL fused grouped MLP. + + Only the dSReLU wrapper can make ``dprob`` bit-exact, and only from cuDNN FE 1.28.0 on. + Anything else must refuse a determinism request rather than run non-deterministically. + """ + + @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: + """``=1`` is the absence of a request, not a request for non-determinism.""" + 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 belongs to the wrapper, not the environment. Needs no GPU.""" + 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,fused_cls", + ( + ("scaled_srelu", grouped_mlp_module.GroupedMLP_CuTeGEMMUnary), + ("scaled_swiglu", grouped_mlp_module.GroupedMLP_CuTeGEMMGLU), + ), + ) + def test_determinism_either_runs_or_refuses( + self, monkeypatch, *, activation, fused_cls + ) -> None: + """A request TE cannot honor must fail loudly; one it can must still be correct.""" + if not fused_cls.is_supported(): + pytest.skip("MXFP8 fused grouped MLP is not supported on this system") + + monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0") + expectation = ( + contextlib.nullcontext() + if fused_cls.grouped_gemm_dactivation_is_deterministic() + else pytest.raises(RuntimeError, match="dprob") + ) + with expectation: + TestGroupedMLPFusedOp().test_grouped_mlp( + bias=False, + hidden_size=128, + quantization="mxfp8", + single_grouped_weight=False, + activation=activation, + ) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_scale_bias_refuses_under_the_torch_flag( + self, monkeypatch, _restore_torch_determinism + ) -> None: + """``scale_bias`` finishes ``dprob`` in a Triton kernel that reads only the env var. + + So the torch flag alone is the combination that used to pass this op's own check and + then reduce nondeterministically anyway, on a front-end new enough to say yes. + """ + 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") + + monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False) + # warn_only so torch's own enforcement cannot raise first and mask what TE does. + torch.use_deterministic_algorithms(True, warn_only=True) + with pytest.raises(RuntimeError, match="dprob"): + TestGroupedMLPFusedOp().test_grouped_mlp( + bias=True, + hidden_size=128, + quantization="mxfp8", + single_grouped_weight=False, + activation="scaled_srelu", + ) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None: + """Repeated identical runs must give a bit-identical ``dprob``. + + An ulp of reordering passes every tolerance in this file, so only an exact + comparison across runs can see it. + """ + 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 + # Measured on GB300, determinism off, 8 launches per shape (job 538058): this shape + # gives 7/7 runs differing from run 0, so the assertion below can actually fail. + # Shapes matter more than they look -- l=8 with the same n and tokens/group varies + # only 2/7, which an 8-run sample reports as stable often enough to be useless, and + # cudnn-frontend#521 measured its own l=4 / [256]*4 / n=512 as never varying. + group_size = 16 + hidden_size = 2048 + tokens_per_group = 1024 + split_sizes = torch.tensor([tokens_per_group] * group_size, dtype=torch.int, device=device) + num_tokens = tokens_per_group * group_size + + recipe = make_recipe("mxfp8") + + # Plain random tensors, not make_reference_and_test_tensors: this test compares two + # runs against each other, never against a reference, so the fp64 companion and the + # MXFP8 representability round-trip would both be allocated and thrown away. + def _rand(*shape, requires_grad=True) -> torch.Tensor: + out = torch.empty(shape, dtype=dtype, device=device).uniform_(-0.25, 0.25) + return out.requires_grad_() if requires_grad else out + + x = _rand(num_tokens, hidden_size) + dy = _rand(num_tokens, hidden_size, requires_grad=False) + probs = _rand(num_tokens) + + # No bias, or probs.grad comes from the Triton dbias kernel instead of cuDNN. + 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() + + runs = [_run()] + # Without the fusion there is no cuDNN dprob and the comparison proves nothing. + forward_ops = module._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], fused_cls) + # More than two, as cudnn-frontend#521 does: the cross-CTA order that determinism + # removes is set by the scheduler, so two runs can agree by luck. + runs += [_run() for _ in range(int(os.getenv("NVTE_TEST_DETERMINISM_REPEATS", "4")) - 1)] + torch.cuda.synchronize() + + assert torch.isfinite(runs[0]).all(), "dprob is not finite; the comparison would be moot" + # Bytes, not values: torch.equal calls +0.0 and -0.0 equal, and a change in reduction + # order can produce exactly that. Weight grads are excluded from the comparison -- + # the CuTe DSL wgrad kernel has its own K-split atomics, which this change leaves. + for index, later in enumerate(runs[1:], start=1): + assert torch.equal( + runs[0].contiguous().view(torch.uint8), later.contiguous().view(torch.uint8) + ), ( + f"dprob differs between run 0 and run {index} under determinism; max |delta| =" + f" {(runs[0].float() - later.float()).abs().max().item()}" + ) + + def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None: if not mxfp8_available: pytest.skip(reason_for_no_mxfp8) diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index dca1119800..adf6bfb949 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py @@ -137,6 +137,17 @@ def _cudnn_frontend_supports_single_group_runtime_offsets( ) +def _deterministic_algorithms_required() -> bool: + """Whether bit-exact reproducibility was asked for. Same union as ``DotProductAttention``. + + Uncached: both knobs 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 +928,11 @@ 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 produce a bit-exact ``dprob``.""" + return False + @classmethod @functools.lru_cache(maxsize=None) def grouped_gemm_quant_kernel(cls) -> Callable: @@ -2025,6 +2041,28 @@ def fuser_backward( current_stream = torch.cuda.current_stream().cuda_stream unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False)) + # A unit activation scale produces no dprob, so there is nothing to make deterministic. + deterministic_dactivation = ( + not unit_activation_scale and _deterministic_algorithms_required() + ) + if deterministic_dactivation: + # Two kernels write dprob and both have to be exact. The cuDNN dactivation + # epilogue produces it below; then, when scale_bias is set, it is passed to + # compute_grouped_dbias_dscales as the ``dscales`` accumulator and atomically + # added into (see triton/grouped_dbias_dscales.py). That Triton kernel is never + # deterministic, so scale_bias rules out a bit-exact dprob on its own. + dprob_is_deterministic = ( + self.grouped_gemm_dactivation_is_deterministic() and not scale_bias + ) + if not dprob_is_deterministic: + raise RuntimeError( + "Deterministic execution was requested" + " (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or" + " torch.use_deterministic_algorithms), but the scale gradient (dprob) is" + " accumulated with nondeterministic atomics on this configuration." + " A bit-exact dprob requires the scaled-SReLU activation," + " nvidia-cudnn-frontend 1.28.0 or later, and an FC2 without scale_bias." + ) scales_f32 = None scales_tensor = None dscales_tensor = None @@ -2079,6 +2117,9 @@ def fuser_backward( "use_dynamic_sched": True, } dactivation_kernel = self.grouped_gemm_dactivation_kernel() + if deterministic_dactivation: + # Never passed to a wrapper that would reject it -- the check above raises first. + 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 +2648,19 @@ 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: + """Feature-detect the dSReLU wrapper's ``deterministic`` argument (cuDNN FE 1.28.0+).""" + try: + kernel = cls.grouped_gemm_dactivation_kernel() + except ImportError: + return False + try: + return "deterministic" in inspect.signature(kernel).parameters + except (TypeError, ValueError): + return False + def fuse_ops( ops: list[FusibleOperation],