[MCore] Add global-batch quantile balancing by rkarimimahab · Pull Request #7394 · NVIDIA/Megatron-LM

rkarimimahab · 2026-09-16T09:01:54Z

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

What does this PR do?

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

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:

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
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"

When your PR is ready, click Ready for Review.
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 · 2026-09-16T09:01:58Z

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.

github-actions · 2026-09-16T09:02:04Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

jaredcasper · 2026-09-23T00:21:31Z

/ok to test bf2428b

jaredcasper · 2026-09-23T18:11:16Z

/ok to test 8895c55

jaredcasper · 2026-10-08T19:00:11Z

#7986 will make this failing test report only, after that goes in we can update and retry.

Signed-off-by: Harry Zhou <hhanyu@nvidia.com> Signed-off-by: rkarimimahab <rkarimimahab@nvidia.com>

Signed-off-by: Rabeeh Mahabadi <rkarimimahab@nvidia.com>

Signed-off-by: Rabeeh Mahabadi <rkarimimahab@nvidia.com> Signed-off-by: rkarimimahab <rkarimimahab@nvidia.com>

Signed-off-by: Rabeeh Mahabadi <rkarimimahab@nvidia.com>

Signed-off-by: Rabeeh Mahabadi <rkarimimahab@nvidia.com> Signed-off-by: rkarimimahab <rkarimimahab@nvidia.com>

Signed-off-by: rkarimimahab <rkarimimahab@nvidia.com>

jaredcasper · 2026-10-22T16:53:53Z

/ok to test eb9dfdb

jaredcasper · 2026-10-22T17:55:58Z

/ok to test d7cd30d

nemo-automation-bot · 2026-10-22T20:06:06Z

🔄 Merge queue validation started!

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

Seventeen upstream commits since 9e5aad0, none touching a file this branch changes; no conflicts. The merge brings in NVIDIA#8414 ("Fix import error by DSA spec"), which repairs the import of megatron.core.transformer.experimental_attention_variant.dsa_layer_specs that made the Installation Test fail on main from 2026-10-21 up to and including 9e5aad0, and therefore on this branch's head 09cced3. Also in: NVIDIA#8404 removed eight gpt functional-test cases (the corpus-parity gate now covers 297 configs), NVIDIA#7394 added two MoE router flags (outside the data group; the data-arguments snapshot is unchanged), NVIDIA#7926 removed an inference flag. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: asolergi-nv <asolergibert@nvidia.com>

Repairs semantic interactions between the imported main slice (tip f9d9bdc) and dev dd8f93b found by the baseline sweep and review: - experimental_attention_variant/dsa_gqa.py: mark main NVIDIA#7156's DSGQACoreAttention as a DSA indexer-loss metric writer (logs_dsa_indexer_loss = True, as DSAttention and CSA do) so dev's fixed-capacity DSA metric tracker reserves capacity for it. - moe/experts.py: assign TEGroupedMLP.activation_recompute (and the fp8/fp4 set_save_original_input(linear_fc2) hook) before the fused-impl support check reads it; read main's _grouped_linear_kwargs with a getattr default in _make_fused_ops, matching the existing getattr(self, "activation_recompute", False) convention. - moe/router.py: restore main NVIDIA#7394's global-batch branch in TopKRouter.quantile_balancing(), gated on dev's use_quantile_balancing. - multi_token_prediction.py: black formatting (AST-equivalent). - test_attention_variant_dsa_gqa.py: cover metric-tracker initialization followed by a DSA-GQA training forward, with and without indexer loss. - test_fine_grained_activation_offloading.py: map expert_fc1 + moe_act to main's fused_group_mlp offload-summary group. - test_multi_token_prediction.py: apply main NVIDIA#8211's cfg.finalize() pattern in the dev-only TestMHCMTPIntegration. Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

rkarimimahab requested review from a team as code owners September 16, 2026 09:01

svcnvidia-nemo-ci marked this pull request as draft September 16, 2026 09:02

rkarimimahab marked this pull request as ready for review September 16, 2026 09:03

jaredcasper approved these changes Sep 16, 2026

View reviewed changes

yaox12 approved these changes Sep 17, 2026

View reviewed changes

guihong-nv approved these changes Sep 18, 2026

View reviewed changes

kvareddy approved these changes Sep 18, 2026

View reviewed changes

svcnvidia-nemo-ci added the Final Review PR is in the "final review" stage label Sep 18, 2026

lhb8125 mentioned this pull request Sep 21, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

copy-pr-bot Bot temporarily deployed to public September 23, 2026 00:25 Inactive

ericharper approved these changes Sep 23, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Sep 23, 2026

copy-pr-bot Bot temporarily deployed to public September 23, 2026 18:17 Inactive

ericharper enabled auto-merge September 28, 2026 23:24

auto-merge was automatically disabled September 30, 2026 21:51
Head branch was pushed to by a user without write access

copy-pr-bot Bot temporarily deployed to public October 8, 2026 13:40 Inactive

harryzhou2000 and others added 6 commits October 22, 2026 04:10

[MCore] Add global-batch quantile balancing

cd5b1c1

Signed-off-by: Harry Zhou <hhanyu@nvidia.com> Signed-off-by: rkarimimahab <rkarimimahab@nvidia.com>

Fix CUDA graph capture state and NCCL EP receive buffer reuse

f5e3b3a

Signed-off-by: Rabeeh Mahabadi <rkarimimahab@nvidia.com>

Scope failed CUDA graph cleanup to started TE capture

e3c1a22

Signed-off-by: Rabeeh Mahabadi <rkarimimahab@nvidia.com> Signed-off-by: rkarimimahab <rkarimimahab@nvidia.com>

Document and test stale capture stream override policy

bafd7c5

Signed-off-by: Rabeeh Mahabadi <rkarimimahab@nvidia.com>

Validate quantile balancing top-k in both estimation scopes

d6e4120

Signed-off-by: Rabeeh Mahabadi <rkarimimahab@nvidia.com> Signed-off-by: rkarimimahab <rkarimimahab@nvidia.com>

Preserve quantile balancing replay coverage and fix lint after rebase

eb9dfdb

Signed-off-by: rkarimimahab <rkarimimahab@nvidia.com>

rkarimimahab force-pushed the rkarimimahab/k3-quantile-balancing branch from e423010 to eb9dfdb Compare October 22, 2026 11:14

jaredcasper enabled auto-merge October 22, 2026 16:54

copy-pr-bot Bot temporarily deployed to test October 22, 2026 16:56 Inactive

copy-pr-bot Bot temporarily deployed to public October 22, 2026 16:59 Inactive

Merge branch 'main' into rkarimimahab/k3-quantile-balancing

d7cd30d

copy-pr-bot Bot temporarily deployed to test October 22, 2026 17:58 Inactive

copy-pr-bot Bot temporarily deployed to public October 22, 2026 17:59 Inactive

jaredcasper added this pull request to the merge queue Oct 22, 2026

Merged via the queue into NVIDIA:main with commit 92641a4 Oct 22, 2026
152 of 154 checks passed

fanshiqing mentioned this pull request Oct 23, 2026

[GTP] Record all asynchronous gather completion events at issuance #8428

Merged

svcnemo-autobot mentioned this pull request Oct 23, 2026

chore: AUT-2616 nightly sync main into dev (2026-10-23) #8461

Merged

Navigation Menu

Repository navigation

[MCore] Add global-batch quantile balancing - #7394

[MCore] Add global-batch quantile balancing#7394
jaredcasper merged 7 commits into
NVIDIA:mainfrom
rkarimimahab:rkarimimahab/k3-quantile-balancing

rkarimimahab commented Sep 16, 2026

Uh oh!

copy-pr-bot Bot commented Sep 16, 2026

Uh oh!

github-actions Bot commented Sep 16, 2026

Uh oh!

jaredcasper commented Sep 23, 2026

Uh oh!

jaredcasper commented Sep 23, 2026

Uh oh!

jaredcasper commented Oct 8, 2026

Uh oh!

jaredcasper commented Oct 22, 2026

Uh oh!

jaredcasper commented Oct 22, 2026

Uh oh!

nemo-automation-bot Bot commented Oct 22, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants

Conversation

rkarimimahab commented Sep 16, 2026

What does this PR do?

Issue tracking

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Sep 16, 2026

Uh oh!

github-actions Bot commented Sep 16, 2026

Uh oh!

jaredcasper commented Sep 23, 2026

Uh oh!

jaredcasper commented Sep 23, 2026

Uh oh!

jaredcasper commented Oct 8, 2026

Uh oh!

jaredcasper commented Oct 22, 2026

Uh oh!

jaredcasper commented Oct 22, 2026

Uh oh!

nemo-automation-bot Bot commented Oct 22, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

11 participants