feat(megatron): add nccl_reshard M-to-N refit for Megatron generation

[wdykas]

What does this PR do?

This PR extends the existing nccl_reshard M-to-N weight-refit transport to support Megatron generation and consolidates Megatron refit selection under the common policy.generation.refit_transport setting.

For Megatron generation:

refit_transport	Behavior
null	NeMo-RL packed collective for non-colocated generation. This remains the default.
mcore	Megatron-Core native refit. The copy service is selected by mcore_generation_config.refit_backend.
nccl_reshard	NeMo-RL shard-to-shard M-to-N refit using xferdtensor, with Megatron as the destination.

Native MCore refit supports gloo, nccl, and nccl_m2n. The nccl_m2n backend is non-colocated only. nvshmem remains rejected while #3646 is unresolved.

Design

The nccl_reshard path retains its existing dual-path design:

• Bulk FFN weights are transferred directly from each training shard to the corresponding generation shard. No rank materializes the full bulk tensor.
• Embeddings, attention weights, normalization weights, routers, scales, and other miscellaneous tensors use the existing packed collective.
• Training pipeline stages receive dedicated bulk communicators so source PP layouts can be routed correctly.
• Megatron generation workers are explicitly marked as refit destinations while training workers retain source behavior; both share the same prepare_refit_info, build_hf_to_local_param_map, and refit entry points.

The destination declares its required payload representation through RefitPayloadMode:

• hf_export preserves the universal HF-named, backend-independent representation used by vLLM and other destinations.
• logical_weights is an explicit Megatron-to-Megatron exception. Megatron generation can assume a Megatron training source and request logical weights suitable for MCore destination storage.

This keeps the existing vLLM FP8 export behavior unchanged and prevents physical FP8 payloads from being mixed with incorrectly materialized BF16 weights.

Precision and parallelism

Megatron generation supports:

• BF16 training weights into BF16 generation.
• Supported Transformer Engine blockwise-FP8 or MXFP8 training parameters.
• BF16 or supported FP8 sources into MXFP8 generation storage.
• Training-side TP, PP, EP, and ETP layouts.
• Generation-side TP and EP, with independent source and destination placements.
• Generation ETP where supported; inference_optimized MoE generation is normalized to ETP=1 because MCore inference linears do not currently implement ETP greater than one.

For a Megatron destination, quantized sources are transferred as logical BF16. MXFP8 destinations receive into bounded, short-lived BF16 staging buffers and quantize into persistent MCore MXFP8 storage. The wire representation remains BF16 because Transformer Engine and MCore MXFP8 layouts are not byte-compatible.

Native MCore staging is configurable with:

policy:
  generation:
    refit_transport: mcore
    mcore_generation_config:
      refit_backend: nccl
      refit_execution_batch_bytes: null

null uses NeMo-RL's existing dynamic packed-buffer target: 2% of total GPU memory, capped at 5 GiB. A positive byte count overrides it.

Compatibility

• The top-level refit_transport default remains null.
• Shipped Megatron-generation recipes that depend on native MCore refit now select refit_transport: mcore explicitly, preserving their current behavior.
• refit_backend and refit_execution_batch_bytes are only used with refit_transport: mcore.
• Existing vLLM nccl_reshard payload behavior is preserved.
• Rank offsets allow non-colocated training and generation workers to join a shared process group without changing existing callers.

Current limitations

• nccl_reshard requires non-colocated Megatron training.
• Generation-side PP greater than one is not yet supported by nccl_reshard.
• Custom training PP layouts and VPP greater than one are not yet supported.
• SGLang, TRT-LLM, and ModelOpt real-quant destinations are not supported by this transport.
• MXFP8 refit uses BF16 on the wire and therefore does not provide MXFP8-sized transfer bandwidth.
• Native MCore refit has no worker-side abort hook, so a configured refit deadline is rejected rather than silently ignored.

Validation

• Added unit coverage for transport selection, source/destination behavior, process-group membership, native nccl_m2n, staging configuration, payload modes, MXFP8 handling, and TP/PP/EP/ETP geometry.
• Added grpo_megatron_nccl_reshard_refit.sh with BF16 coverage and a GB200 MXFP8 mode.
• Added GB200 nightlies for packed-collective MXFP8 refit and packed-collective plus FlashInfer MXFP8 refit on NanoV3 MoE.
• Wired BF16 and MXFP8 coverage into the corresponding functional test runners.
• 11 focused selector, synchronizer, ETP, and checkpoint tests passed.
• 320 reference/config validation cases passed.
• Ruff and Pyrefly pass on the changed Python files.
• A four-GPU Nemotron-3 Nano BF16-to-MXFP8 run completed two training/refit steps with numerical checks passing.

Before your PR is "Ready for review"

• Read and followed the contributor guidelines
• Added necessary unit and functional coverage
• Updated refit documentation and configuration examples
• Ran focused local validation
• Full PR CI

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

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: fc47b72 (PR #3739 from nemorl-native-megatron-inference-refit)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[youngeunkwon0405]

Thanks for adding Megatron generation support to the NCCL-reshard path. I have one architecture concern about how the backend extension point is represented.

The intended design separates backend-agnostic nccl_reshard_info from the per-backend hf_to_local_param_map. Each source/destination backend should implement build_hf_to_local_param_map(refit_info), own the resulting map, and use its LocalParamSpec entries to prepare and commit each refit payload. This is also the contract described by RefitBuilderInterface.

This implementation creates an equivalent hf_name -> LocalParamSpec map, but through a parallel generation-specific API and state: _build_megatron_bulk_refit_map, _generation_hf_to_local_param_map, prepare_nccl_reshard_generation_refit_info, nccl_reshard_generation_refit, and _generation_nccl_reshard_refit_info. The same pattern also appears in the collective path with prepare_generation_refit_info and update_generation_weights_from_collective. MegatronGeneration translates the public interface names into these worker-specific names, so the current call path works, but the common backend contract is no longer usable uniformly.

This appears to come from MegatronPolicyWorkerImpl serving both the train/source and generation/destination roles: it already owns the source-side methods with the standard names, so the generation mixin introduces parallel names to avoid collisions. As a result, a generation worker structurally exposes the source-side build_hf_to_local_param_map, while its actual destination builder is hidden behind a different private method. It also leads to source-side checks such as _uses_megatron_generation(), meaning that adding another destination backend or wire representation may require modifying the training worker again.

Could we preserve the existing extension point by separating the source and destination refit builders—either as backend-owned components or through another explicit role boundary—with each exposing the standard build_hf_to_local_param_map(refit_info) contract and owning its corresponding state? The payload representation could then be selected from the backend-agnostic per-parameter metadata and implemented through LocalParamSpec, keeping future backend additions localized to the new backend.

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 6b29a9e (PR #3739 from nemorl-native-megatron-inference-refit)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: cb35c4c (PR #3739 from nemorl-native-megatron-inference-refit)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[wdykas]

Thanks for adding Megatron generation support to the NCCL-reshard path. I have one architecture concern about how the backend extension point is represented.

The intended design separates backend-agnostic nccl_reshard_info from the per-backend hf_to_local_param_map. Each source/destination backend should implement build_hf_to_local_param_map(refit_info), own the resulting map, and use its LocalParamSpec entries to prepare and commit each refit payload. This is also the contract described by RefitBuilderInterface.

This implementation creates an equivalent hf_name -> LocalParamSpec map, but through a parallel generation-specific API and state: _build_megatron_bulk_refit_map, _generation_hf_to_local_param_map, prepare_nccl_reshard_generation_refit_info, nccl_reshard_generation_refit, and _generation_nccl_reshard_refit_info. The same pattern also appears in the collective path with prepare_generation_refit_info and update_generation_weights_from_collective. MegatronGeneration translates the public interface names into these worker-specific names, so the current call path works, but the common backend contract is no longer usable uniformly.

This appears to come from MegatronPolicyWorkerImpl serving both the train/source and generation/destination roles: it already owns the source-side methods with the standard names, so the generation mixin introduces parallel names to avoid collisions. As a result, a generation worker structurally exposes the source-side build_hf_to_local_param_map, while its actual destination builder is hidden behind a different private method. It also leads to source-side checks such as _uses_megatron_generation(), meaning that adding another destination backend or wire representation may require modifying the training worker again.

Could we preserve the existing extension point by separating the source and destination refit builders—either as backend-owned components or through another explicit role boundary—with each exposing the standard build_hf_to_local_param_map(refit_info) contract and owning its corresponding state? The payload representation could then be selected from the backend-agnostic per-parameter metadata and implemented through LocalParamSpec, keeping future backend additions localized to the new backend.

ok I think I am addressing this better now. I am just waiting on some other PRs going into bridge to get merged and then I can clean up and add more functional tests for review

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 351061d (PR #3739 from nemorl-native-megatron-inference-refit)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[wdykas]

/ok to test 351061d

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 65275a1 (PR #3739 from nemorl-native-megatron-inference-refit)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 28e22aa (PR #3739 from nemorl-native-megatron-inference-refit)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

[chtruong814]

@wdykas Sorry, I cancelled the previous CI run because it was hanging on this unit test job for 3 hours.
https://github.com/NVIDIA-NeMo/RL/actions/runs/34350677418/job/102506074692

Please resolve merge conflicts and we can try again.

[wdykas]

@wdykas Sorry, I cancelled the previous CI run because it was hanging on this unit test job for 3 hours. https://github.com/NVIDIA-NeMo/RL/actions/runs/34350677418/job/102506074692

Please resolve merge conflicts and we can try again.

ok will do now

[yuki-97]

Scope: re-review of the three fix commits since the last pass (407aca1) — ea842d7, b95991f, a8d2bcd.

Checked:

• Open threads — all 7 re-audited against the current head rather than the fix diffs; all resolved, each replied on individually.
• Config resolution — every --config-bearing script under tests/ run through load_config_with_inheritance + parse_hydra_overrides, and every shipped config resolving to Megatron generation walked for refit_transport.
• Docs — refit.md and the nccl-reshard design doc checked claim-by-claim against the code at head.

Remaining comments are all on the new commits: two config regressions introduced by the fixes, and one doc claim the new ETP guard contradicts.

[wdykas]

/ok to test 6a2b83d

[wdykas]

/ok to test 1475066

[wdykas]

/ok to test 66a132d

[wdykas]

/ok to test efff98b

[wdykas]

/ok to test 4194bf5

[wdykas]

/ok to test a168c70