[PyTorch] DeepSeekV3Layer: full MoE transformer layer (MLA + DeepSeek MoE) - #36
Draft
pggPL wants to merge 44 commits into
Draft
[PyTorch] DeepSeekV3Layer: full MoE transformer layer (MLA + DeepSeek MoE)#36pggPL wants to merge 44 commits into
pggPL wants to merge 44 commits into
Conversation
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
MultiLatentAttention: low-rank q/kv latents (RMSNorm fused into LayerNormLinear up-projections), decoupled RoPE/NoPE head split with a shared key rope head, DotProductAttention with kv_channels=(qk, v) for the cuDNN fused backend. DeepSeekV3MoE: fused sigmoid router with aux-loss-free expert bias and grouped top-k, routed experts as te.ops GroupedLinear+ScaledSwiGLU+ GroupedLinear (CuTe fused grouped MLP on supported HW), probs applied per-token in the activation, local permute/unpermute or NCCL expert parallelism via ep_dispatch/ep_combine, optional shared expert. DeepSeekV3Layer: pre-RMSNorm + MLA and dense LayerNormMLP (RMSNorm, swiglu) or MoE with residual connections. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
run_deepseek_ep.py checks the EP path against the all-experts-local path numerically (forward, input/gate grads, all-reduced expert wgrads) and smoke-tests the full layer with EP. Also size the default EP recv capacity for per-expert alignment padding and the fused grouped MLP's row-count requirement. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The per-expert wgrad check called all_reduce on different tensors per rank (rank-local experts), corrupting the reference grads; reduce every expert's grad on every rank instead. Also pass zero-filled recv/grad buffers to ep_dispatch/ep_combine so alignment-padding rows inside the grouped-GEMM m_splits can never poison expert wgrads. Verified on lyris (4x GB300, arm64): run_test_deepseek_ep.sh passes on all ranks (EP forward/dgrad/gate-grad/expert-wgrad match the all-local reference; full-layer EP smoke passes). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Move the Triton MLA RoPE kernels (Megatron-LM fused_mla_yarn_rope_apply port) from tests/pytorch/attention/ mla_rope_utils.py into models/deepseek_v3/mla_rope.py and use them in MultiLatentAttention: the q kernel rotates the rope slice in place and the kv kernel assembles key/value in a single pass, removing the torch.cat/expand/contiguous copies (~10% of layer GPU time). PyTorch fallback (same convention) covers missing Triton and bshd. Fix a latent bug from the test util: the q backward kernel assumed a contiguous incoming gradient, but cuDNN attention backward can hand over a strided one (allocator-state dependent IMA). The old test file stays as a compat shim. Add a Triton-vs-PyTorch parity test. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Maps HF DeepseekV3DecoderLayer weights into DeepSeekV3Layer (GLU interleave for routed experts, fused latent norms) and checks forward and input grads match within bf16 tolerance. Expose layernorm_epsilon on MultiLatentAttention (HF latent RMSNorms use 1e-6). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
docs/api/pytorch_models.rst: usage (local and EP), fused-path notes, HF checkpoint weight mapping, and the class API; linked from the PyTorch API page via a toctree entry. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Keep the verified weight-mapping table in the docs; the comparison itself stays as an out-of-tree script. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…scripts Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…eek_v3.mla_rope directly Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…rical comparison Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…hell launcher Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
… grouped GEMM in local MoE path Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…; expose MLA softmax_scale; clean docstrings Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…ters Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
… routed experts Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…ad of syncing bincount Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…p standard layers, autocast and other utilities Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…4M3 (NVIDIA#3352) * [PyTorch] Decline fused grouped MLP when the backward format is not E4M3 The fused grouped MLP packs the incoming activation gradient by reinterpreting its storage as E4M3, conditioned only on NVFP4 and never on the FP8 format. Under MXFP8BlockScaling(fp8_format=Format.HYBRID) the backward quantizers emit E5M2, so those bytes are read as the wrong format rather than converted, and every gradient out of the fusion is wrong. The forward pass is unaffected, so this shows up as a model that trains too slowly instead of one that fails. Fall back to the unfused ops when the recipe's backward format is not E4M3, and raise instead of reinterpreting if such a gradient reaches the kernel path. Signed-off-by: William <wilyan090@gmail.com> * [PyTorch] Gate the grouped MLP backward format check on MXFP8 fp8_format describes the FP8 formats of an MXFP8 recipe. NVFP4BlockScaling carries one too, pinned to E4M3, but its gradients are quantized to FP4 and the value says nothing about them, so testing it for an NVFP4 recipe reached the right answer for the wrong reason. Restrict the check to recipes where it means something. The runtime check at the pack site is unchanged: it reads the grad output quantizer's own dtype on the non-NVFP4 branch, not the recipe. Signed-off-by: William <wilyan090@gmail.com> * Test grouped MLP format fusion with real ops Signed-off-by: Przemek Tredak <ptredak@nvidia.com> --------- Signed-off-by: William <wilyan090@gmail.com> Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
… torch version (NVIDIA#3466) [PyTorch] Resolve EP symm-mem window offset against both torch symm-mem layouts Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* world_group() used a single hardcoded init_method="file:///tmp/rdzv",
shared across every world_size this test parametrizes over
([device_count(), 1, 2]), and neither world_group() nor
test_distributed_fuser_ops ever calls destroy_process_group() or
removes the file afterwards. A world_size=N run's leftover FileStore
content can then corrupt a differently-shaped world_size=M run that
reuses the same path in the same pytest session, surfacing as a
confusing NCCL bootstrap failure:
torch.distributed.DistBackendError: NCCL error ...
ncclOsSocketPollConnect: connect to <self> ... Connection refused,
exceeded error retry count after 35 attempts
instead of a clear rendezvous error. Confirmed 100% deterministic:
running the full file fresh (no pre-existing /tmp/rdzv) still fails
test_distributed_fuser_ops[2] every time, because the [4] parametrization
(which runs first) leaves /tmp/rdzv behind for [2] to trip over.
This was previously masked by older bundled NCCL (2.30.7), which
apparently tolerated the stale/mismatched FileStore well enough to
still succeed; a newer NCCL (2.31.2) surfaces it as a hard failure.
See the investigation writeup for the full comparison:
Fix: key the rendezvous path by world_size
(file:///tmp/rdzv_test_fusible_ops_{world_size}), and defensively
remove any pre-existing file at that path in test_distributed_fuser_ops
before launching each subprocess job, to also cover a leftover file
from an earlier crashed run of the same world_size.
Verified on a GB200 node
NCCL 2.31.2, the container that previously failed 2 of 3 parametrizations):
all 3 world_size parametrizations now pass, including two runs of the
full file back to back with no manual cleanup in between.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Signed-off-by: Wei Wang <weiwan@nvidia.com>
* Use unique rendezvous files in fusible ops tests
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
---------
Signed-off-by: Wei Wang <weiwan@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
* Reduce grouped MLP fuser CPU overhead Reuse fused operation plans when full activation recompute changes grad mode, and avoid redundant CUDA current-device discovery for grouped MLP stream lookups. Co-authored-by: Ting-Yang Kao <tingyangk@nvidia.com> Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * resolve comments Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * fix cutedsl wgrad crash Signed-off-by: tingyangk <tingyangk@nvidia.com> * resolve comments Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> --------- Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> Signed-off-by: tingyangk <tingyangk@nvidia.com> Co-authored-by: Ting-Yang Kao <tingyangk@nvidia.com>
* [PyTorch] Ask cuDNN for a deterministic dprob under NVTE_ALLOW_NONDETERMINISTIC_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> * Honor torch.use_deterministic_algorithms too, not just the env variable _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> * Test that dprob is actually bit-exact, not just within tolerance 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> * Probe the dsrelu wrapper's signature instead of the frontend version _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> * Revert the unrelated nccl-extensions submodule bump `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> * Raise instead of warning, and cut the change down to what it needs 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> * Apply suggestion from @vthumbe1503 Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * Match the feature-detection idiom main just landed 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> * Cut the comments down to the file's own register 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> * Flatten the determinism check and the runs-or-refuses test 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> * Close the second dprob producer, and stop discarding fp64 test tensors 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> * Add the regression test for the scale_bias hole 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> * Make the bit-exactness test capable of failing 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> * Pick the bit-exactness config by measuring it, not by borrowing one 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> * Make the scale_bias half of the dprob check readable The condition was written as `not (A and not B)` with `scale_bias` as B, which gives a reader no way to tell why an FC2 bias flag decides whether a scale gradient is reproducible -- and the call that makes it relevant is ~250 lines further down. Same logic, named and nested: `dprob_is_deterministic` says what the conjunction means, and the comment names the mechanism instead of gesturing at it. dprob is finished by two kernels, not one -- the cuDNN dactivation epilogue writes it, and then, when scale_bias is set, fuser_backward hands it to compute_grouped_dbias_dscales as the `dscales` accumulator, which atomically adds into it. Its docstring is explicit: "Both outputs use fp32 atomic adds, so pre-populated tensors are accumulated into." No logic change; the truth table is identical. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> --------- Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
* [PyTorch] Release warmup outputs after their last use Signed-off-by: Robin Zhang <robinz@nvidia.com> * [PyTorch] Release buffer-reuse capture temporaries Signed-off-by: Robin Zhang <robinz@nvidia.com> * [PyTorch] Release per-callable state on reset Signed-off-by: Robin Zhang <robinz@nvidia.com> * [PyTorch] Bundle per-callable lifecycle helpers Signed-off-by: Robin Zhang <robinz@nvidia.com> --------- Signed-off-by: Robin Zhang <robinz@nvidia.com>
) * Remove redundant runtime checks for activation recompute in MLP from _ScaledUnary class in activation.py Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> * Add warning for activation recompute in MLP outside fused path in _ScaledUnary class Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> * Add test for Scaled SReLU activation recompute warning outside fused MLP path Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> * Enhance activation recompute warning in ScaledSReLU: Update test to verify multiple warnings during backward passes and refactor warning mechanism to ensure it survives Dynamo tracing. Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> * Remove redundant logic to only warn once Signed-off-by: Tim Moon <tmoon@nvidia.com> * Remove unhelpful test Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Tim Moon <tmoon@nvidia.com>
* fix: remove unreachable fused-attn backend skip and add regression test
Root cause: FusedAttnRunner._check_configs skipped with 'Unsupported
inputs combination or device compute capability.' unless the backend was
NVTE_F16_arbitrary_seqlen. The later elif re-testing
self.backend != NVTE_F16_arbitrary_seqlen could therefore never be
reached, so its skip message ('B1SS, BHSS and 11SS bias shapes are only
supported for the F16_arbitrary_seqlen backend') was dead code.
Fix: drop the dead elif arm. The padding-mask skip in the sibling if arm
is retained. A regression test locks the remaining behavior: a
non-1HSS post-scale-bias config (BiasShape._B1SS) that passes
_check_configs selects NVTE_F16_arbitrary_seqlen, proving the removal is
behaviorally invisible and the earlier guard is the sole gate.
Testing: not run locally - the JAX test stack (jax,
transformer_engine_jax) is not installed on this machine. The suite runs
in NVIDIA TransformerEngine CI. Contribution: tests/jax tests target
real GPU/cuDNN fused-attention kernels and skip otherwise; the new test
will follow that path via CI.
Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
* Remove unnecessary test
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
---------
Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
…IDIA#3482) * Use the shmem alignment operator consistently Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Warn Linear argument documentation Signed-off-by: Evgeny <etsykunov@nvidia.com> * Update docstring, remove warning Signed-off-by: Evgeny <etsykunov@nvidia.com> --------- Signed-off-by: Evgeny <etsykunov@nvidia.com>
* Fine-grained recipe docs Signed-off-by: Evgeny <etsykunov@nvidia.com> * Update docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.com> * resolve comments Signed-off-by: Evgeny <etsykunov@nvidia.com> * Rework heterogeneous quantization docs into mixed-format quantization; add per-recipe Quantizer sections Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Align mixed-format quantization diagrams with shared diagram-colors.css and dark mode Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Rename mixed-format quantization docs to fine-grained quantization recipes Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Simplify quantizer factory paragraph Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Generalize the fallback-path description Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Evgeny <etsykunov@nvidia.com> Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
# Conflicts: # docs/api/pytorch.rst
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
11 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
Adds
transformer_engine.pytorch.models— a namespace for model-specific layers composed from TE modules — with a DeepSeek-V3 transformer layer analogous toTransformerLayer, exercising TE's MoE and MLA features end to end.MultiLatentAttention(models/deepseek_v3/multi_latent_attention.py): low-rank q/kv latents with RMSNorm fused into the up-projections (LayerNormLinear(normalization="RMSNorm")), decoupled RoPE/NoPE head split with a single shared key rope head, core attention throughDotProductAttentionwith asymmetric head dimskv_channels=(qk_nope+qk_rope, v)so the cuDNN fused attention backend is used where supported. RoPE runs through fused Triton kernels (models/deepseek_v3/mla_rope.py): in-place rotation of the query rope slice and single-pass key/value assembly with the broadcast shared rope head, with a pure-PyTorch fallback. Optional YaRN context extension (rope_scaling_factor,original_max_position_embeddings,beta_fast/beta_slow,mscale/mscale_all_dim) with the matchingsoftmax_scale; verified against the Megatron-Core YaRN implementation. TP via column/row parallel projections.DeepSeekV3MoE(models/deepseek_v3/moe.py): fused sigmoid router with aux-loss-free expert bias and node-limited grouped top-k (fused_topk_with_score_function) plus anupdate_expert_bias()helper; routed experts aste.ops.Sequential(GroupedLinear, ScaledSwiGLU, GroupedLinear), which auto-fuses into the CuTe grouped-GEMM MLP on supported hardware and runs the identical unfused path elsewhere; routing probs applied per-token inside the activation, so merge is plain accumulation. Token routing either local (moe_permute_with_probs/moe_unpermute, with per-expert row alignment under quantization) or expert-parallel over NCCL (ep_dispatch/ep_combine). Optional shared expert built with the same SwiGLU MLP helper.DeepSeekV3Layer(models/deepseek_v3/transformer_layer.py): pre-RMSNorm + MLA, then denseLayerNormMLP(first dense layers) orDeepSeekV3MoE, with residual connections.Type of change
Changes
transformer_engine/pytorch/models/namespace withmodels/deepseek_v3/subpackage (MultiLatentAttention,DeepSeekV3MoE,DeepSeekV3Layer, fused MLA RoPE kernels), exported viatransformer_engine.pytorch.modelsModel-specific layerssection on the PyTorch API page (docs/api/pytorch.rst), with the page regrouped into standard layers / model-specific layers / othertests/pytorch/test_models.py: Triton MLA RoPE vs PyTorch reference (fwd/bwd), YaRN tables andsoftmax_scale, MoE vs dense reference overtopk× grouped routing × shared experttests/pytorch/test_sanity.py::test_sanity_deepseek_v3_layer: dense/MoE layer fwd+bwd across dtypes and all quantization recipestests/pytorch/distributed/test_models.py+run_models.py(added to L1): fullDeepSeekV3Layerwith EP vs all-experts-local, comparing output, input grad, non-expert param grads and all-reduced expert wgrads; verified on 4x GB300tests/pytorch/attention/mla_rope_utils.pyremoved;test_linear_mxfp8_attention.pyusesmodels.deepseek_v3.mla_ropedirectlyChecklist:
🤖 Generated with Claude Code