Skip to content

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

Merged
ericharper merged 6 commits into
NVIDIA:devfrom
WanZzzzzz:qiyu/param_gather_eval_fix_dev
May 7, 2026
Merged

[MXFP8/FP4-param-gather] Post processing after forced param AG in eval#4563
ericharper merged 6 commits into
NVIDIA:devfrom
WanZzzzzz:qiyu/param_gather_eval_fix_dev

Conversation

@WanZzzzzz

@WanZzzzzz WanZzzzzz commented Apr 30, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

What does this PR do ?

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:
image

⚠️ 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"

  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.

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.

@WanZzzzzz
update colwise data after param AG in eval
14ff992 
Signed-off-by: qiyuw <qiyuw@nvidia.com>
@WanZzzzzz
WanZzzzzz requested review from a team as code owners April 30, 2026 20:34
@copy-pr-bot

copy-pr-bot Bot commented Apr 30, 2026

Copy link
Copy Markdown

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.

@WanZzzzzz WanZzzzzz mentioned this pull request May 6, 2026
Merged
5 tasks
WanZzzzzz added 2 commits May 6, 2026 16:20
@WanZzzzzz
fix similar issue for mxfp8 param gather
a61efb0 
Signed-off-by: qiyuw <qiyuw@nvidia.com>
@WanZzzzzz
tighten the edge case
c8e5883 
Signed-off-by: qiyuw <qiyuw@nvidia.com>
@WanZzzzzz WanZzzzzz changed the title Update colwise data after param AG in eval [MXFP8/FP4-param-gather] Post processing after forced param AG in eval May 6, 2026
@ericharper
ericharper enabled auto-merge May 6, 2026 23:40
WanZzzzzz added 3 commits May 7, 2026 13:29
@WanZzzzzz
add eval transition unit test in test_fp8_param.py
be2d62b 
Signed-off-by: qiyuw <qiyuw@nvidia.com>
@WanZzzzzz
add eval transition unit test in test_fp4_param.py
433d8fb 
Signed-off-by: qiyuw <qiyuw@nvidia.com>
@WanZzzzzz
lint
da05c93 
Signed-off-by: qiyuw <qiyuw@nvidia.com>
@WanZzzzzz

WanZzzzzz commented May 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test da05c93

@ericharper
ericharper added this pull request to the merge queue May 7, 2026
@svcnvidia-nemo-ci

svcnvidia-nemo-ci commented May 7, 2026

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

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

Merged via the queue into NVIDIA:dev with commit 4401ac8 May 7, 2026
64 of 65 checks passed
@WanZzzzzz WanZzzzzz mentioned this pull request May 8, 2026
Merged
4 of 5 tasks
Phlip79 added a commit that referenced this pull request May 11, 2026
@Phlip79
fix: restore dev-only changes lost by main-priority merge
dc4ca1f 
Brings forward six regressions previously fixed manually on PR #4659
(commits 676f3fa, d019432, 2207908) that the new sync of main
into dev re-introduced, plus one further dropped import that surfaced
as a runtime NameError on gpt_dynamic_inference CI:

- transformer_layer.py: add input_ids=None to _forward_mlp_router so the
  body's input_ids reference (line 2103) resolves; callers already pass
  input_ids from _forward_mlp.
- param_and_grad_buffer.py: replace inline fp8_params else-branch in
  finish_param_sync with self._post_param_sync() to match the three
  other call sites and dev's refactor (PR #4563).
- moe_layer.py: call _maybe_record_overload_factor() on dispatched_input
  before the routed-experts conditional, as on dev.
- token_dispatcher.py: pass
  num_sms_preprocessing_api=self.config.moe_hybridep_num_sms_preprocessing
  to _HybridEPManager constructor.
- datasets/readme.md: restore the dev-only "Packing Scheduler"
  documentation section before main's "Offline cache preparation".
- gpt_dynamic_inference_with_coordinator.py: restore
  `from megatron.training.arguments import parse_and_validate_args`
  import; the call at line 216 raised NameError on the
  gpt_dynamic_inference_tp2_pp2_dp2_583m_logitsmatch_zmq CI job.

Verified: ast.parse OK on all 5 changed .py files; ruff F821 reports
only pre-existing TYPE_CHECKING-guarded names; CODEOWNERS and the
pyproject.toml/uv.lock/Dockerfile.ci.dev triple unchanged from
origin/dev.
Phlip79 added a commit that referenced this pull request May 11, 2026
@Phlip79
fix: restore dev-only changes lost by main-priority merge
bfce39d 
Brings forward six regressions previously fixed manually on PR #4659
(commits 676f3fa, d019432, 2207908) that the new sync of main
into dev re-introduced, plus one further dropped import that surfaced
as a runtime NameError on gpt_dynamic_inference CI:

- transformer_layer.py: add input_ids=None to _forward_mlp_router so the
  body's input_ids reference (line 2103) resolves; callers already pass
  input_ids from _forward_mlp.
- param_and_grad_buffer.py: replace inline fp8_params else-branch in
  finish_param_sync with self._post_param_sync() to match the three
  other call sites and dev's refactor (PR #4563).
- moe_layer.py: call _maybe_record_overload_factor() on dispatched_input
  before the routed-experts conditional, as on dev.
- token_dispatcher.py: pass
  num_sms_preprocessing_api=self.config.moe_hybridep_num_sms_preprocessing
  to _HybridEPManager constructor.
- datasets/readme.md: restore the dev-only "Packing Scheduler"
  documentation section before main's "Offline cache preparation".
- gpt_dynamic_inference_with_coordinator.py: restore
  `from megatron.training.arguments import parse_and_validate_args`
  import; the call at line 216 raised NameError on the
  gpt_dynamic_inference_tp2_pp2_dp2_583m_logitsmatch_zmq CI job.

Verified: ast.parse OK on all 5 changed .py files; ruff F821 reports
only pre-existing TYPE_CHECKING-guarded names; CODEOWNERS and the
pyproject.toml/uv.lock/Dockerfile.ci.dev triple unchanged from
origin/dev.
This was referenced May 12, 2026
Merged
Closed
@zhongbozhu zhongbozhu mentioned this pull request May 26, 2026
Merged
1 of 6 tasks

This branch was previously deployed

1 inactive deployment
test — da05c937 Deployed May 7, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #13032
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yaox12 yaox12 yaox12 approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@WanZzzzzz @svcnvidia-nemo-ci @yaox12 @ericharper