fix(mfsdp): refresh V2 compute weights after optimizer step

[wujingyue]

Summary

• Refresh MFSDP V2 BF16 compute weights through FullyShardedOptimizer's native post-step copy hook.
• Extract sync_model_weights_from_main_weights() so the checkpoint-load path and the fully_shard_optimizer() step post-hook refresh compute weights through one implementation instead of two.
• Extend the MFSDP V2 optimizer integration test from three to ten steps so missing compute-weight refreshes are detected.

Root cause

#5949 moved main-weight to model-weight synchronization from pre-forward to optimizer post-step. FullyShardedOptimizer inherited the post-step extension point but left its implementation as a no-op, so optimizer weights updated while subsequent forwards continued to use stale compute weights.

The existing three-step integration test used a 1% relative tolerance. The resulting loss drift remained below that threshold, so the test passed despite the missing refresh.

Also fixes a rank-order deadlock in fully_shard_optimizer()

The experimental step_post_hook collected parameter groups into a set and iterated it:

fsdp_parameter_groups: set[FsdpParameterGroup] = set()
...
for parameter_group in fsdp_parameter_groups:
    parameter_group.sync_model_weight_from_main_weight()

FsdpParameterGroup is @dataclass(frozen=True, eq=False), so it hashes by id() and set iteration order varies per process. sync_model_weight_from_main_weight() issues a collective whenever main_weight and model_weight are distinct buffers, so ranks could issue the per-group collectives in different orders and hang.

Routing the hook through the shared helper fixes this as a consequence of the refactor rather than as a special case: the helper iterates the parameters, whose order is identical on every rank, and uses a set only to skip groups already synced.

The bug is latent on main today because it needs both conditions at once — two or more parameter groups and main_weight is not model_weight. No in-tree test combines them: configurations whose optimizer and parameter placements match (e.g. HSDP) alias the two buffers, so the sync early-returns without a collective.

Validation

• uv run python -m torch.distributed.run --nproc_per_node=2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py::TestMcoreAdapter::test_build_train_and_step
• Full tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py on 4x H100: 32 passed. The one failure, test_overlaps_communication_and_compute[default], reproduces identically on unmodified main — it is a launch-bound profiler assertion, unrelated to this change.
• Deadlock repro: an HFSDP configuration (optimizer [Flat, Flat], parameter [Replicate, Flat]) over three parameter groups on 4 GPUs hung indefinitely before this change and completes in ~9s after it.
• ruff check and black --check on edited files

Related to #5949 and #6325.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[wujingyue]

/ok to test 1252bd6

[skyw]

Nit: single letter "s" difference is hard to spot and easily miss. May consider better name.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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