Fix mxfp8 param gather numerical issue when DP overlap is off by WanZzzzzz · Pull Request #4800 · NVIDIA/Megatron-LM

WanZzzzzz · 2026-05-14T16:24:52Z

What does this PR do ?

With MXFP8 grad-buffer reuse and non-overlap param gather, each DistOpt stages its own updated main-param shards into its param buffers during step. However, param sync is a DDP model-chunk operation: model_chunk.start_param_sync() gathers both dense and expert bucket groups, copies gathered values into model weights, and zeros the shared MXFP8 param/grad buffers. For MoE, dense and expert DistOpts may share the same model chunk, so defer param sync until all chained optimizers have staged their params, then sync each model chunk once. Otherwise the expert DistOpt may copy the zeroes into expert model weights.

⚠️ 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 the @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
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 the @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.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

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

copy-pr-bot · 2026-05-14T16:24:56Z

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-05-14T16:25:01Z

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.

WanZzzzzz · 2026-05-18T16:36:57Z

/ok to test fc3471a

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

deepakn94 · 2026-05-22T14:22:16Z

+        deferred_bucket_groups = []
+        deferred_bucket_group_ids = set()
+
+        if defer_param_sync:


Can we re-factor this into a separate method? This is making the optimizer.step_with_ready_grads method extremely complicated...

Phlip79 · 2026-05-22T15:50:44Z

/ok to test 0479d84

Phlip79 · 2026-05-22T15:53:37Z

/ok to test c2b4f6d

Phlip79 · 2026-05-22T16:31:43Z

/ok to test 752d246

svcnvidia-nemo-ci · 2026-05-22T19:56:18Z

🔄 Merge queue validation started!

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

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

* origin/main: (50 commits) Drain predecessor reduce-scatter at dispatch time (NVIDIA#4940) ci: Add allow_failure flag to gpt and moe recipes that are failing in nightlies (NVIDIA#4905) fix(tests): initialize num_microbatches calculator in vision cudagraph tests (NVIDIA#4986) test: re-enable test_pp2_create_cudagraphs_first_stage on TE 2.15+ (NVIDIA#4985) ci: Add support for MBridge job gating based on PR labels (NVIDIA#4926) test(ci): re-enable 8experts2parallel_multi_dist_optimizer_instances_1node (NVIDIA#4984) test: re-enable paged stashing MoE tests (NVIDIA#4978) Fix elastification unwrap_model import (NVIDIA#4972) Avoid offsetting functional test master port (NVIDIA#4973) test: enable NVTE_CUTEDSL_FUSED_GROUPED_MLP via pytest fixture (NVIDIA#4931) chore(beep boop 🤖): Bump (main) (2026-05-25) test(release): add release goldens for deepseekv3/nemotron3 and set tp2pp2 exit-interval (NVIDIA#4932) Fix `get_batch` return order to ignore BlendedDataset provenance fields (NVIDIA#4952) ci: restore perf test torchrun logs (NVIDIA#4951) Various training utils (NVIDIA#4872) ci: Update training script paths in BERT and T5 (NVIDIA#4939) [MXFP8/FP4-param-gather] Post processing after forced param AG in eval (NVIDIA#4562) Fix mxfp8 param gather numerical issue when DP overlap is off (NVIDIA#4800) Add TEFusedDenseMLP for Dense+Grouped GEMM fusion on SM100+ (NVIDIA#4318) (NVIDIA#4786) Fix paged stashing test submodules lookup (NVIDIA#4925) ... # Conflicts: # megatron/training/training.py

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>

WanZzzzzz added 4 commits May 14, 2026 09:23

fix mxfp8 without overlap

af02072

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

minor

79d0bc8

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

unit test for mxfp8 on moe model

097f85f

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

lint

fc3471a

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

WanZzzzzz requested review from a team as code owners May 14, 2026 16:24

svcnvidia-nemo-ci marked this pull request as draft May 14, 2026 16:25

Phlip79 added the 26.06 label May 14, 2026

Phlip79 requested a review from a team May 14, 2026 17:58

WanZzzzzz marked this pull request as ready for review May 14, 2026 17:59

svcnvidia-nemo-ci requested a review from a team May 14, 2026 17:59

svcnvidia-nemo-ci added Final Review PR is in the "final review" stage complexity: low labels May 14, 2026

Phlip79 removed the Final Review PR is in the "final review" stage label May 15, 2026

This was referenced May 15, 2026

Fix mxfp8 param gather numerical issue when DP overlap is off #4769

Merged

[bug] Investigate convergence of performance features with Qwen3.5 VL as proxy model NVIDIA-NeMo/Megatron-Bridge#3801

Closed

copy-pr-bot Bot temporarily deployed to public May 18, 2026 16:37 Inactive

copy-pr-bot Bot temporarily deployed to test May 18, 2026 16:38 Inactive

fix conflicts

fc3fb7f

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

deepakn94 reviewed May 22, 2026

View reviewed changes

Refactor chained optimizer param sync handling

0479d84

copy-pr-bot Bot temporarily deployed to public May 22, 2026 15:51 Inactive

Format optimizer refactor

c2b4f6d

Phlip79 requested a review from deepakn94 May 22, 2026 15:54

copy-pr-bot Bot temporarily deployed to public May 22, 2026 15:54 Inactive

copy-pr-bot Bot temporarily deployed to public May 22, 2026 16:32 Inactive

copy-pr-bot Bot temporarily deployed to test May 22, 2026 16:32 Inactive

deepakn94 approved these changes May 22, 2026

View reviewed changes

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

Phlip79 enabled auto-merge May 22, 2026 16:34

copy-pr-bot Bot temporarily deployed to public May 22, 2026 16:35 Inactive

copy-pr-bot Bot temporarily deployed to public May 22, 2026 16:43 Inactive

ericharper approved these changes May 22, 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 May 22, 2026

Phlip79 added this pull request to the merge queue May 22, 2026

Merged via the queue into NVIDIA:main with commit 5f79118 May 22, 2026
81 of 83 checks passed

ko3n1g mentioned this pull request May 26, 2026

fix(optimizer): gate ChainedOptimizer MXFP8 defer-sync on DDP-level overlap_param_gather #4982

Merged

4 tasks

santhnm2 pushed a commit to santhnm2/Megatron-LM that referenced this pull request May 26, 2026

Fix mxfp8 param gather numerical issue when DP overlap is off (NVIDIA…

530f665

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

sbhavani mentioned this pull request May 26, 2026

[ROADMAP][2026 Q2] Megatron Core Roadmap #4997

Closed

janEbert pushed a commit to janEbert/Megatron-LM that referenced this pull request Jun 2, 2026

Fix mxfp8 param gather numerical issue when DP overlap is off (NVIDIA…

942b867

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

mathemakitten pushed a commit to mathemakitten/Megatron-LM that referenced this pull request Jun 12, 2026

Fix mxfp8 param gather numerical issue when DP overlap is off (NVIDIA…

2ac0215

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

kunlunl mentioned this pull request Jul 7, 2026

[dev] Muon fp8 param gather for decoupled layout #5470

Merged

6 tasks

terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026

Fix mxfp8 param gather numerical issue when DP overlap is off (NVIDIA…

80580db

…#4800) Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

Navigation Menu

Uh oh!

Fix mxfp8 param gather numerical issue when DP overlap is off - #4800

Fix mxfp8 param gather numerical issue when DP overlap is off#4800
Phlip79 merged 8 commits into
NVIDIA:mainfrom
WanZzzzzz:fix_mxfp8_overlap_off

WanZzzzzz commented May 14, 2026 •

edited by zhongbozhu

Loading

Uh oh!

copy-pr-bot Bot commented May 14, 2026

Uh oh!

github-actions Bot commented May 14, 2026

Uh oh!

WanZzzzzz commented May 18, 2026

Uh oh!

deepakn94 May 22, 2026

Uh oh!

Phlip79 commented May 22, 2026

Uh oh!

Phlip79 commented May 22, 2026

Uh oh!

Phlip79 commented May 22, 2026

Uh oh!

svcnvidia-nemo-ci commented May 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

WanZzzzzz commented May 14, 2026 • edited by zhongbozhu Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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 May 14, 2026

Uh oh!

github-actions Bot commented May 14, 2026

Uh oh!

WanZzzzzz commented May 18, 2026

Uh oh!

deepakn94 May 22, 2026

Choose a reason for hiding this comment

Uh oh!

Phlip79 commented May 22, 2026

Uh oh!

Phlip79 commented May 22, 2026

Uh oh!

Phlip79 commented May 22, 2026

Uh oh!

svcnvidia-nemo-ci commented May 22, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

WanZzzzzz commented May 14, 2026 •

edited by zhongbozhu

Loading