[MCore] Add global-batch quantile balancing

[harryzhou2000]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Extend main's existing Quantile Balancing router with Kimi K3 global-batch histogram estimation while preserving the current exact micro_batch estimator and its default behavior.

Paired dev-target PR: #6637

Quantile estimator

For token i and expert j, let s[i,j] be the raw sigmoid router score and let alpha[i] be the (k+1)-th largest value of s[i,:] + b[:]. K3 defines:

r[i,j]       = alpha[i] - s[i,j]
b_next[j]    = quantile(r[:,j], k / num_experts)
b_next[:]   -= mean(b_next[:])

The histogram range for the next global batch is [min(b_next)-1, max(b_next)+1]. The default is 1000 uniform bins per expert, matching the report.

Reference: Kimi K3 technical report, Quantile Balancing.

Implementation

• Add --moe-router-quantile-balancing-estimation-scope {micro_batch,global_batch}:
  • micro_batch remains the default and retains the existing exact estimator;
  • global_batch enables K3's pooled histogram estimator.
• Require --moe-aux-loss-coeff 0 for QB.
• Accumulate one persistent-address int32 histogram [num_experts, num_bins] per MoE router across all gradient-accumulation microbatches.
• Keep the additive global-batch QB bias and histogram bounds as persistent FP32 module buffers; keep the temporary histogram non-persistent in checkpoints.
• Reuse the router's existing expert_bias buffer without enabling --moe-router-enable-expert-bias; that flag continues to select the independent signed-count updater.
• At gradient finalization, sum histograms across the appropriate TP/DP/CP groups, recover the pooled quantile, update bias and bounds in place exactly once, and reset the histogram without changing buffer addresses.
• Preserve the existing GTP weight-rematerialization reduction topology.
• Reset QB histograms before a paged-stash dropless capacity rerun so replayed microbatches are counted exactly once.
• Treat a negative expert-capacity sentinel as disabled for QB, reject active capacity for QB, and preserve the existing validation for non-QB modes such as Sinkhorn.
• Support both router implementations:
  • fused routing feature-detects the histogram API from [Common][PyTorch] Add QB router histogram paths TransformerEngine#3395 and selects its single-kernel fused_atomic path;
  • non-fused routing performs top-(k+1) and histogram accumulation with PyTorch operations.

No full score tensor is retained. The fused path adds shared-memory bin classification and histogram atomics to the existing router kernel. The global-batch finalizer communicates num_experts * num_bins int32 counters per router; the non-fused fallback additionally launches PyTorch top-(k+1) and scatter operations.

Validation

Validated on NVIDIA B300 with NVIDIA/TransformerEngine#3395 installed:

• 18 focused single-GPU QB tests passed (the six skips are eight-rank-only cases), covering existing micro_batch behavior plus report-derived histogram recovery, pooled-global-batch semantics, fused-vs-non-fused parity, paged-stash retry reset, and real MCore router finalization.
• The real-router global-batch finalization test passed under DP2 on both ranks with an actual NCCL histogram all-reduce.
• The fused test uses qb_histogram_mode="fused_atomic" and compares its histogram exactly with the non-fused implementation.
• Rebased-tip regressions: 38 router tests passed, 6 skipped; 4 finalizer tests passed; 3 paged-stash tests passed.

Additional eight-rank validation on NVIDIA B300 exercised the production finalize_model_grads path for:

• TP1/EP4: dense DP8 and expert DP2.
• TP1/EP8: dense DP8 and expert DP1.
• TP4/EP2 with sequence parallelism: dense DP2, expert DP1, and expert TP4.

For every topology, the test verifies model-parallel group sizes, per-microbatch accumulation, non-identical pre-reduction rank-local histograms, the all-reduced histogram against an independent quantile-update oracle, identical bias and bounds on all ranks, in-place reset, stable CUDA-graph-visible addresses, and exactly one gradient synchronization. All six fused/unfused parameterizations passed on all eight ranks. The focused Transformer Engine router suite passed 38 tests with 2 skips.

The commits are SSH-signed and signed off. git diff --check and the repository tools/autoformat.sh pass; pylint rates the changed files 10/10 and Ruff reports no errors.

Scope and dependencies

The existing micro_batch path and non-fused global_batch path do not depend on NVIDIA/TransformerEngine#3395. If fused global-batch QB is requested, MCore checks the installed TE signature and reports a targeted error when the histogram API is unavailable.

This PR does not depend on SiTU-GLU or #6674. It has no SiTU activation imports or configuration coupling; the paired integration checkout was used only to verify compatibility between the independent branches.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue: TODO — link a Megatron-LM feature-request issue before marking this draft ready for review.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

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"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
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, the Final Review label 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 Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[harryzhou2000]

Applied the same review fixes as the paired dev PR: #6637.

Changes:

• Prevent duplicate CLI registration and test the complete Megatron argument parser.
• Document quantile_balancing, both supported accumulation scopes, the required zero auxiliary-loss coefficient, and histogram-bin control.
• Improve invalid bin_bounds diagnostics.
• Add real non-fused MoELayer forward/backward/finalization coverage.
• Cover eval, frozen-bias, no-grad, and activation-checkpoint recomputation behavior.

Validation on NVIDIA B200 with PyTorch 2.13 nightly and Transformer Engine 2.19 development:

• Focused single-GPU QB suite: 14 passed, 3 skipped (the skips are the 8-rank-only topology cases).
• Eight-rank distributed cases: 3 passed per rank, covering EP4 with two data replicas, EP8, and TP4+EP2.
• Formatting and static checks passed for the changed files.

[harryzhou2000]

Mirrored the QB review fixes from the paired dev PR #6637 in 5a88a7cd0: documented sigmoid routing, singleton-list normalization, explicit legacy process-group fallback, corrected diagnostics/invariant comment, QB-only capacity handling, and unfused plus fused-atomic eight-rank coverage. Final NVIDIA B300 results are 18 passed single-GPU plus all 6 distributed topology cases on every rank.