[Dev] Add global-batch quantile balancing

[harryzhou2000]

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

What does this PR do?

Add Kimi K3 global-batch Quantile Balancing (QB) as an auxiliary-loss-free MoE routing mode on dev.

Paired main-target PR: #6638

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 quantile_balancing as a sole router load-balancing mode and require --moe-aux-loss-coeff 0.
• Add --moe-router-quantile-balancing-estimation-scope global_batch and --moe-router-qb-num-bins.
• Accumulate one persistent-address int32 histogram [num_experts, num_bins] per MoE router across all gradient-accumulation microbatches.
• Keep the additive 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 TP/DP/CP, recover the pooled quantile, update bias and bounds in place exactly once, and reset the histogram without changing buffer addresses.
• Reset QB histograms before a paged-stash dropless capacity rerun so replayed microbatches are counted exactly once.
• 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.
• 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.

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 report-derived histogram recovery, pooled-global-batch semantics, fused-vs-non-fused histogram parity, paged-stash retry reset, configuration validation, and real MCore router forward/backward/finalization.
• The real-router 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: 39 router tests passed, 2 skipped; the expert-bias finalizer test 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 non-fused QB path has no dependency on NVIDIA/TransformerEngine#3395. If --moe-router-fusion is enabled, 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 #6673. 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.

[yuzhongw-nvidia]

/claude review

[claude]

Light review — one blocking issue plus a few smaller notes.

Blocking: duplicate CLI flags will break parse_args()

--moe-router-quantile-balancing-estimation-scope and --moe-router-qb-num-bins are added manually in _add_moe_args, but the corresponding TransformerConfig fields are not in the exclude list at megatron/training/arguments.py:2729-2824. ArgumentGroupFactory therefore already generates both flags in _add_network_size_args (which runs first), so argparse hits a duplicate option string at parser-construction time. See the inline comment for the fix.

This is presumably why the CLI test in the PR builds a bare parser with only _add_moe_args rather than the full Megatron parser — that test passes while real training would fail at startup. Worth switching test_qb_cli_exposes_global_batch_scope_and_histogram_bins to use the full add_megatron_arguments(parser) so a regression like this is caught.

Documentation

megatron/core/transformer/moe/README.md documents the load-balancing strategies but was not updated. Two spots:

Load Balancing Strategies table (around line 300) — add a row such as:

| **quantile_balancing** | Kimi K3 aux-loss-free global-batch histogram quantile bias updates | `--moe-router-load-balancing-type quantile_balancing --moe-aux-loss-coeff 0` |

Router Arguments table (around line 531): the --moe-router-load-balancing-type row still reads "aux_loss, sinkhorn, seq_aux_loss, none" — it already omits global_aux_loss, and now quantile_balancing too. Adding rows for the two new flags would help as well.

Test coverage

The unit tests are thorough for the histogram math and the finalizer, but three of the six substantive tests are gated behind fused_topk_with_score_function_supports_qb, which is False against released TE. In CI today, effectively only the pure-math and non-fused paths run. Two gaps worth closing with tests that do not need the TE patch:

1. A real-router non-fused (moe_router_fusion=False) forward/backward plus finalize test. test_qb_mcore_router_accumulates_microbatches_and_finalizes is the only test that constructs an actual MoELayer, and it requires the TE QB API — so the Router.__init__ buffer registration and routing() wiring are currently untested in CI.
2. The accumulate_qb_histogram gating in Router.routing() — specifically that eval mode and frozen_expert_bias=True do not accumulate.

Also, no functional test was added (the PR checklist notes this). Given this changes the router bias-update path, at least one tests/functional_tests/test_cases/moe/ case would be valuable before this leaves draft.

Minor

_update_router_expert_bias_with_quantile(model, config, tp_dp_cp_group=tp_dp_cp_group) at finalize_model_grads.py:601-603 fits on one line under the 100-char limit; tools/autoformat.sh (black) will collapse it.

The unrelated blank-line removal at router.py:199 touches code outside the PR scope — per docs/developer/contribute.md, worth reverting.

One thing I could not verify locally: whether torch.is_grad_enabled() is sufficient to exclude activation recomputation from histogram accumulation. Raised inline.

[harryzhou2000]

Addressed the latest review on the paired dev/main changes (main counterpart: #6638).

Changes:

• Prevent duplicate CLI registration by excluding the two manually defined QB fields from ArgumentGroupFactory; the CLI test now constructs the complete Megatron argument parser.
• Document quantile_balancing, its required zero auxiliary-loss coefficient, the global-batch scope, and histogram-bin control.
• Improve invalid bin_bounds diagnostics with the exact expected shape.
• Exercise the real non-fused MoELayer router path through forward, backward, and bias finalization.
• Cover eval, frozen-bias, and no-grad accumulation gates.
• Cover activation checkpointing with the real router, asserting that the no-grad original forward records nothing and the grad-enabled backward recomputation records exactly one token sample set.
• Restore the unrelated whitespace change and apply the requested formatting cleanup.

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

• Focused single-GPU QB suite: 15 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.
• py_compile, git diff --check, isort, and Black checks passed for the changed Python files.

I kept the coverage focused in the QB unit/distributed suite rather than adding a separate golden-value functional recipe: the updated tests now execute the production MoELayer routing, backward, finalization, checkpoint-recompute, and multi-rank reduction paths without requiring a full training recipe.

[yuzhongw-nvidia]

/claude review

[FDecaYed]

A few non-blocking comments from the current head.

[FDecaYed]

/ok to test 2d6521b

[claude]

Light review — the implementation looks coherent and the unit tests are thorough. Two things worth addressing before this leaves draft, plus one inline nit.

1. The distributed and fused tests are gated on an unmerged TE PR, so neither runs in CI

test_qb_world8_ep_topologies_finalize_model_grads is the only test that exercises the actual finalize_model_grads reduction across EP/TP topologies, but it carries:

@pytest.mark.skipif(
    not torch.cuda.is_available() or not fused_topk_with_score_function_supports_qb,
    ...
)

and configures moe_router_fusion=True. Since fused_topk_with_score_function_supports_qb is False until TE#3395 lands, this test — and test_qb_fused_atomic_matches_unfused_histogram — will skip on every CI run. The non-fused QB path has no TE dependency at all, so the topology test can cover the distributed histogram all-reduce today. Consider parametrizing fused the way test_qb_mcore_router_accumulates_microbatches_and_finalizes already does, keeping the TE gate only on the True case:

@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
@pytest.mark.internal
@pytest.mark.parametrize(
    "fused",
    [
        False,
        pytest.param(
            True,
            marks=pytest.mark.skipif(
                not fused_topk_with_score_function_supports_qb,
                reason="requires the Transformer Engine QB fused-router API",
            ),
        ),
    ],
)
@pytest.mark.parametrize(
    "tp_size,ep_size,dense_dp_size,expert_dp_size", [(1, 4, 8, 2), (1, 8, 8, 1), (4, 2, 2, 1)]
)
def test_qb_world8_ep_topologies_finalize_model_grads(
    fused, tp_size, ep_size, dense_dp_size, expert_dp_size
):
    ...
    config = _config(..., moe_router_fusion=fused, ...)

That way the reduction-group, in-place-reset, and cross-rank-agreement assertions all become live CI coverage instead of skips.

2. moe_expert_capacity_factor validation is silently loosened for non-QB configs

transformer_config.py:2148-2152 changes if isinstance(...) to elif isinstance(...):

if self.moe_expert_capacity_factor is not None:
    if self.moe_expert_capacity_factor < 0:
        self.moe_expert_capacity_factor = None
    elif isinstance(self.moe_router_load_balancing_type, list):   # was `if`

Previously a negative sentinel was normalized to None and the load-balancing-type check still ran, so e.g. moe_expert_capacity_factor=-1 with sinkhorn raised. It now passes. Skipping the check when capacity is disabled is defensible, but it is a behavior change for configs unrelated to quantile balancing — worth calling out in the commit message, or scoping to the QB path if unintended.

---

Minor: get_updated_expert_bias_with_quantile reads total_tokens = tokens_per_expert[..., 0], relying on expert 0's total standing in for the whole batch. That holds because every token contributes exactly one margin sample per expert, but a short comment stating the invariant would keep the next reader from treating it as a bug.

Nothing flagged on process-group hygiene: the new parallel_state.get_tensor_and_data_parallel_group() call at finalize_model_grads.py:598 mirrors the existing moe_router_enable_expert_bias fallback a few lines above and only fires when pg_collection is None, so it sits inside the established compatibility path.

[harryzhou2000]

Addressed the current QB review findings on both target branches in eeb522774 (dev) and 5a88a7cd0 (main):

• Added the required sigmoid score-function flag to the documented command.
• Normalized the direct-config singleton QB list without accepting actual load-balancing combinations.
• Kept pg_collection.tp_dp_cp as the modern process-group path and explicitly labeled the global lookup as legacy compatibility.
• Corrected the expert-bias shape diagnostic and documented why every expert histogram has the same total.
• Parametrized the eight-rank topology fence over unfused and TE fused-atomic routing.
• Scoped negative-capacity handling to QB, rejected active QB capacity, and added a Sinkhorn regression proving non-QB behavior is unchanged.

QB continues to reuse the router's expert_bias buffer without enabling --moe-router-enable-expert-bias; that flag selects the separate DeepSeek-style signed-count updater.

Final NVIDIA B300 results on each branch: 18 passed in the single-GPU suite (six world-size-gated skips), then all 6 EP4/EP8/TP4+EP2 fused/unfused cases passed on every rank.

[FDecaYed]

/ok to test eeb5227