Batch-invariant mode: native NVLS multimem option, and training-side mirror for true on-policy RL - #6634
Draft
utkarsh530 wants to merge 25 commits into
Draft
Batch-invariant mode: native NVLS multimem option, and training-side mirror for true on-policy RL#6634utkarsh530 wants to merge 25 commits into
utkarsh530 wants to merge 25 commits into
Conversation
Under batch_invariant_mode, graphed steps must execute norms/GEMMs in the
same M-alignment class as eager steps (TE rmsnorm switches reduction
codepaths at M % 32; the eager path pads token counts to TOKEN_ROUNDER=64).
The bucket auto-sizing paths (exponential ladder endpoint tp_size; linear
-1 ladder [1, 2, 4, 8, ...]) inject 1- and 2-token decode buckets whose
graphed norms execute in a different bit-class, breaking cross-batch
bit-equality.
Floor every generated bucket token count to a 64-multiple (min 64) when
batch-invariant mode is enabled: in both sizing distributions and centrally
in add_if_valid (with duplicate suppression after collisions). Request
counts are untouched, mirroring eager TOKEN_ROUNDER semantics.
Verified on the NeMo-RL true-on-policy determinism campaign: without the
floor, graphed decode diverges bitwise from eager scoring at M in {1,2}
buckets; with it, full-CUDA-graph decode is bitwise-exact vs the TE
scoring path over 20-step GRPO runs (gen_kl == 0.0).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
batch_invariant_mode hard-wired the MoE activation to squared_relu_with_probs, so gated-unit models (e.g. Qwen3 MoE, enabled for inference_optimized by NVIDIA#5700) fail with a grouped-GEMM K mismatch (FC1 output is 2*ffn wide and never halved). Add swiglu_with_probs — the gated-SiLU counterpart with the same graph-safe fixed-CTA structure — and select it from mcore_fused_moe when activation_type is SWIGLU. SiLU(gate)*up*prob is computed in FP32 with a single BF16 round, matching the training fused weighted-swiglu rounding. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
batch_invariant_mode previously forced inference_grouped_gemm_backend='torch' (DeepGEMM grouped GEMM), which costs ~6x generation throughput vs the vLLM Triton fused-MoE backend on the same engine (measured on Qwen3-30B-A3B, EP8/TP1, 8xB200: ~4.1k vs ~25.4k tok/s at BS256/OSL1024). This change makes the vLLM backend batch-invariant and allows it under the flag; the equivalent deterministic configuration reaches ~24k tok/s (0.99x of the engine's non-invariant throughput) and is certified bitwise-identical to the Megatron training forward over full-learning-rate GRPO runs (gen_kl == 0.0) in NeMo-RL true on-policy training. Three changes, active only under batch_invariant_mode: - vllm_fused_moe: pin the kernel launch config. _get_default_config selects tile shapes from the token-count hint, so different co-batch sizes change the fp32 accumulation grouping (batch-variant bits). Grid sizing may still use the hint: the kernel strides, so grid size never changes per-tile math. - SwiGLU activation: apply routing probabilities at the activation (before FC2) with the training kernel's exact rounding sequence, via a new device-bounded weighted SiLU-mul kernel (persistent grid, bounded by the valid_tokens*topk device scalar; CUDA-graph safe). Matches the convention already used by the torch-backend batch-invariant activation. - _moe_sum: unit weights when probs were applied at the activation, and fp64 accumulation so the topk reduction is order-independent by precision. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
…ion) Add a third batch-invariant GEMM backend that keeps the NATIVE cuBLASLt kernels for every dense GEMM (aten and TE) and obtains batch invariance by starving the cuBLASLt workspace (~1KB): split-K reduction variants require workspace, so starving it disqualifies them and pins every M to the same serial-K reduction recipe. M/N tile selection may still vary with M, but tiling does not affect bits — bf16 products are exact in fp32, so only the K-reduction order matters. This is the lowest-overhead invariant dense path: no kernel substitution, native speed. It is also the configuration certified bitwise-identical to the Megatron/TE training forward in NeMo-RL true on-policy GRPO (gen_kl == 0.0 over full-learning-rate runs, Qwen3-30B-A3B EP8/TP1 8xB200). Notes: - TE (<= 2.15 verified) hardcodes a 32MiB workspace in get_cublas_workspace_size_bytes() and ignores CUBLASLT_WORKSPACE_SIZE, so the env pin alone never engages for TE-launched GEMMs; the backend patches the size fn and clears its lru_cache. - Under te_native, aten::mm/addmm are left unpatched and TE general_gemm is not substituted (skip_gemm); the non-GEMM batch-invariant patches (log_softmax, mean, RMSNorm, attention gate, Mamba autotuner pins) still apply. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Add TransformerConfig.batch_invariant_backend ('deepgemm' | 'triton' |
'te_native') and plumb it through training initialization, so the backend
added in the previous commit is reachable from the CLI
(--batch-invariant-backend) instead of being hardcoded to the default.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
TE RMSNorm switches its reduction codepath (and therefore bit pattern on rare-value rows) at M % 32. The default batch-invariant backends handle this by substituting a batch-invariant RMSNorm kernel; the te_native backend can instead keep the NATIVE kernel, because the 64-multiple alignment discipline (CUDA-graph bucket floor, eager TOKEN_ROUNDER, scoring-side sequence-length rounding) holds every launch in the same M%32 bit-class. Native RMSNorm is faster and is the configuration certified bitwise-identical to the training forward. Applies to both the inference-optimized RMSNorm call site and the TE class/module-level patches (skip_rmsnorm under te_native). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
The previous change pinned the entire launch config to the large-M tuning, which penalizes decode (small-M) steps. The fused-MoE kernel accumulates in fp32 with no split-K, so its bit pattern depends only on the K-loop grouping: pin BLOCK_SIZE_K and keep the M/N tile shapes, tile grouping and pipeline depth hint-adaptive. bf16 products are exact in fp32, so tile shape changes reorder nothing in the accumulation — only the K addition order matters. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Unit tests for the batch-invariant vLLM-backend work: - CUDA-graph bucket token counts are 64-multiples under batch-invariant mode (all sizing distributions incl. num_cuda_graphs=-1 auto), with a guard test documenting that the non-BI auto ladder still injects 1/2-token buckets. - swiglu_with_probs / weighted_silu_mul_bounded: value correctness (tolerance-based vs a torch reference), bitwise repeat-determinism, row-locality across co-batch sizes, and device-bound soundness under NaN-poisoned tails. - _moe_sum apply_weights/acc_fp64 options: unit-weight fp64 reduction is bitwise-exact vs an fp64 reference; the default weighted fp32 path is unchanged. - End-to-end batch invariance: the same tokens produce bitwise-identical outputs across co-batch sizes/hint classes (exercises the pinned K-reduction recipe with adaptive M/N tiling and warp counts). - te_native backend registration and enable/disable round-trip (native aten::mm stays unpatched). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
…_floor Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
te_native keeps the native cuBLASLt kernels (invariance via workspace starvation), has the lowest overhead, needs no extra dependencies, and is the configuration verified bitwise-identical to the TE training forward. Also scope the DeepGEMM MoE requirement to the backend combinations that actually use it (deepgemm/triton backends, or the torch inference grouped-GEMM path); te_native with the vLLM inference backend — and the training path, where TE grouped GEMM stays native — no longer require DeepGEMM to be installed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
APPLY_WEIGHTS was nested inside the ACC_FP64 branch, so apply_weights=False was silently ignored on the fp32 path (weights applied anyway). Latent today — the single call site couples the flags — but the docstring advertises them as independent. Hoist APPLY_WEIGHTS out and add the missing (apply_weights=False, acc_fp64=False) test case. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Validate batch dimensions BEFORE the 64-multiple alignment: the floor is padding, mirroring the eager path's TOKEN_ROUNDER (which already yields token counts above what the requests produce), so validity is judged on the unpadded dims and request budgets are untouched. Previously a non-64-multiple max_requests (e.g. 100) lost its largest decode bucket to the is_valid token-sufficiency check. Regression test with max_requests=100 added. Also rename _batch_invariant_token_floor -> _batch_invariant_token_align (it rounds up). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Save and restore TE's get_cublas_workspace_size_bytes across enable/disable so the 1KB workspace cannot leak into non-batch-invariant work in the same process, and force-pin CUBLASLT_WORKSPACE_SIZE (with save/restore, a warning when overriding a preset value, and a warning when enabling after CUDA initialization — earlier cuBLASLt handles may retain their original workspace). Round-trip test asserts full restoration. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Squared-ReLU's weighted fp32 reduction is already order-deterministic, so fp64 accumulation is only enabled for the SwiGLU (unit-weight) path. batch_invariant_backend becomes Literal[...] for argparse choices and config validation. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
num_programs=1184 was the Inductor launch config captured on B200 (148 SMs x 8 waves) — device-specific. Derive it as SMs * 8 with the explicit argument kept as an override, and document why grid size is bit-inert (persistent elementwise loop, disjoint per-program ranges). Add a test asserting bitwise-identical output across grid sizes (1, 148, 1184, 4096). GPU-validated: 14/14 (job 2390801). Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
The Literal annotation covers argparse, but programmatic TransformerConfig construction bypassed it and only failed inside enable_batch_invariant_mode() after model init. Guard at construction time against _BATCH_INVARIANT_BACKENDS (single source of truth). GPU-validated: 24/24 (job 2391673). Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Remaining 'floor' comment wording said the value could only shrink — the exact misreading behind the dropped-bucket bug. Reworded to 'align up'. The 64 literal now derives from DynamicInferenceContext.TOKEN_ROUNDER (lazy import; dynamic_context imports this module at top level) so the graph-bucket alignment and the eager path's padding multiple cannot drift apart. Behavior-identical. GPU-validated: 24/24 (job 2391676). Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
The early return also skipped _te_patch_general_grouped_gemm (a GEMM concern); moved it under skip_gemm where it belongs — under te_native grouped GEMMs stay native, covered by the same workspace starvation. Documented skip_rmsnorm in the docstring and rewrote the wrong call-site comment (te_native substitutes no TE kernels; the attention gate is a standalone assert, not part of this patch function). Both existing call sites behave identically. GPU-validated: 24/24 (job 2395658). Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Module-level pytestmark so CPU-only or Triton-less environments (e.g. the lts lane) report a legible skip instead of hard-failing at cuda allocation. Reuses the library's own HAVE_TRITON flag. GPU-validated: 24/24 (job 2395658). Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
The field is auto-registered from TransformerConfig so argparse always populates it; a silent fallback would mask a real wiring bug (and its hardcoded default had already drifted once). GPU-validated: 24/24 (job 2395658). Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
'Accumulates in fp32' was stale after acc_fp64 was added; it is the sentence a caller reads to judge reduction precision. Docstring-only. Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Add batch_invariant_collective ('ordered' | 'multimem') selecting the
cross-rank EP combine under batch_invariant_mode.
'ordered' (default) is the existing explicit fixed rank-order fp32 Triton
kernel: deterministic by construction on any hardware.
'multimem' keeps the native NVLS in-switch reduce-scatter used by the
non-invariant path. Measured properties (8xB200, symmetric-memory
microbenchmarks): the fp32 fabric reduction returns the correctly-rounded
exact sum -- bitwise-equal to an fp64 reference over 16.7M channels of
adversarial wide-spread data, where every software fp32 summation order
diverges (57-61% agreement) due to intermediate roundings. It is therefore
deterministic and batch-invariant, and any software path that must match it
bitwise (e.g. a training-side combine mirror) should accumulate in fp64 and
round once. Perf is equal to 'ordered' at EP8 (~1.02x on the collective) and
scales better at larger NVLink domains (one in-network reduction vs
WORLD_SIZE remote loads per channel).
The default stays 'ordered' pending confirmation that the exactly-rounded
fabric reduction is a design guarantee across NVSwitch generations rather
than a measured property of this platform. bf16 multimem remains excluded
(the bf16 fabric adder is a truncating fixed-width reducer).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
… exactly The batch-invariant training unpermute already groups contributions by EP rank to match the inference combine structurally; make it match BITWISE by mirroring the engine's precision at each stage of the tree: - Within-rank: when routing probabilities were applied at the activation (probs is None here — the gated/SwiGLU convention), the inference engine's _moe_sum accumulates in fp64 with unit weights; mirror in fp64 and round each rank partial to fp32, the dtype the engine stores into the RSV buffer before the collective. The weighted (probs-in-unpermute) path stays fp32, matching the engine's weighted fp32 reduction. - Cross-rank: match the configured batch_invariant_collective. 'multimem' returns the correctly-rounded exact fp32 sum of the partials (measured bitwise-equal to an fp64 reference over 16.7M adversarial channels), so the mirror accumulates partials in fp64 and rounds once; 'ordered' is an ascending rank-order fp32 chain, mirrored as-is. Plumb the collective selection through enable_batch_invariant_mode (and CLI) as the runtime source of truth so config-less call sites agree with the dispatcher. This is the training-side half of bitwise rollout==training (true on-policy RL); certified end-to-end (gen_kl == 0.0 over full-learning-rate GRPO) with the equivalent mirror in NeMo-RL. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Utkarsh Utkarsh <uutkarsh@nvidia.com>
Contributor
|
@utkarsh530 I am going to pull this into my fork to push through. It should maintain your authorship, Just easier for me to get it through ci |
Contributor
|
6 tasks
This was referenced Aug 27, 2026
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.
What does this PR do?
tacked on #6521. The first 23 commits are that PR; only the top 2 commits are new here. Will rebase to a clean 2-commit diff once #6521 merges; until then, please review only the top 2.
Batch-invariant mode currently forces the ordered software reduce-scatter for the expert-parallel combine, on the assumption that the in-switch NVLS multimem reduction has unspecified summation order and therefore unspecified bits. We measured its semantics, and it is stronger than ordered: the fp32 multimem reduce-scatter returns the correctly-rounded exact sum of its operands — its output matched fp32(fp64_exact_sum) on 100.000% of 16.7M adversarially-constructed channels. An exact sum has no order to vary, so it is batch-invariant by construction, at native speed and with no kernel to maintain.
This commit adds batch_invariant_collective: Literal["ordered", "multimem"] = "ordered" to TransformerConfig. The default is unchanged, and the ordered kernel remains the fallback for non-NVLS paths. bf16 multimem stays excluded — its accumulation truncates and the exact-sum property does not hold there.
Batch-invariance generation with vLLM fused-MoE backend at improved performance #6521 makes generation batch-invariant; true on-policy RL needs more: the training forward must reproduce generation's bits for the sampled tokens, so the importance ratio is exactly 1 and the policy gradient unbiased. The last structural difference between the two forwards is this combine: inference reduces in-switch, training transports and sums locally — two different summation trees over the same addends.
The exact-sum property is what makes closing this gap simple: an exact sum is reproducible by any fp64 accumulation order, so training does not need to imitate the switch — it only needs to sum exactly. Under batch-invariant mode, the training/scoring unpermute now accumulates in the matching tree: within-rank top-k partials in fp64 where routing probabilities were pre-applied at the activation, one fp32 rounding at the reduce-scatter boundary, cross-rank accumulation matched to the selected collective (multimem → fp64 exact-sum mirror; ordered → the ordered tree). Default paths untouched.
Result: generation ≡ training bitwise, certified in full-learning-rate GRPO — generation–training KL error exactly 0.0 at every step, per-token importance ratio identically 1.0, at a 2% generation overhead vs the same engine with determinism off (NeMo-RL integration: NVIDIA-NeMo/RL#3531).
Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.