Implement Quantile Balancing in MoE

[Mellonta]

Adding Quantile Balancing (QB) MoE.

Standard MoE load balancing adds an auxiliary loss whose gradient nudges the router toward uniform expert usage. QB instead enforces balance directly through a per-expert routing bias, with no auxiliary loss and no gradient contribution, load balancing is handled entirely by a bias that is updated once per global batch from a quantile estimate. From marin's report: https://openathena.ai/blog/quantile-balancing/

Conceptually, it can be understood as a two-part process run after each step, based on the scores that each token generates for each expert.

1. For each token, compute the threshold score required to become activated.
2. For each expert, using the result from step 1, determine the bias that would activate a balanced number of tokens.

Usage

--moe-router-load-balancing-type quantile_balancing
--moe-router-quantile-balancing-ema 0.0   # optional; 0.0 (default) = no memory, e.g. 0.9 to smooth
--moe-aux-loss-coeff 0                     # aux loss must be off

Verified on 2B models

[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.

[Mellonta]

/claude review

[claude]

LGTM

[Mellonta]

/ok to test b4dde89

[Mellonta]

/ok to test 4b2f1d7

[Mellonta]

/ok to test 2c1fadf

[Mellonta]

/ok to test 6e0ef96

[Mellonta]

/ok to test c1123ab

[Mellonta]

/ok to test f1dddce

[Mellonta]

/ok to test cd4aed8

[Mellonta]

Also fixes a pre-existing order-dependent hang in test_shared_expert_forward_backward: it seeded only CUDA but uses use_cpu_initialization, so inherited CPU-RNG state could drift in some ranks and cause a deadlock. now seeded via _set_random_seed.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29379425584

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/29381411801