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

Assistant: GPT-5.5 xhigh

What does this PR do ?

Follow up of #4818 #4563

This PR fixes a general correctness issue in the interaction between:

--fp8-param-gather
--reuse-grad-buf-for-mxfp8-param-ag
--overlap-param-gather

With MXFP8 grad-buffer reuse, DDP param_data is not a persistent parameter buffer. It is a temporary high-precision all-gather staging buffer shared with the gradient buffer. After param all-gather completes, the gathered values are copied/quantized into TE MXFP8 parameter storage, and the shared buffer is zeroed so the next backward pass can safely accumulate gradients into it.

That means any explicit/forced param sync must first ask the optimizer to restage the updated FP32 master-weight shards into the DDP param buffer. Otherwise DDP may all-gather stale data or zeroes and overwrite model parameters.

Previously checkpoint path did roughly:

  if should_disable_forward_pre_hook(args):
      force_param_sync(model)

  save_checkpoint(...)

Then force_param_sync(model) called:

  model_chunk.start_param_sync(force_sync=True)

DDP does not all-gather from optimizer master weights. It all-gathers from the local shard view of bucket.param_data

So if nobody first did _copy_main_params_to_param_buffer(), then the checkpoint forced AG reads whatever is currently in the shared DDP staging buffer.

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 @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"

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.