fix(sc): pass refit_payload_mode when recovery regenerates refit info

[asolergi-nv]

What does this PR do ?

Fixes the fleet-health recovery on the collective (nccl) refit transport: reconcile_communicator() called Policy.prepare_refit_info() without the refit_payload_mode that #3739 made mandatory, so every recovery died with a TypeError before touching NCCL.

Why it happened

• feat(megatron): add nccl_reshard M-to-N refit for Megatron generation #3739 (5d49fbf4e, Sept 11) made refit_payload_mode a required keyword-only argument of Policy.prepare_refit_info and updated the one caller that existed at the time, CollectiveWeightSynchronizer.init_communicator (collective_weight_synchronizer.py:213).
• feat(sc): restart dead generation shards and re-admit them at the next refit #3592 (53bce0568, Sept 14) added a second caller in CollectiveWeightSynchronizer.reconcile_communicator — the recovery that rebuilds the communicator over the surviving shards and redistributes refit metadata to them (a restarted engine has none). It was written against the pre-feat(megatron): add nccl_reshard M-to-N refit for Megatron generation #3739 signature, and the merge was textually clean and semantically wrong.
• Nothing turned red. The unit fake for the policy in test_reconcile_communicator.py accepted a bare call; every fleet-health recovery variant in L1_Functional_Tests_SingleController.sh self-skips on CI's 2-GPU runners; and the last multi-GPU runs of those variants predate the main sync that brought feat(megatron): add nccl_reshard M-to-N refit for Megatron generation #3739 in.

It surfaced in the fault-injection validation of #3929 — see #3929 (comment) (attribution to #3739/#3592) and the run details in #3929 (comment) (CW-DFW job 18689836: step 3/24, one generation shard killed, then):

CollectiveWeightSynchronizer.reconcile_communicator()
  state_dict_info = self._policy.prepare_refit_info()
TypeError: Policy.prepare_refit_info() missing 1 required keyword-only argument: 'refit_payload_mode'

Effect on main since Sept 14: on the default transport, a run that loses a generation shard ends in the rebuild that was meant to keep it alive (#3591's contract). nccl_reshard is unaffected — its rebuild does not re-prepare refit info — which is why the reshard step-kill scenario in that validation passed while the nccl one and the Gym quarantine one failed.

The fix

• nemo_rl/weight_sync/collective_weight_synchronizer.py: the rebuild now mirrors init_communicator — self._policy.prepare_refit_info(refit_payload_mode=self._generation.get_refit_payload_mode()). get_refit_payload_mode is on GenerationInterface (default "hf_export", overridden by MegatronGeneration), so this is valid for every backend the collective synchronizer can hold.
• tests/unit/weight_sync/test_reconcile_communicator.py: the fake policy is now keyword-only, exactly like Policy.prepare_refit_info, so the signature is pinned — against the old call, 14 of the 38 tests in the file fail with the production error. The fake generation answers get_refit_payload_mode, and a new test_refit_info_is_regenerated_in_the_backend_s_payload_mode asserts the rebuild asked for the backend's mode and pushed the metadata to the fleet once.

Issues

No standalone issue; raised in #3929 (comment). #3929 should re-merge main after this lands and re-run its collective-transport scenarios.

Usage

Not applicable — no user-facing change. To exercise the fixed path end to end on a node with ≥ 3 GPUs:

# default transport (nccl): kills one generation shard at a step boundary and expects the run to
# complete on the survivor. Before this fix it dies at the first rebuild with the TypeError above.
bash tests/functional/grpo_sc_generation_shard_recovery.sh

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 run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests — unit: 1420 passed across tests/unit/weight_sync, tests/unit/single_controller, tests/unit/models/generation/test_fleet_health.py, tests/unit/models/policy/test_worker_refit_signatures.py, tests/unit/models/policy/test_lm_policy_collective.py. Functional: the recovery variant above needs ≥ 3 GPUs and was not run locally; it is the run that found the bug.
• Did you add or update any necessary documentation? — none needed; behaviour is what the design docs already describe.

Additional Information

• The nccl_reshard mid-refit stall reported in the same validation (reshard-refit-kill / reshard-refit-freeze) is a different problem — a timing race between the controller's refit bound and the synchronizer's straggler settle — and is not addressed here.

🤖 Generated with Claude Code

[copy-pr-bot]

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

Contributors can view more details about this message here.

[asolergi-nv]

/ok to test 3d4701f

[terrykong]

@youngeunkwon0405 to review (this pr should target yours)

[youngeunkwon0405]

@youngeunkwon0405 to review (this pr should target yours)

Hi @terrykong, why does this PR have to target #3929? This is just fixing the previous implementation. Shouldn't it be more natural to go right into main?

[terrykong]

@youngeunkwon0405 sorry, it was my confusion. i thought it was tied