[checkpoint_engine][rollout] Delta weight sync over NCCL for disaggregated rollout (+ sharded-snapshot variant)#6974
Open
ChangyiYang wants to merge 11 commits into
Open
Conversation
…saggregated rollout Adds a "delta" checkpoint engine that puts only the changed weights on the trainer->rollout wire, mirroring THUDM/slime's NCCL delta transport. Instead of broadcasting every parameter each sync, the trainer byte-diffs against a pinned-CPU snapshot and broadcasts only the changed (position, value) pairs over the same ray.util.collective group the full-weight NCCLCheckpointEngine uses. Design (follows verl's existing "A" topology: actor rank0 -> rollout workers): - verl/workers/rollout/delta_sync/: framework-agnostic core -- DeltaState (pinned snapshot + bytewise diff, side-stream H2D/D2H pipelining), encode/ decode (indices / gap-deltas, per-param manifest, checksum), and a wrapper that turns verl's (name, tensor) generator into bucketed DeltaFlush objects. - DeltaCheckpointEngine(NCCLCheckpointEngine): send_weights diffs + broadcasts per-flush positions/values (master uses cupy buffers, expandable_segments safe); the rollout worker reconstructs full tensors from the delta into a local mirror and hands them to the standard server_adapter.update_weights. So the trainer->worker wire is sparse, while the worker->engine push is an ordinary full-tensor load -- no SGLang-side delta receiver required. First sync forces a full delta so a dummy-initialized rollout gets a correct base; a per-flush checksum is verified on receipt. Enable with rollout.checkpoint_engine.backend=delta (+ engine_kwargs.delta. encoding). Validated on a 4+4 single-node disaggregated one_step_off GRPO run; unit tests cover encode/decode bit-identity. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Adds test_delta_result_equals_full_sync: seeds the trainer snapshot and the rollout mirror from the same W0, then over several steps diffs W_new, applies only the changed positions onto the mirror (reproducing the rollout worker's receive_weights mirror-combine), and asserts the mirror is byte-equal to W_new -- i.e. the weights a rollout ends up with via delta == what the old full path delivers. CPU-only, no GPU/NCCL/SGLang: the transport only moves bytes, so the lossless guarantee lives entirely in encode/decode/combine. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ngine The delta_sync core (DeltaState / encode / wrapper) was placed under workers/rollout/ back when delta lived in the SGLang rollout ServerAdapter. Its only consumer now is DeltaCheckpointEngine, so move it to verl/checkpoint_engine/delta_sync/ (and the test to tests/checkpoint_engine/) and switch to relative imports. Fixes the awkward checkpoint_engine -> workers.rollout dependency direction; no behavior change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The delta transport is NCCL-only; remove the leftover disk framing. Drops the `deltas_zstd` encoding (it only existed to zstd-wrap the gap stream at safetensors-write time on the disk path -- a no-op alias for `deltas` without disk) and rewrites the docstrings that still referenced disk safetensors / a DeltaSpec-style SGLang receiver. The receiver is now the rollout worker's local decode-into-mirror; no behavior change for the `indices` / `deltas` encodings. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Document the ``delta`` checkpoint-engine backend in the one_step_off (disaggregated) recipe (docs/advance/one_step_off.md) and note it in the checkpoint-engine worker docs (docs/workers/engine_workers.rst), and add a runnable example: grpo_0.6b_gsm8k_fsdp2_sglang_delta_2_6.sh — the SGLang 2+6 disaggregated GRPO recipe with checkpoint_engine.backend=delta. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The `delta` backend still `full_tensor()`-gathers every parameter to rank 0 before diffing, and rank 0 keeps a full-model pinned-CPU snapshot. This adds a `delta_sharded` backend that pushes the diff below the all-gather: each actor rank pins a snapshot of only its FSDP shard, byte-diffs the shard locally, and gathers just the changed (position, value) pairs to rank 0. The gather volume drops from the full parameter to the sparsity ratio (~1-3%) and no rank needs a full-model snapshot -- memory and gather traffic both shard with the world size. - delta_sync/sharded.py: local_shard_view computes each shard's absolute offset in the full flattened parameter purely locally from the DTensor spec (compute_local_shape_and_global_offset, no collective); handles uneven shards, and 2D FSDP x Replicate meshes (ulysses SP) by only letting the replicate coord-0 rank contribute, so replicated params are not double-counted. shard_delta_indices does the bytewise (view-as-int) diff; gather_v_to_rank0 does a count-exchange + pad-to-max gather to rank 0. - DeltaShardedCheckpointEngine (backend="delta_sharded"): assembles rank 0's gathered deltas into the SAME indices-encoded flush + zmq manifest + cupy broadcast + per-flush checksum as `delta`. The assembled delta is bit-identical to the parent's, so the rollout-side receiver is reused unchanged. - FSDP get_per_tensor_param_shard yields the per-rank local shards; engine_workers dispatches delta_sharded to it instead of the full-tensor generator. Scope: FSDP2 Shard(0) params + replicated / non-DTensor params; other shard dims raise NotImplementedError.
test_sharded_delta.py: CPU unit tests that the local bytewise shard diff (shard_delta_indices) reproduces the global bytewise diff and is empty when the shard is unchanged, plus local_shard_view on a plain (non-DTensor) param. sharded_delta_multigpu_check.py: a torchrun check (nproc_per_node=4) that runs the real sharded module against a full_tensor()-gather baseline over uneven Shard(0) shapes on a 1D FSDP mesh and a 2D FSDP x SP (replicate) mesh, and asserts the gathered (position, value) set is bit-identical to the full diff.
Add a "Sharded snapshot (delta_sharded)" subsection under Delta Weight Sync in one_step_off.md explaining the per-shard snapshot + gather-only-the-changes design, that it is bit-identical to delta, and its scope.
ChangyiYang
requested review from
ArronHZG,
eric-haibin-lin,
vermouth1992 and
wuxibin89
as code owners
|
|
Contributor
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
ChangyiYang
changed the title
wuxibin89
reviewed
Jul 8, 2026
| > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node > physical_gpus_per_node`, | ||
| > the required node count is `trainer.nnodes + rollout.nnodes` | ||
|
|
||
| ### Delta Weight Sync |
Collaborator
There was a problem hiding this comment.
Move to a separate delta weight sync design doc: megatron, fp8, etc should be include in future.
| n_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout)) | ||
|
|
||
|
|
||
| python3 -m verl.experimental.one_step_off_policy.main_ppo \ |
Collaborator
There was a problem hiding this comment.
Please test on new V1 trainer:
trainer.v1.trainer_mode=separate_asyncverl/experimental/one_step_off_policy and verl/experimental/fully_async_policy are going to be deprecated in future release.
| # The sharded delta engine diffs each rank's local FSDP shard (no all-gather), | ||
| # so it consumes the sharded param generator instead of the full-tensor one. | ||
| if effective_mode == "delta_sharded": | ||
| per_tensor_param, _ = self.actor.engine.get_per_tensor_param_shard() |
Collaborator
There was a problem hiding this comment.
Add a common function get_per_tensor_param_shard to BaseEngine.
…or_param_shard on BaseEngine - Move the delta weight sync section out of one_step_off.md into a standalone docs/advance/delta_weight_sync.md design doc with a roadmap (Megatron, fp8, native in-engine apply), per review. - Declare get_per_tensor_param_shard on BaseEngine so the delta_sharded consumption in engine_workers is a formal engine interface; FSDP implements it (Megatron implementation comes with the mcore backend). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…SGLang's custom weight loader Two structural upgrades to the delta backends, validated end-to-end on Qwen2.5-7B across 2 nodes (trainer node -> rollout node): Sender: stream the bucket flushes. Instead of materializing every flush before dispatch, each bucket-sized flush is broadcast the moment it is produced and freed (is_last-terminated stream), so sender peak memory is ~2 buckets regardless of model size -- previously the first full-seed sync held the entire delta on rank 0 and OOMed at 7B. Oversized per-param deltas (e.g. the embedding on the full seed) are sliced into <=64M-element manifest entries to bound the receiver's decode transient. Receiver: drop the full-model mirror; apply in place inside SGLang. The rollout CheckpointEngineWorker no longer reconstructs full tensors (which staged a second full model on the rollout GPU). It hands its local copy of the sparse payload to the colocated SGLang TP worker over same-GPU update_weights_from_tensor IPC, where a verl-shipped loader -- registered automatically through SGLang's stock --custom-weight-loader hook, no SGLang fork or patch required -- verifies the flush checksum, densifies each param's delta into a NaN-masked tensor (int32-view position decode, 8 B/elem transient), and overwrites only the changed positions in place on the live weights via a masked-copy load. The radix cache is flushed once per sync (on the stream's last flush) rather than per bucket. Receiver peak memory is now one bucket plus one decode chunk, independent of model size: 7B runs at gpu_memory_utilization=0.7 with zero OOM (the mirror receiver required 0.4), and steady-state sync drops ~2x (7s -> 3.5s; full-NCCL baseline 2.75s on the same 2-node setup). Rewards/KL match the NCCL baseline; checksums verified on every flush. Also declares get_per_tensor_param_shard usage in docs and adds loader round-trip unit tests (bit-identity for both encodings, masked-apply isolation, checksum fail-loud). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…embly Profiling the delta send at 7B showed ~70% of steady-state send time (and 83s of the one-time full seed) spent in _assemble_flush's host round-trip: positions went GPU -> .cpu().numpy().tobytes() -> bytes join -> frombuffer, only for _publish_flush to move them straight back to the GPU for the NCCL broadcast. Assemble now concatenates the int32 positions on the GPU and bitcasts to the uint8 wire view directly; bytes and checksum are unchanged. 7B 2-node effect (same config as before): sender 3.2-3.8s -> ~0.9s (diff 0.17s + gather 0.47s + pack 0.01s + broadcast 0.09s); full seed send 97s -> 10s; per-step sync_rollout_weights 3.5-4.0s -> 1.8s, now ~35% faster than the full-NCCL baseline (2.75s) on the same cluster. The sparse gather itself is 3x cheaper than the full all-gather (0.47s vs 1.4s), so the delta backend now wins end to end even on a fast interconnect. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
Author
Benchmark: delta vs full-NCCL weight sync — Qwen2.5-7B, 2×8 GPU disaggregated (cross-node)Setup: one-step-off disaggregated GRPO on GSM8K — 8-GPU FSDP2 trainer (node A) + 8-GPU SGLang rollout,
|
| step | delta_sharded |
nccl (full) |
|---|---|---|
| 1 (Δ≈0) | 1.26 | 3.14 |
| 2 | 1.61 | 2.52 |
| 3 | 1.65 | 2.86 |
| 4 | 1.73 | 2.60 |
| 5 | 1.73 | 2.65 |
| 6 | ~1.7 | 2.63 |
| steady-state mean | ~1.67 s | ~2.65 s |
~35% faster than the full broadcast, sustained over a 50-step long run (delta sync mean 1.73 s vs nccl 2.64 s; whole-step 7.42 s vs 8.35 s). Step 1 ships a near-empty delta (the weights synced before the first optimizer step are identical to the seed), which shows the fixed diff+gather floor (~1.3 s) vs nccl's constant full broadcast.
Where the time goes (profiled): the sparse gather is ~3× cheaper than the full all-gather (0.47 s vs 1.40 s of exposed GPU time); the broadcast itself is fully overlapped in both backends, so the interconnect is not the bottleneck on this cluster — the win comes from moving 1/13 of the bytes through the gather/pack pipeline.
Correctness
- Bit-exact by construction (integer-view diff), per-flush checksum verified inside the receiver (fail loud); 0 mismatches across all runs.
- 50-step training-equivalence run: reward/KL curves co-move with the nccl baseline (per-step reward correlation 0.76; mean |diff| 0.087, i.e. temperature-1.0 sampling noise on 32 samples/step).
- Receiver stages no full-model mirror (applies in place through SGLang's stock
--custom-weight-loaderhook), so rollout runs atgpu_memory_utilization=0.7— receiver peak memory is one bucket + one decode chunk, independent of model size.
W&B report (both 50-step runs, curves overlaid)
(delta-sharded-50step = this PR's delta_sharded backend; nccl-50step = the stock full-broadcast baseline.)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds delta weight sync for the disaggregated (one-step-off) path: after each training step the trainer broadcasts only the parameters that changed since the previous sync — RL updates leave >99% of BF16 weight bytes unchanged step-over-step — cutting weight-sync traffic to the sparsity ratio while staying bit-exact (a per-flush checksum is verified on the receiver). Two backends:
delta— the trainer byte-diffs each full parameter against a pinned-CPU full-model snapshot on rank 0 and broadcasts the changed(position, value)pairs over the existing NCCL collective group; the rollout worker reconstructs the full tensors locally and applies them through the ordinary weight-update path (no SGLang-side delta receiver required).delta_sharded— pushes the diff below the all-gather: each actor rank pins a snapshot of only its FSDP shard, byte-diffs the shard locally, and gathers only the changed(pos, val)to rank 0. Gather volume drops from the full parameter to the sparsity ratio and no rank holds a full-model snapshot. The assembled delta is bit-identical todelta, so the wire format, checksum, and rollout-side receiver are all reused unchanged.Key files
verl/checkpoint_engine/delta_checkpoint_engine.py(sender/receiver over NCCL) +verl/checkpoint_engine/delta_sync/(encode / diff-state / flush).verl/checkpoint_engine/delta_sync/sharded.py(local_shard_view/shard_delta_indices/gather_v_to_rank0) +DeltaShardedCheckpointEngine; FSDPget_per_tensor_param_shard.Scope
Disaggregated (
hybrid_engine=False) + SGLang rollout, BF16. Sharded path covers FSDP2Shard(0)+ replicated / non-DTensor params (incl. 2D FSDP×SP meshes); other shard dims raiseNotImplementedError.Validation
delta==full-sync byte-equality; sharded local-shard diff reproduces the global bytewise diff.(pos, val)set is byte-identical to afull_tensor()-gather baseline over 1D FSDP and 2D FSDP×SP meshes.