fix(megatron): synchronize parameters before inference refit

[wdykas]

Summary

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

With overlap_param_gather=True, updated compute-weight shards are not assembled until the next training forward. Refit runs between the optimizer step and that forward, so every DP rank exports its own updated shard next to stale copies of the others. This affects every refit path that reads Megatron parameters directly: the collective HF export, the native Megatron reshard, and the colocated CUDA-IPC stream.

Changes:

• MegatronPolicyWorkerImpl.sync_params_before_refit gains a plain-BF16 branch: with a DistributedDataParallel model and overlap_param_gather, finalize any in-flight async checkpoint, force a synchronous parameter all-gather, and device-sync. The existing MXFP8 shared-buffer path is unchanged.
• refit_policy_generation (legacy callers without a WeightSynchronizer) now calls sync_params_before_refit() before both the colocated and non-colocated branches, in the same order as IPCWeightSynchronizer.sync_weights (sync, offload, prepare). Colocated vLLM attaches no synchronizer and previously took this branch with no gather.
• Base-worker and worker comments updated; Megatron-FSDP wrappers are intentionally not handled (they gather in their own module hooks).

Behavior note: non-colocated Megatron refits with overlap_param_gather=True now block on in-flight async checkpoint writes before the gather, since the gather rewrites the parameter buffer the writer may still reference.

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: gather/device-sync ordering verified; forward hooks unchanged; subsequent parameter updates bit-identical to the synchronous-gather control.
• Unit tests: legacy colocated and non-colocated refit ordering (tests/unit/algorithms/test_grpo.py); BF16 branch event sequence and no-op guards (tests/unit/models/policy/test_megatron_worker.py).
• Ruff lint and format checks pass.

These are isolated local checks. 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

[wdykas]

/ok to test 50864a2

[wdykas]

Two findings: the colocated legacy refit branch in grpo.py still skips the new param sync (default colocated vLLM path), and the _copy_main_params_to_param_buffer call in the new BF16 branch is a guaranteed no-op that the new unit test nonetheless asserts. Ruff passes; the grpo unit test passes locally. Worker unit tests could not be run here (aarch64 box lacks causal_conv1d_cuda).

[wdykas]

Second pass at 538e0ce, after the fixes for the two earlier findings.

Both prior findings are addressed: sync_params_before_refit() now runs before both legacy branches (sync, offload, prepare, matching IPCWeightSynchronizer), and the MXFP8-only staging call is gone from the BF16 branch with the test asserting the real sequence plus assert_not_called on the staging helper.

Re-verified against Megatron-LM source: start_param_sync(force_sync=True) runs a synchronous coalesced all-gather over every bucket group (including expert-parallel), leaves each with param_gather_dispatched=True, param_gather_handle=None, so the next forward pre-hook is a no-op on already-current weights; param_gather_dispatched is reset in finish_grad_sync at the end of the next backward, so the post-step gather still fires. The optimizer step writes updated shards into bucket.param_data, which the gather reads, so no staging is needed. finalize_async_save is idempotent across the sync -> offload double call. No caller of refit_policy_generation already calls sync_params_before_refit itself, and DTensor policies inherit the base no-op.

No new findings at high confidence. Two notes for the record:

• Behavior change: non-colocated Megatron refits with overlap_param_gather=True now block on in-flight async checkpoint writes (previously only the colocated offload_before_refit did). Required for correctness; noted in the description.
• Local verification: the two grpo unit tests pass here. The worker unit tests could not run on this box (missing causal_conv1d_cuda extension and a stale Megatron-LM submodule checkout), so they rely on CI.

LGTM once CI is green.

[wdykas]

/ok to test 538e0ce

[yuki-97]

Scope: the Megatron pre-refit parameter sync, both refit entry points that reach it (refit_policy_generation and the SingleController), and the two unit-test surfaces the PR adds.

CI at 538e0ce is red, and the failure is in this PR's new code. Four functional-test jobs fail — L1_Functional_Tests_Megatron_4 and L1_Functional_Tests_SingleController, on both runners (one job). All four raise the same error on the initial refit, inside the new BF16 branch:

megatron_policy_worker.py:4161 in sync_params_before_refit
    self.model.start_param_sync(force_sync=True)
  -> param_and_grad_buffer.py:557  shard_buffer(bucket.param_data, ...)
RuntimeError: setStorage: sizes [40907008], strides [1], storage offset 0, and itemsize 2
              requiring a storage size of 81814016 are out of bounds for storage of size 0

param_data has been resized to 0, i.e. the policy params are offloaded at that point. The fifth red check is the CI-label gate and is unrelated.

Also checked:

• Upstream semantics — start_param_sync(force_sync=True) traced in the pinned Megatron-LM (1e7598cbf) for the in-flight handle wait, _post_param_sync, and the param_gather_dispatched lifecycle across the following train step.
• Entry-point coverage — every sync_weights call site in the repo and every synchronizer under nemo_rl/weight_sync/, for whether the sync runs at all and where it sits relative to offload and KV-scale calibration.
• Reachability — shipped Megatron recipe config chains resolved down to use_distributed_optimizer and overlap_param_gather.
• Branch coverage — the (MXFP8 x hook-state) matrix against the three sync_params_before_refit tests.
• Doc surfaces — docs/ and the weight_sync lifecycle docstrings re-read against the widened semantics; no contradictions found.

Remaining comments are on two transports the sync does not reach today and where the call belongs so it cannot be missed again, plus two nits.

[wdykas]

/ok to test 2b890b3

[wdykas]

/ok to test e8ef8de

[wdykas]

/ok to test e8ef8de

[yuki-97]

Scope: third pass, over e8ef8de — its two test-double fixes, the parameter-buffer residency question the previous pass left open, and the MXFP8 grouped-GEMM change.

Checked:

• The test doubles updated for the new sync_params_before_refit call — test_grpo_checkpoint_engine.py and test_checkpointing.py. Both are required by the call sites at grpo.py:2548 and single_controller.py:4070, and both look right. All 103 checks are green at head.
• Parameter-buffer residency across the forced gather — every refit_policy_generation site in ppo.py and grpo.py traced against the guard at megatron_policy_worker.py:4118-4124. Each one either onloads first (prepare_for_training / prepare_for_lp_inference) or runs before the first train step with the hooks dropped at :765-766, so start_param_sync(force_sync=True) cannot observe a released buffer. Nothing to raise there.
• The MXFP8 grouped-GEMM change against the Megatron-LM pin, which this PR leaves at 1e7598cbf.

One comment, on the MXFP8 backend set.

[wdykas]

/ok to test 16eb1b0