[feat] nccl_reshard refit: support generation-side pipeline parallelism (vLLM and Megatron generation)
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
`refit_transport: nccl_reshard` (the M-to-N shard-to-shard refit) rejects any generation engine with pipeline parallelism, for both vLLM and Megatron destinations. Training-side PP is supported; what is missing is generation-stage-aware destination routing. The legacy collective broadcast path has no such gate, so a PP>1 engine can only be refit on the slow path today. No recipe in the repo runs vLLM with `pipeline_parallel_size > 1`, so the gap is not exercised by CI.
## Where it is stated and enforced (upstream `main` at 5d49fbf4e)
- Design doc, "Generation-side PP > 1 is not supported by this refit transport yet": https://github.com/NVIDIA-NeMo/RL/blob/5d49fbf4e74bdd3a6627c48759c4a66b8920743b/docs/design-docs/nccl-reshard-refit.md#L97
- Config check, vLLM destination: https://github.com/NVIDIA-NeMo/RL/blob/5d49fbf4e74bdd3a6627c48759c4a66b8920743b/nemo_rl/weight_sync/nccl_reshard_utils.py#L846
- Config check, Megatron destination: https://github.com/NVIDIA-NeMo/RL/blob/5d49fbf4e74bdd3a6627c48759c4a66b8920743b/nemo_rl/weight_sync/nccl_reshard_utils.py#L772
- Why: the bulk-path comm groups are built per *training* PP stage and every generation worker joins all of them, because the destination map assumes each generation rank holds a TP shard of every layer: https://github.com/NVIDIA-NeMo/RL/blob/5d49fbf4e74bdd3a6627c48759c4a66b8920743b/nemo_rl/models/generation/vllm/vllm_backend.py#L522
- Related: vLLM PP>1 additionally requires the async engine (https://github.com/NVIDIA-NeMo/RL/blob/5d49fbf4e74bdd3a6627c48759c4a66b8920743b/nemo_rl/models/generation/vllm/vllm_generation.py#L173); #2180 was a PP>1 race on the collective path, not this gap.
## Reproduce
Any nccl_reshard recipe with `policy.generation.vllm_cfg.pipeline_parallel_size: 2` fails at setup with
```
nccl_reshard_refit cannot be enabled with the current config:
- policy.generation.vllm_cfg.pipeline_parallel_size must be 1 (got 2).
```
Before #3739 the same configuration was not rejected at setup.
## Why it matters
Models whose weights do not fit a node's TP degree need PP on the engine, and the fast refit is the only transport that keeps the weight sync at a few seconds at this scale. Two data points from GH200 nodes (4 GPUs each, Slingshot):
- Apertus 70B, training TP2 PP4 to vLLM TP4 PP1: 135.6 GiB per refit in about 3 s.
- GLM-5.1, training TP2 PP18 EP16 to vLLM TP32 EP32 PP1: 1,385.6 GiB per refit (97.5% on the bulk path) in 6.1 to 7.1 s, with the engine's tensor parallelism spanning eight nodes because PP on the engine is not available on this transport.
Today such models must choose between a PP>1 engine on the collective path or a cross-node TP layout.
## What would close it
- The destination map (`build_hf_to_local_param_map` on the vLLM and Megatron generation sides) describes only the parameters a generation rank's pipeline stage owns, and the reshard planner routes each training-stage shard to those ranks only.
- Refit value and logprob parity with a generation PP2 engine against training PP1 and PP4, plus clean failure and teardown, in the existing nccl_reshard tests.
We track the same gap on our fork as https://github.com/Alvorecer721/Nemo-RL/issues/30. We can test candidate changes on GH200 4-GPU-node layouts at whatever scale is useful.
Contributor guide
Assessment
This issue has not been assessed yet.