[WanZzzzzz]

What does this PR do ?

When using fp8_param_gather=True + reuse_grad_buf_for_mxfp8_param_ag=True + overlap_param_gather=True with full-iteration CUDA graph (cuda_graph_impl=local, cuda_graph_scope=full_iteration), the first training iteration after warmup produces zero grad norm and incorrect loss.

Root cause: _handle_mxfp8_param_buffer_copy() guards the _copy_main_params_to_param_buffer() call behind a forward_pre_hook_enabled check. The main training loop intentionally disables forward pre-hooks for the first iteration (line 283-284 of train.py), so the copy is skipped. Without CUDA graphs this is safe — params are already in param.data and no all-gather runs.

However, with full-iteration CUDA graph, the all-gather (and the subsequent copy-to-param.data + zero of param_data) are baked into the captured graph and replay unconditionally. On the first iteration:

zero_grad_buffer() zeros param_data
_handle_mxfp8_param_buffer_copy() is skipped (hooks disabled)
Graph replays → AG gathers from param_data (all zeros) → copies zeros to param.data
Forward runs with zero weights → zero gradients

Changelog

• Add specific line by line info of high level changes in this PR.

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

Summary by CodeRabbit

Release Notes

• Bug Fixes
  • Improved parameter buffer handling during distributed training with CUDA graph acceleration, ensuring correct execution across all training iteration scenarios.

[copy-pr-bot]

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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 577ffee6-ed57-4322-b13d-8fec8dc04e49

📥 Commits

Reviewing files that changed from the base of the PR and between 5addaf6 and b5e8c46.

📒 Files selected for processing (1)

• src/megatron/bridge/training/train.py

---

📝 Walkthrough

Walkthrough

The change modifies _handle_mxfp8_param_buffer_copy to invoke parameter buffer copying not only when forward pre-hooks are enabled, but also when a full-iteration CUDA graph for training has already been captured. Documentation is added explaining the requirement for correct parameter data during CUDA graph replay.

Changes

Cohort / File(s)	Summary
MXFP8 Parameter Buffer Handling src/megatron/bridge/training/train.py	Modified condition in _handle_mxfp8_param_buffer_copy to trigger DistributedOptimizer._copy_main_params_to_param_buffer() when full-iteration training CUDA graphs are captured, with documentation explaining that all-gather and param_data zeroing execute unconditionally during graph replay.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Test Results For Major Changes	⚠️ Warning	PR fixes convergence bug but lacks test results or validation evidence demonstrating the fix works and no regressions occur.	Add test results and convergence validation evidence to PR description showing the zero gradient/loss issue is resolved.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically summarizes the main change: fixing a convergence issue that occurs when combining mxfp8 parameter handling with full CUDA graph execution.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[yaoyu-33]

/claude review

[claude]

LGTM

[yaoyu-33]

/ok to test b5e8c46