[MXFP8/FP4-param-gather] Post processing after forced param AG in eval by WanZzzzzz · Pull Request #4562 · NVIDIA/Megatron-LM

WanZzzzzz · 2026-04-30T20:19:38Z

What does this PR do ?

The same PR to dev has been merged: #4563
Fixes FP8/MXFP8/FP4 parameter-gather convergence when eval runs between training steps.

When eval disables forward pre-hooks, start_param_sync(force_sync=True) can complete parameter all-gather synchronously. In that path there is no async gather handle for the next train forward hook to wait on (line), so the usual post-all-gather processing was skipped in next train step. In eval, the post processing is also skipped if the job enables overlap_param_gather (line). With FP8 (blockwise scaling) or FP4 param gather, this can leave columnwise weight storage stale while rowwise storage has been refreshed, causing the next train backward to use staled weights. With MXFP8 param+reuse_grad_buff_for_mxfp8_param_ag, the AG'ed weights will not be copied(quantized) from param buffer to model weights, which is not correct.

In this PR we add post processing in bucket group's start_param_sync for synchronous gather, hence making sure the above cases are correctly handled.

Validated the fix on DSv3 proxy model pretraining for FP4 param gather:

⚠️ 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-04-30T20:19:42Z

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-04-30T20:19:48Z

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.

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

WanZzzzzz · 2026-05-18T16:53:45Z

/ok to test b7a1243

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

deepakn94 · 2026-05-22T14:04:50Z

/claude review

claude

LGTM

Phlip79 · 2026-05-22T16:34:27Z

/ok to test a853e77

svcnvidia-nemo-ci · 2026-05-22T20:49:01Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-05-22T21:23:44Z

🔄 Merge queue validation started!

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

NVIDIA#4562) Signed-off-by: qiyuw <qiyuw@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

NVIDIA#4562) Signed-off-by: qiyuw <qiyuw@nvidia.com>

update colwise data after param AG in eval

a38ef9b

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

WanZzzzzz requested review from a team as code owners April 30, 2026 20:19

svcnvidia-nemo-ci marked this pull request as draft April 30, 2026 20:19

WanZzzzzz marked this pull request as ready for review April 30, 2026 20:31

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

svcnvidia-nemo-ci requested a review from a team April 30, 2026 20:32

svcnvidia-nemo-ci added the complexity: low label Apr 30, 2026

WanZzzzzz mentioned this pull request Apr 30, 2026

[MXFP8/FP4-param-gather] Post processing after forced param AG in eval #4563

Merged

5 tasks

WanZzzzzz added 6 commits May 8, 2026 11:13

fix similar issue for mxfp8 param gather

74be067

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

tighten the edge case

3e8fd8a

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

add eval transition unit test in test_fp8_param.py

c7849a5

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

add eval transition unit test in test_fp4_param.py

3b4a151

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

fix conflict

07324bc

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

lint

3ec63c9

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

WanZzzzzz changed the title ~~Update colwise data after synced param AG~~ [MXFP8/FP4-param-gather] Post processing after forced param AG in eval May 8, 2026

Phlip79 added the 26.06 label May 12, 2026

mirror fixes in mbridge

b7a1243

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

WanZzzzzz requested review from a team as code owners May 14, 2026 18:57

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

This was referenced May 15, 2026

[MXFP8] Mirror fixes in Mbridge for mxfp8 param gather #4818

Merged

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

Open

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

maanug-nv approved these changes May 19, 2026

View reviewed changes

WanZzzzzz added 2 commits May 21, 2026 13:57

Merge branch 'main' into qiyu/param_gather_eval_fix_main

13ac1c1

remove duplicated block

a853e77

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

WanZzzzzz requested a review from a team as a code owner May 21, 2026 21:10

claude Bot approved these changes May 22, 2026

View reviewed changes

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:34 Inactive

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

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

copy-pr-bot Bot temporarily deployed to public May 22, 2026 16:46 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 08bad7a May 22, 2026
82 of 84 checks passed

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

[MXFP8/FP4-param-gather] Post processing after forced param AG in eval (

2feb385

NVIDIA#4562) Signed-off-by: qiyuw <qiyuw@nvidia.com>

sbhavani mentioned this pull request May 26, 2026

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

Open

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

[MXFP8/FP4-param-gather] Post processing after forced param AG in eval (

808ac42

NVIDIA#4562) Signed-off-by: qiyuw <qiyuw@nvidia.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MXFP8/FP4-param-gather] Post processing after forced param AG in eval#4562

[MXFP8/FP4-param-gather] Post processing after forced param AG in eval#4562
Phlip79 merged 10 commits into
NVIDIA:mainfrom
WanZzzzzz:qiyu/param_gather_eval_fix_main

WanZzzzzz commented Apr 30, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Apr 30, 2026

Uh oh!

github-actions Bot commented Apr 30, 2026

Uh oh!

WanZzzzzz commented May 18, 2026

Uh oh!

deepakn94 commented May 22, 2026

Uh oh!

claude Bot left a comment

Uh oh!

Phlip79 commented May 22, 2026

Uh oh!

svcnvidia-nemo-ci 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!

6 participants

Conversation

WanZzzzzz commented Apr 30, 2026 • edited 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 Apr 30, 2026

Uh oh!

github-actions Bot commented Apr 30, 2026

Uh oh!

WanZzzzzz commented May 18, 2026

Uh oh!

deepakn94 commented May 22, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Phlip79 commented May 22, 2026

Uh oh!

svcnvidia-nemo-ci 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!

6 participants

WanZzzzzz commented Apr 30, 2026 •

edited

Loading