[Training][Rollout] Add direct NCCL delta weight updates - #392
[Training][Rollout] Add direct NCCL delta weight updates#392ShuoleiWang wants to merge 1 commit into
Conversation
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements a direct delta-weight-update (DWU) MVP for non-colocated Megatron trainers and VIME-launched vLLM rollout engines, utilizing direct NCCL transport instead of the previously planned disk-based transport. It introduces a CheckpointDeltaSource to compute and send sparse weight updates, a VimeDeltaNCCLWeightTransferEngine to receive and apply these patches on the vLLM side, and a WeightUpdateCoordinator to manage the control plane sequence across engines. The review comments correctly identify a critical runtime issue in vime/backends/vllm_utils/checkpoint_delta.py where PyTorch Stream objects are passed to PyNcclCommunicator.broadcast instead of raw CUDA stream pointers, which will cause a TypeError on GPU.
| self.model_update_group.broadcast( | ||
| positions, | ||
| src=0, | ||
| stream=torch.cuda.current_stream(), | ||
| ) |
There was a problem hiding this comment.
In vLLM, the PyNcclCommunicator.broadcast method expects a raw CUDA stream pointer (an integer) rather than a PyTorch Stream object. Passing torch.cuda.current_stream() directly will result in a TypeError at runtime on GPU.
Use torch.cuda.current_stream().cuda_stream to pass the raw stream pointer.
| self.model_update_group.broadcast( | |
| positions, | |
| src=0, | |
| stream=torch.cuda.current_stream(), | |
| ) | |
| self.model_update_group.broadcast( | |
| positions, | |
| src=0, | |
| stream=torch.cuda.current_stream().cuda_stream, | |
| ) |
| self.model_update_group.broadcast( | ||
| values, | ||
| src=0, | ||
| stream=torch.cuda.current_stream(), | ||
| ) |
There was a problem hiding this comment.
In vLLM, the PyNcclCommunicator.broadcast method expects a raw CUDA stream pointer (an integer) rather than a PyTorch Stream object. Passing torch.cuda.current_stream() directly will result in a TypeError at runtime on GPU.
Use torch.cuda.current_stream().cuda_stream to pass the raw stream pointer.
| self.model_update_group.broadcast( | |
| values, | |
| src=0, | |
| stream=torch.cuda.current_stream(), | |
| ) | |
| self.model_update_group.broadcast( | |
| values, | |
| src=0, | |
| stream=torch.cuda.current_stream().cuda_stream, | |
| ) |
1 participant
Summary
This PR adds an initial direct delta-weight-update (DWU) path from a non-colocated Megatron trainer to VIME-launched vLLM rollout engines.
Instead of broadcasting the full actor checkpoint after every training step, the new path:
int32checkpoint indices for later sparse updates;model.load_weights()mappings; andThe PR also consolidates the control plane used by the existing full-weight NCCL and CUDA IPC paths into a shared
WeightUpdateCoordinator, and adds argument validation, English and Chinese documentation, and protocol/lifecycle tests.No vLLM source files are modified by this PR.
Motivation
Weight synchronization is on the repeated critical path in disaggregated RL training. A full sync always transfers the complete model, even when only a small fraction of BF16 values changed after the optimizer step.
In the measured Qwen3-30B-A3B run below, a sparse step changed 4.50% of the checkpoint elements. The direct delta path produced an 8.24 GB logical NCCL tensor payload instead of the 61.06 GB dense BF16 payload, a reduction of 86.5%. A no-change step produced no tensor payload.
The implementation deliberately keeps checkpoint-to-runtime mapping inside vLLM. VIME produces patches in Hugging Face checkpoint coordinates, while the model's native loader remains responsible for tensor-parallel slicing, packed QKV and gate/up projections, and model-specific name mapping.
Design
Source-side encoding
CheckpointDeltaSourcestores one committed CPU snapshot for each deterministic HF export bucket.For every update it:
An update with no changed values sends only the final manifest and advances the logical version without moving tensor data.
Receiver-side application
The VIME-owned vLLM weight-transfer backend validates the session's base and target versions, sequence order, encoding, and final-manifest lifecycle. For data chunks, it drains the NCCL payload before consulting worker-local version/sequence state so a stale worker does not strand healthy peers in the same collective.
Received metadata and tensors are converted into
CheckpointWeightPatchobjects. Dense seeds use vLLM's layerwise reload lifecycle. Sparse patches modify initialized runtime storage without rerunning dense post-load processing.Patch payloads contain absolute values rather than arithmetic deltas. The versioned session protocol still rejects replay, mixed versions, and out-of-order chunks; a failed session is not retried in place.
Shared update lifecycle
WeightUpdateCoordinatorprovides one control plane for the existing full paths and the new delta path:Intermediate chunks carry a candidate version but do not publish it. The trainer-side version and rollout-engine version marker advance only after the complete transfer and source commit succeed. Rank-zero control-plane failures are broadcast to the other trainer ranks; transport callbacks remain collective operations and are required to surface failures collectively.
GPU writes are not transactionally rollback-safe. Recovery is fail-stop: do not serve or retry an affected session in place; restart the task/affected rollout fleet and establish a new dense seed.
Effect on existing full-weight modes
This PR also changes the lifecycle implementation used by existing full-weight NCCL and CUDA IPC synchronization:
set_weight_versionpublishes the server-side tag when available. For non-delta modes, an older vLLM that returns 404/405 falls back to the local marker and logs one warning.The non-colocated full NCCL path was exercised by the H20 baseline below. Colocated CUDA IPC behavior is covered by the CPU adapter/lifecycle tests but was not included in this H20 A/B.
Scope
Validated in this PR:
Rejected by argument validation:
Not claimed by the current validation evidence: multimodal models, rollout-side expert parallelism/EP filtering, and multi-model deployments with multiple independently configured rollout models.
The source remains wire-sparse rather than end-to-end sparse: every sync still performs a full canonical-HF export and compares against a CPU snapshot. The receiver may also construct full checkpoint-shaped staging tensors while applying a sparse patch.
vLLM dependency
This PR requires a vLLM build containing the checkpoint-coordinate patch API introduced by:
vllm-project/vllm#50723
The H20 testbed used vLLM 0.26. Its
checkpoint_weight_patch.pyimplementation was verified byte-identical to the file at vLLM commit:The patch and layerwise-reload APIs are checked when the Direct-DWU engine is constructed. Default full-weight mode does not import the delta receiver and does not require the checkpoint patch API, although it does use the shared coordinator introduced by this PR.
Direct-DWU also requires the vLLM server's
/update_weight_versionendpoint to publish the committed version. Unlike the patch and reload helpers, this capability is exercised at the final publish step; a missing endpoint is fatal for delta mode. Non-delta modes retain the 404/405 fallback described above.H20 validation
The A/B used Qwen3-30B-A3B in BF16 on two nodes:
eth0, with IB disabled.The timings below use the outer
Timer update_weightswall time. Payload sizes are logical tensor payloads, not NIC counters; they exclude HTTP metadata and NCCL protocol overhead.The easiest way to read the result is that the full path always moves the whole model, while the delta path still inspects the whole model but moves only the changed values.
First sync (one-time setup)
Later sync with no BF16 weight change
Later sync with 4.50% of BF16 elements changed
The main result is therefore straightforward: for the measured sparse update, the tensor payload fell by 86.5% (61.06 GB to 8.24 GB), but end-to-end sync time only fell by about 5% (20.47 s to 19.4 s). The producer still performs a full HF export and diff before sending the sparse payload. The 8.6 s no-change result also shows substantial fixed work even when tensor payload is zero; this experiment did not profile the individual components separately.
Full and delta were separate runs, not paired measurements of the same optimizer step. Their training trajectories diverged because vLLM sampling is not deterministic across independent runs, so training-curve agreement is not used as a correctness claim.
Correctness was checked separately with the real checkpoint patch implementation on vLLM TP4. Sparse and native dense loading produced bit-identical runtime tensors for 14 representative checkpoint tensors spanning fused MoE experts, router weights, packed QKV, embeddings, norms, and
lm_head; reverting the perturbations also restored the two greedy test outputs exactly.The delta run successfully committed the dense seed, no-change update, and nonzero sparse update and completed all 3/3 update-and-rollout cycles without protocol errors. It later failed in the unrelated final asynchronous actor-checkpoint save. The full baseline disabled final checkpoint saving outside the measured update boundary and completed successfully. Testbed-only changes also relaxed the local NumPy-version guard and pinned NCCL/Gloo to
eth0; they did not modify the weight-sync implementation.CPU validation
Focused suites:
pytest tests/utils/test_direct_dwu_roundtrip.py \ tests/utils/test_weight_update_coordinator.py \ tests/utils/test_update_weight_from_distributed.py \ tests/utils/test_update_weight_from_tensor.py \ tests/utils/test_vllm_engine.py \ tests/test_megatron_argument_validation.pyThe main round-trip suite feeds the real source's
update_info()andwire_tensors()into the real receiver against a fake model. It covers dense seed, nonzero sparse update, no-change update, source abort, failed-session poisoning, version/sequence errors, layout drift, NaN rejection, and final manifest behavior.Latest full
tests/utilscomparison in the same local environment:The six failures were not branch-specific; no additional failure was introduced relative to
origin/mainin that environment.Usage
The trainer must start from update version 0, and VIME must launch and own the non-colocated vLLM rollout engines. See
examples/delta_weight_sync/README.mdfor the runtime requirements and verification checklist.Follow-up
The current producer gathers and exports the full HF checkpoint before computing the delta. A follow-up shard-local producer can diff Megatron shards in parallel and gather only changed positions and values, reducing the producer-side contribution to the fixed overhead visible in the no-change run.
AI assistance disclosure
This PR includes code developed with assistance from Claude and Codex. The human submitter reviewed the changes, ran the validation described above, and remains responsible for this PR.