[PyTorch] Ask cuDNN for a deterministic dprob when determinism is requested - #3407
Open
ZhiyuLi-Nvidia wants to merge 5 commits into
Open
[PyTorch] Ask cuDNN for a deterministic dprob when determinism is requested#3407ZhiyuLi-Nvidia wants to merge 5 commits into
ZhiyuLi-Nvidia wants to merge 5 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 propagates PyTorch or Transformer Engine determinism requests to the supported cuDNN grouped-GEMM dSReLU backward and warns when deterministic dprob accumulation is unavailable.
Confidence Score: 5/5The PR appears safe to merge. The previously reported coverage gap is resolved by executing the fused backward path twice with identical inputs and comparing the resulting dprob gradients exactly; no blocking failure remains. Important Files Changed
Reviews (4): Last reviewed commit: "Revert the unrelated nccl-extensions sub..." | 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
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
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
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 a determinism request did not reach it — the run trained fine and was silently not reproducible.NVIDIA/cudnn-frontend#521 (merged) added a
deterministicargument togrouped_gemm_dsrelu_wrapper_sm100: each N-subtile's partial result parks in its own slot and the slots are summed in a canonical order, fordproband fordbias. This PR passes it when the user has asked for reproducibility.Important
On availability. #521 landed on
developafterv1.27.0was tagged, so it ships in cuDNN frontend 1.28.0, which is not yet released — latest tag and latest PyPI are both1.27.0.TE therefore does not gate on the version.
develophas called itself1.28.0since shortly after that tag, eleven days before #521 merged, so a build from that window would pass a version check and then die withTypeError: ... unexpected keyword argument 'deterministic'. That is the same coarseness that already bituse_single_group_runtime_offsetson a cuDNN reporting 1.27.0 without implementing it. The gate asks the installed function instead:"deterministic" in inspect.signature(grouped_gemm_dsrelu_wrapper_sm100).parameters.Consequence for review: on any released front-end the new branch is inert,
deterministic=Trueis never passed, andtest_dprob_is_bit_exact_across_runsskips whiletest_deterministic_dactivation_is_numerically_correctexpects the warning on both arms. On adevelopbuild carrying #521 it activates immediately, with no code change and no release-tracking edit. Reviewers should know the new branch is currently covered by construction on public builds, not by a green CI run.Three details worth calling out:
DotProductAttentionalready uses (dot_product_attention.py:716):NVTE_ALLOW_NONDETERMINISTIC_ALGO=0ortorch.use_deterministic_algorithms. They 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. Keying on the variable alone would leave the torch flag half-honored: the SReLU path would still come out right by delegation (the wrapper's own default reads the torch flag), but TE would stay silent about the GLU path it cannot fix — the exact failure mode the warning exists to prevent.True, and only once known to be accepted. It does not exist on the dGLU wrapper or on a front-end older than 1.28.0, where passing it at all — even asFalse— is aTypeError.grouped_gemm_dglu_wrapper_sm100has no equivalent, so a GLU activation stays non-deterministic however new the installed front-end is. That case, and an SReLU op on an older front-end, warn instead — once per distinct reason, since the remedies differ. Merging them into one message would tell a SwiGLU user to upgradenvidia-cudnn-frontend, which would not help them.Out of scope: the CuTe DSL grouped-GEMM wgrad kernel has its own cross-CTA atomic accumulation over its K-split and is not addressed here.
Type of change
Changes
_deterministic_algorithms_required()—NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or torch.are_deterministic_algorithms_enabled(), matchingDotProductAttention._cudnn_frontend_supports_deterministic_dprob()probes the wrapper's signature for the argument, rather than comparing front-end versions.grouped_gemm_dactivation_is_deterministic()—Falseon the base op, the version gate onGroupedMLP_CuTeGEMMUnary.deterministic=Truewhen both hold; otherwise_warn_nondeterministic_cudnn_dprob(reason)warns once per reason.docs/envvars.rst:NVTE_ALLOW_NONDETERMINISTIC_ALGOnow documents what it does for the fused grouped MLP, that the torch flag is honored equally, and the two cases where neither can be honored.TestGroupedMLPDeterminismintests/pytorch/test_grouped_mlp.py.Checklist:
Notes on the last box, so nothing is overstated:
blackand a syntax check. The three CPU-only tests inTestGroupedMLPDeterminismneed neither a GPU nor cuDNN; the two end-to-end tests need an SM100 GPU.developbuild carrying [JAX] Fix JAX distributed unit tests #521.check_support()rather than from running them:deterministic=Truewithuse_dsrelu_reuse=True(TE passes this when recomputing the FC2 input) and with NVFP4. Neither appears in [JAX] Fix JAX distributed unit tests #521's own tests. If a reviewer knows of a constraint there, I would rather hear it than find it at runtime.