[MXFP8]Update param buffer before AG in eval

[WanZzzzzz]

What does this PR do ?

Since we do a forced-sync param AG in eval, we need to copy the main param to param buffer before AG to ensure we gather the updated weights.

A related fix in Megatron-LM (NVIDIA/Megatron-LM#4563): Also in eval, after param AG, we need to copy the values in param buffer back to model weights (param.data) to update the weights. This is required for mxfp8+fp8-param-gather+reuse-grad-buff-for-mxfp8-param-ag. If we do not do this, the eval step will use the staled weights( those from last iteration).

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)

[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.

[Phlip79]

/ok to test 1f1353f

[yaoyu-33]

@WanZzzzzz feels like it could use some tests to make sure the behvior is correct

[gautham-kollu]

codecov is failing. Please add a test.

One example for a recipe test https://github.com/NVIDIA-NeMo/Megatron-Bridge/blob/main/tests/functional_tests/test_groups/recipes/test_deepseek_recipes_pretrain.py#L55-L56

[WanZzzzzz]

@yaoyu-33 @gautham-kollu
I added coverage in this MCore PR instead of a MBridge recipe test because the behavior being fixed is in MCore post-AG processing path. This fix in Mbridge just ensures the MCore code path will work correctly. Also I looked at the existing MBridge recipe tests, and they don’t currently expose the DDP knobs needed to exercise this bug, specifically overlap_param_gather, fp8_param_gather/fp4_param_gather, and reuse_grad_buf_for_mxfp8_param_ag in a train -> eval -> train flow.

The new MCore tests explicitly simulate train -> eval -> train with overlap_param_gather enabled and compare both train and eval loss against the non-param-gather reference:

MXFP8: tests/unit_tests/test_fp8_param.py::TestFP8Param::test_mxfp8_eval_transition
NVFP4: tests/unit_tests/test_fp4_param.py::TestFP4Param::test_nvfp4_eval_transition

[gautham-kollu]

/ok to test f71fe72

[cuichenx]

training code coverage is tested in mcore

[WanZzzzzz]

Test code is here:

NVIDIA/Megatron-LM#4563