Use cuDNN's deterministic dprob in the fused grouped MLP - #3407
Open
ZhiyuLi-Nvidia wants to merge 15 commits into
Open
Use cuDNN's deterministic dprob in the fused grouped MLP#3407ZhiyuLi-Nvidia wants to merge 15 commits into
ZhiyuLi-Nvidia wants to merge 15 commits into
Conversation
…ERMINISTIC_ALGO=0 The cuDNN grouped-GEMM dactivation backward that the CuTe DSL fused grouped MLP calls accumulates the scale gradient (dprob) with cross-CTA atomic adds, so its floating-point summation order follows the tile scheduler and varies run to run. Until now there was no way to switch that off, and NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 did not reach it: the run trained fine and was silently not reproducible. cuDNN frontend 1.28.0 (NVIDIA/cudnn-frontend#521) added a `deterministic` argument to grouped_gemm_dsrelu_wrapper_sm100 that parks each N-subtile's partial result in its own slot and sums the slots in a canonical order, for dprob and for dbias. Pass it from the TE flag. Passed as True or not at all, never as False. The wrapper's own default is None, which follows torch.use_deterministic_algorithms; sending an explicit False would override that and take determinism away from a caller who asked torch for it without setting the TE variable. The capability is reported per subclass rather than per environment variable, because grouped_gemm_dglu_wrapper_sm100 has no equivalent argument -- a GLU activation stays non-deterministic however new the installed front-end is. That case, and an SReLU op on a front-end older than 1.28.0, warn instead, once per distinct reason since the remedies differ. The warning is raised from where dprob is actually produced: with a unit activation scale the epilogue never runs its atomic accumulation, so there is nothing to make deterministic and nothing to warn about. Tests: TestGroupedMLPDeterminism covers the env-var parse, that only the SReLU op reports the capability and that it tracks the front-end version (no GPU or cuDNN needed for either), that the warning fires once per reason, and an MXFP8 end-to-end run under determinism for both SwiGLU and SReLU that checks numerics and pins which of the two arms warns. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Contributor
Greptile SummaryThe PR makes fused grouped-MLP scale-gradient computation honor both TransformerEngine and PyTorch determinism requests.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains. Important Files Changed
Reviews (13): Last reviewed commit: "Pick the bit-exactness config by measuri..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
_deterministic_algorithms_required() copied the narrow check from transformer_engine.pytorch.triton.grouped_dbias_dscales, which reads NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else. DotProductAttention takes the union instead -- the variable OR torch.use_deterministic_algorithms -- and that is the right precedent here. The two knobs answer different questions. The variable is set once in a job launcher, applies uniformly across ranks, and is the only one TE's C++ layer can read. The torch flag is the framework standard, is togglable at runtime, and is what a user who wants reproducibility usually reaches for; most have never heard of the variable. Keying on the variable alone left the torch flag half-honored. The SReLU path happened to come out right, but by delegation rather than by decision: TE passed nothing and the wrapper's own default read torch.are_deterministic_algorithms_enabled(). The GLU path did not -- TE stayed silent about an atomic dprob it cannot fix, for a user who had asked torch for reproducibility. That silence is the exact failure mode the warning exists to prevent, so it was the one case that most needed to warn. Passing the argument only as True, never as False, now needs a different justification than the one the first commit gave: with the union in place the two are equivalent, since the wrapper's default reads the same torch flag TE just read. The reason that survives is narrower and firmer -- the argument does not exist on the dGLU wrapper or on a front-end older than 1.28.0, where passing it at all, even as False, is a TypeError. Tests: the env-var parametrization becomes the two-knob truth table, including the row that motivates the change (torch flag set, NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -- the variable's default is the absence of a request, not a request for non-determinism, so the torch flag still wins). A fixture restores the process-global torch flag. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
ZhiyuLi-Nvidia
changed the title
Review caught that nothing in the suite tested the property this change exists for. The end-to-end test runs the op once and checks numerics against a reference with rtol=0.125 / atol=0.25; reordering the same atomic adds moves dprob by about an ulp, so a run that is silently not reproducible passes it comfortably. The tolerance check proves the deterministic path is correct, which is worth keeping, but it cannot prove the path is deterministic. Add a second run. Same module, same inputs, grads cleared between passes, probs.grad compared with torch.equal. Three things the test has to get right to be worth having: * hidden_size 1024, not the 128 used elsewhere. dprob's reduction is over that extent and the tile is 256 wide, so 128 gives a single N-tile, one writer per token, and nothing to reorder -- the assertion would hold by construction and test nothing. * No bias. With an FC2 scale_bias the scale gradient is finished by the Triton grouped dbias/dscales kernel, which refuses to run under determinism, and probs.grad would stop being the dprob under test. * An assertion that the fusion happened, since dprob only comes from the cuDNN epilogue on the fused path. Skipped rather than xfailed on a front-end older than 1.28.0: there the kernel has no deterministic mode and is expected to vary, which is not a failure of this change. Weight gradients are deliberately left out of the comparison -- the CuTe DSL wgrad kernel has its own K-split atomics that this PR does not address. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
github-actions
Bot
added
the
community-contribution
_cudnn_frontend_supports_deterministic_dprob() gated on
_cudnn_frontend_version_at_least("1.28.0"). That check is too coarse to answer the
question it is asked, and would have raised at runtime on a build TE is actually run
against.
NVIDIA#521 merged after v1.27.0 was tagged, so `deterministic` ships in 1.28.0. But
cudnn-frontend's develop branch has called itself 1.28.0 since shortly after that tag --
eleven days before the merge. Any front-end built from develop in that window reports
1.28.0 and does not accept the argument, so the version check passes, TE adds
`deterministic=True` to the call, and the backward dies with
TypeError: grouped_gemm_dsrelu_wrapper_sm100() got an unexpected keyword argument
'deterministic'
This is not hypothetical, and not new. The same coarseness already bit
use_single_group_runtime_offsets: a cuDNN reporting 1.27.0 that did not implement 1.27.0's
arguments failed the identical way, in fuser_forward, before any backward code ran.
Version numbers describe a release; they do not describe whatever happens to be installed.
Ask the function instead. `"deterministic" in inspect.signature(...).parameters` is exact,
cannot drift, and needs no maintenance when the release lands. The import is wrapped the
way _grouped_gemm_dsrelu_backward_supported() already wraps it, so a missing cuDNN answers
False rather than raising. Cached, since the call site runs every backward.
This also removes the version constant from the code path entirely -- 1.28.0 now appears
only in user-facing text, where a release number is the useful thing to say.
Tests: a smoke test that the probe returns a bool without raising, with or without cuDNN
installed, since reading a signature has more ways to fail than comparing two version
strings. It deliberately does not assert which answer -- that depends on the installed
front-end, and pinning it would only restate the implementation.
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
`git add -u` in the previous commit swept in a local 3rdparty/nccl-extensions pointer change that has nothing to do with this PR. Restore it to main's commit so the branch touches only the three files it means to. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
vthumbe1503
reviewed
Aug 21, 2026
| " only from cuDNN frontend 1.28.0 on; upgrade" | ||
| " nvidia-cudnn-frontend to get a bit-exact dprob" | ||
| ) | ||
| _warn_nondeterministic_cudnn_dprob(reason) |
Collaborator
There was a problem hiding this comment.
We should rather throw an error here
Author
There was a problem hiding this comment.
raises RuntimeError now.
Comment on lines
+2035
to
+2046
| if self._cudnn_dact_func is not None: | ||
| reason = ( | ||
| "grouped_gemm_dglu_wrapper_sm100 has no deterministic mode, so" | ||
| " only the scaled-SReLU activation can be made bit-exact" | ||
| ) | ||
| else: | ||
| reason = ( | ||
| "grouped_gemm_dsrelu_wrapper_sm100 takes a deterministic argument" | ||
| " only from cuDNN frontend 1.28.0 on; upgrade" | ||
| " nvidia-cudnn-frontend to get a bit-exact dprob" | ||
| ) | ||
| _warn_nondeterministic_cudnn_dprob(reason) |
Collaborator
There was a problem hiding this comment.
The solution for both cases, from the user's point of view is to upgraded cudnn-frontend to 1.28.0 or later. Can we just have that as the reason shown in the error? I think checking for self._cudnn_dact_func is an overkill
Author
There was a problem hiding this comment.
single message, branch deleted
Review asked for two things on the unsupported path: make it an error rather than a warning, and stop branching on self._cudnn_dact_func to pick a message. Both are right, and taking them removes most of the machinery this PR had accumulated. Raising matches what TE already does elsewhere: the Triton grouped dbias/dscales kernel refuses to run under determinism rather than running non-deterministically. It also matches what the variable documents -- "only deterministic algorithms are allowed" is not "prefer deterministic algorithms". A silently non-reproducible run is the failure this PR exists to prevent, so continuing past a request TE cannot honor was the wrong default. Checked that no existing determinism test hits this path: test_hybrid_quantization sets the variable for an attention recipe, and test_fusible_ops_with_userbuffers for linear ops. One message, no branch. The two cases did have different remedies, which is why the branch was there, but a single sentence states both facts -- "needs the scaled-SReLU activation and nvidia-cudnn-frontend 1.28.0 or later" -- without telling a SwiGLU user to go upgrade. What that let me delete: * _warn_nondeterministic_cudnn_dprob and its per-reason lru_cache, the two reason strings and the branch selecting them: 30 lines at the call site and above it, down to a single raise. * _cudnn_frontend_supports_deterministic_dprob as a standalone function. The probe now lives in GroupedMLP_CuTeGEMMUnary.grouped_gemm_dactivation_is_deterministic(), which reaches the wrapper through grouped_gemm_dactivation_kernel() -- the import and its ImportError handling already existed there, so folding it in dropped a duplicate import and an indirection. * The warn-once cache-clearing fixture in the tests, and the two tests that existed only to cover the warning. Tests: test_deterministic_dactivation_is_numerically_correct becomes test_determinism_either_runs_or_refuses -- it expects RuntimeError where the request cannot be honored and runs the full numerical check where it can, so both arms assert something either way. The bit-exactness and two-knob tests are unchanged in substance. Net: transformer_engine/pytorch/ops/fused/grouped_mlp.py goes from +106 to +68, all of it addition, no line of pre-existing code touched. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
vthumbe1503
reviewed
Aug 21, 2026
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Collaborator
|
/te-ci pytorch |
vthumbe1503
reviewed
Aug 21, 2026
vthumbe1503
left a comment
vthumbe1503
left a comment
Collaborator
There was a problem hiding this comment.
LGTM @ZhiyuLi-Nvidia. I would want to retest this PR post cudnn-1.28 release to make sure functionality works correctly with TE after the cudnn upgrade
The SiTU-GLU merge (NVIDIA#3402) brought _cudnn_frontend_supports_grouped_gemm_situglu() into this file, which asks inspect.signature(wrapper).parameters for the arguments it needs rather than comparing frontend versions -- the same conclusion this branch reached independently, now the house style. Two things to match. Guard the signature call with `except (TypeError, ValueError)`: a callable that is not introspectable answers "no" instead of raising out of a backward pass. I had left this out on the grounds that the wrapper is a plain undecorated function, which is true today but is not a property this code controls. And say "feature-detect" in the docstring summary, as the neighbor does. Also dropped the sentence about use_single_group_runtime_offsets from the docstring. The neighbor now demonstrates the pattern in the same file, so the cautionary tale is no longer what makes the choice legible. `import inspect` came in with the merge, so this branch no longer adds it. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
The new code carried multi-paragraph docstrings into a file whose 44 functions have a
median docstring of one line. Measured before and after:
grouped_mlp.py _deterministic_algorithms_required 10 -> 3 lines
grouped_gemm_dactivation_is_deterministic (base) 5 -> 1
grouped_gemm_dactivation_is_deterministic (unary) 7 -> 1
test_grouped_mlp.py four new tests 4-6 -> 1-4
four inline comment blocks 2-3 -> 1 each
Before this, the three new functions were the 2nd, 3rd and 5th longest docstrings in
grouped_mlp.py; only fuse_grouped_mlp_ops, which has a full Parameters block, was longer.
In the test file, 63 pre-existing tests have a median docstring of zero lines.
Most of what came out was rationale, not explanation: why the union matches
DotProductAttention, why feature detection beats a version compare, which cuDNN release
window motivated it. That belongs in the commits that made those choices, where it already
is, and it reads as noise next to _cudnn_frontend_supports_grouped_gemm_situglu -- the
neighbor doing the very same feature detection in a one-line docstring with no rationale
at all.
What stayed is what the code cannot say itself: that the check sits inside the
non-unit-scale branch because a unit scale produces no dprob; that hidden_size must exceed
one N-tile or the bit-exactness test is vacuous; that bias would reroute probs.grad through
Triton; that weight grads are excluded because wgrad has its own atomics. Each is now one
line.
No behavior change -- comments, docstrings and one local variable's reading order only.
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Structural cleanups from the review pass. grouped_mlp.py: the check was nested two deep inside `if not unit_activation_scale`, and assigned deterministic_dactivation only to immediately test its own assignment. Hoisted to two flat statements right after unit_activation_scale is computed. `not unit_activation_scale and _deterministic_algorithms_required()` now says in the expression what the comment had to say in prose, and the separate `= False` initializer is gone. The local itself stays -- the kwargs dict is built about sixty lines further down. Also shortened the error: the tile-scheduler detail was not actionable, and "this activation's cuDNN dactivation kernel" is more accurate than naming the grouped-GEMM backward, since which kernel it is depends on the activation. test_grouped_mlp.py: fused_cls was derived from `activation` by a five-line conditional inside the test; it is now the second half of the parametrize pair. That also fixes the skip guard, which asked GroupedMLP_CuTeGEMMGLU.is_supported() on both parametrizations including the SReLU one -- the sibling test three functions down already gets this right. The _run closure existed only so an if/else could call it twice; a contextlib.nullcontext / pytest.raises choice removes the closure and the branch. nullcontext is used in ten test files here, so it is the local idiom rather than a new one. Not taken: dropping the `isinstance(..., bool)` assertion. It looks vacuous but it is the only coverage of the ImportError branch in the capability probe, which is the branch that runs on every machine without cuDNN -- including CI. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Contributor
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
Two findings from the review pass. dprob has two producers in this backward, and the check only covered one. The cuDNN epilogue produces grad_scales at fuser_backward, and when scale_bias is set compute_grouped_dbias_dscales accumulates into it further down -- the Triton kernel that grouped_dbias_dscales.py documents as nondeterministic atomic adds. That kernel's own guard reads NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else. So the hole opened exactly where this branch widened the trigger. With torch.use_deterministic_algorithms(True) and the variable unset -- the case the union exists to start honoring -- SReLU on a 1.28.0 front-end with scale_bias passed the new check, set deterministic=True, raised nothing, and then routed dprob through the nondeterministic path anyway. Env-var users were never exposed: the Triton guard fires for them. It was reachable only via the torch flag, which is to say only through what this branch added. The test picked bias=False and so never crossed it. scale_bias is computed ~130 lines earlier in the same scope, so the fix is to require both producers rather than one. Still one condition and one message, per review -- the message now lists all three requirements instead of two. Separately, the bit-exactness test built its tensors with make_reference_and_test_tensors and discarded the reference every time. That helper allocates an fp64 CPU companion, quantizes and dequantizes for MXFP8 representability, then copies back D2H with an implicit sync -- about 16 MB of host allocation across the two (1024, 1024) calls, for a test that compares run 1 against run 2 and never against a reference. Twelve of the file's other fifteen uses keep the reference; this one had no use for it. Plain uniform_ tensors instead. Also dropped a .item() sync for a token count already known in Python. Not taken, with reasons: * Hoisting _deterministic_algorithms_required into pytorch/utils.py so the Triton guard reads the same union. That is the deeper fix and it is correct, but broadening that guard changes behavior for callers this PR does not touch (ops/basic/grouped_linear.py, module/grouped_linear.py) -- users who set only the torch flag would start seeing RuntimeError where they now get silent nondeterminism. Worth doing deliberately, not as a side effect of this branch. * Extracting the signature-probe shared with _cudnn_frontend_supports_grouped_gemm_situglu. The overlap is about four lines and the two are not interchangeable; refactoring working code outside the diff to save them is not this PR's job. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
The previous commit fixed a real bug and shipped it with no test. Every test in the class used bias=False and every end-to-end one set the env var, so neither half of the bug was reachable: not scale_bias, and not the torch-flag-only trigger. Both halves are load-bearing. With the env var the Triton kernel raises on its own, so an env-var test would have passed before the fix as well as after and pinned nothing. Only torch.use_deterministic_algorithms with the variable unset reaches the state where this op's check said yes and the Triton reduction then ran nondeterministically. warn_only=True so torch's own enforcement cannot raise first and be mistaken for TE's refusal. are_deterministic_algorithms_enabled() still reports True in that mode -- the separate is_deterministic_algorithms_warn_only_enabled() getter exists precisely because the two are independent -- so the predicate under test sees what it should. Not executed: no GPU or torch on the machine this was written on. Formatting and syntax only, like the rest of the branch. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
ZhiyuLi-Nvidia
changed the title
Followed cudnn-frontend#521's own test work and found this test had the flaw its commit 88c7fab was written to fix, at the same config. That commit measured 16 launches per shape and found that at l=4 / [256]*4 / n=512 the NONDETERMINISTIC dprob is already bit-stable: the assertion cannot fail there, so a pass certifies nothing. It varies 15/15 at l=8 / [1024]*8 / n=2048. This test used l=4 / [256]*4 / n=1024 -- the vacuous shape, one power of two along n. Moved to the shape that actually varies. n > 256 was necessary but not sufficient, which is what the old comment got wrong. Spanning several N-tiles exercises the within-CTA subtile ordering; making the cross-CTA reduction unstable needs the larger token count and expert count too. Also took the rest of NVIDIA#521's discipline for these comparisons: * Repeat rather than compare a pair. The order determinism removes is set by the tile scheduler, so two runs can match by luck. Four by default, NVTE_TEST_DETERMINISM_REPEATS to raise it, matching that file's DETERMINISM_REPEATS. * Compare bytes, not values. torch.equal treats +0.0 and -0.0 as equal, and a change in reduction order produces exactly that; upstream's bitwise_bits views as uint8 for the same reason. * Assert the output is finite first, so a NaN run cannot be read as a determinism result. Not copied: asserting that the nondeterministic path *does* vary. It is the thing that makes the config meaningful, but as an assertion it is timing-dependent and would flake. Upstream settled this by measuring once and pinning the config; the comment now cites that measurement so the next person does not shrink the shape back. Not run yet: job 535935 is building the previous revision of this test. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Measured on GB300 across five shapes, determinism off, 8 launches each (job 538058), counting how many runs differ from run 0: l=8 tok/grp=1024 n=2048 2/7 max|d| 5.96e-08 l=8 tok/grp=1024 n=4096 2/7 max|d| 9.54e-07 l=16 tok/grp=1024 n=2048 7/7 max|d| 1.19e-07 l=8 tok/grp=2048 n=2048 5/7 max|d| 7.63e-06 l=4 tok/grp=512 n=8192 6/7 max|d| 1.91e-06 Moved to l=16, the only shape where every run differs, so the assertion cannot pass by luck. The previous choice, l=8, varies 2/7 -- an eight-run sample calls it stable often enough to be a poor detector, and an earlier control run (537313) did exactly that and reported 0/7 at this shape. I took that single sample as proof the config was vacuous and said so; it was a sampling artifact, and the shape does vary, just weakly. That earlier shape came from cudnn-frontend#521's own measurement, which was taken on its direct wrapper test. It does not transfer to TE's path -- different scheduler settings, different quantization -- so borrowing the number was the mistake underneath both errors. This config is measured through the fused grouped MLP itself. Two things the same job settled that are worth recording: * Without NVIDIA#521 the values genuinely move: 6e-08 to 8e-06 absolute across these shapes. Small, but nonzero every time, and the reason the refusal exists rather than a warning. * The refusal cannot be exercised against the stock 1.27.0 frontend on this image at all. TE's forward passes prob_tensor=None because _cudnn_frontend_version_at_least("1.27.0") reports optional-prob support that a stock 1.27.0 does not implement, so the op dies in fuser_forward with "prob_tensor is required" before any determinism code runs. Same version-gate-too-coarse failure this PR avoids for its own argument, on a gate it does not own. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
cuDNN's grouped-GEMM dSReLU backward accumulates the scale gradient (
dprob) with cross-CTA atomic adds, so its summation order follows the tile scheduler and the result varies run to run. NVIDIA/cudnn-frontend#521 added adeterministicargument togrouped_gemm_dsrelu_wrapper_sm100that makes it bit-exact.This passes that argument when the user asks for determinism, and raises when
dprobcannot be made bit-exact instead of running anyway.Determinism is "asked for" when
NVTE_ALLOW_NONDETERMINISTIC_ALGO=0ortorch.use_deterministic_algorithmsis set — the same unionDotProductAttentionuses.The argument ships in cuDNN frontend 1.28.0, which is not released yet. The gate feature-detects it on the installed wrapper rather than comparing versions, because the version string does not track the feature: #521 merged on 2026-08-17 and
developwas only bumped to 1.28.0 on 2026-08-19 (#668). A build from that window — or #521's own branch, which reports1.27.0— acceptsdeterministicwhile a>= 1.28.0check would say no and silently drop determinism, which is the failure this change exists to prevent.Type of change
Changes
deterministic=Truetogrouped_gemm_dsrelu_wrapper_sm100when determinism is requested and the installed wrapper accepts it.RuntimeErrorwhen determinism is requested butdprobcannot be bit-exact — a GLU activation (nodeterministicargument upstream), a cuDNN frontend older than 1.28.0, or an FC2scale_bias(which finishesdprobin a nondeterministic Triton kernel).TestGroupedMLPDeterminismintests/pytorch/test_grouped_mlp.py.Out of scope: the CuTe DSL grouped-GEMM wgrad kernel has its own cross-CTA atomics and is unchanged.
Checklist:
The last box is unchecked deliberately: this was written on a machine with no GPU and no PyTorch, so the tests have not been run. The deterministic path also cannot be exercised on any released cuDNN frontend yet — it needs a
developbuild carrying #521.