fix(megatron): synchronize parameters before inference refit

[wdykas]

Summary

Complete deferred distributed-optimizer parameter all-gathers before non-colocated inference refit reads model weights.

With overlap_param_gather=True, updated compute-weight shards need not be assembled until the next training forward. Neither the collective HF export nor native Megatron reshard runs that forward. Consequently, refit can read partly stale parameters even though the transfer itself succeeds. The existing gather in prepare_for_generation covers the colocated path, not these source-worker exports.

This patch adds synchronization before collective broadcast and before refit buffer offload. The existing shared-buffer preparation is preserved, and device synchronization makes the gathered weights ready for side-stream exporters. Training forward hooks remain unchanged.

Scope: one file, 13 added lines. No configuration, dependency, submodule, parser, or kernel changes.

Validation

• Four-GPU BF16 precision-aware distributed Adam reproduction, TP2/DP2 with parameter-gather overlap: native-planner destinations differ before gather and become exactly equal afterward. The fixed rank-0 export source also contains stale values before gather.
• Three consecutive refit points: copy/gather/device-sync ordering verified; forward hooks unchanged; subsequent parameter updates bit-identical to the synchronous-gather control.
• Exact final inline blocks validated for DDP/non-DDP and overlap-enabled/disabled guards.
• Ruff lint and format checks pass.

These are isolated local checks; diagnostic scripts are intentionally not included in this minimal experiment commit. Full RL convergence and quantized shared-buffer behavior have not been validated.

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

[wdykas]

/ok to test 7ef3d21

[wdykas]

/ok to test 7ef3d21

[wdykas]

/ok to test e521a5b