Repository navigation
[ckpt] feat: add nccl_parallel checkpoint engine (all actor ranks send) - #7263
Closed
savaresejeremy wants to merge 1 commit into
Closed
savaresejeremy wants to merge 1 commit into
savaresejeremy wants to merge 1 commit into
Conversation
Add a parallel-sender NCCL checkpoint engine: instead of actor rank 0
broadcasting every bucket, all S actor ranks send concurrently, each
owning a stripe of the global bucket sequence over its own NCCL group.
Receivers join all S groups and consume buckets in ascending global
order, so the reassembled stream is identical to the stock engine's.
Concurrent trainer-side broadcasts next to live FSDP collectives are
made safe by three mechanisms, documented in the module docstring:
bulk-synchronous phases that pin cross-communicator enqueue order,
non-blocking comm streams (patched ray stream factory, fail-closed)
with per-transfer CUDA events, and a zmq PUB/SUB subscription handshake
against the slow-joiner drop. The phase kwarg
(engine_kwargs: {nccl_parallel: {phase: N}}) caps concurrent in-phase
transfers for receivers that cannot co-schedule the full fan-out. The
per-group bucket is the configured budget divided by the sender count,
floored at 32MB and rounded down to a 16-byte boundary so a bucket-sized
cut through a large tensor cannot leave the tensors packed behind it at
a dtype-misaligned offset.
Measured (2-node, 8xB200 FSDP2 trainer to 8xB200 vLLM rollout,
Qwen3.5-27B 55.6GB bf16, GEN_TP=4, bucket 2048MB, EFA, n=20 refits):
2.51s +/-0.27 per refit vs 3.01s +/-0.25 for the stock nccl engine;
plateau accuracy (steps 61-100 of 100-step GSM8K runs) matched stock
within 0.0005. Also validated on 4-node topologies (1/2/3 trainer
nodes, S=8/16/24, 10 refits each, phase=8): all completed with normal
learning curves, one run per topology.
Tests: a correctness test mirroring test_correctness_on_gpu.py (6 GPU
cases: S=1 degenerate, phase defaulting, non-divisor and small phase<S
marker-drain paths, rebuild_group=True; 3 update rounds each) plus two
CPU checks (out-of-range phase rejected, bucket-size alignment across
S in {1,3,7,8,16,24,192}), and a 1-GPU mechanism test for the
stream-factory patch, comm-stream lookup, and cupy event completion.
No new dependencies: cupy is imported with the same guarded idiom as
the existing nccl engine.
Co-authored-by: Claude
savaresejeremy
requested review from
PeterSH6,
eric-haibin-lin,
tongyx361 and
vermouth1992
as code owners
August 4, 2026 19:28
|
|
This was referenced Aug 4, 2026
1 task done
Collaborator
|
@savaresejeremy Thanks for you contribution, busy on v0.9.0 release for now, will review it ASAP. |
2 tasks done
Author
|
@wuxibin89 I also raised #7318, an independent fix for the mixed-dtype packing alignment issue mentioned in this PR's text. It applies to the existing engines and does not depend on this PR. |
wuxibin89
pushed a commit
that referenced
this pull request
Aug 14, 2026
…ine (#7291) ### What does this PR do? Adds an opt-in `multi_sender` mode to the NCCL checkpoint engine that admits actor rank 0's node-local peers (and hence probably connected via high-speed interconnects) into the broadcast group as relays, so NCCL can drive more than one NIC per node during weight sync. Addresses #7167. With a single sender, broadcast bandwidth is capped by the one NIC reachable from the source GPU (~25 GB/s on a GH200 node with 4 NICs and ~100 GB/s bisection). Admitting the source's NVLink peers gives NCCL's broadcast tree fan-out points it can reach over NVLink and push on from over their own NICs. The relays carry no data the rollout needs — rank 0 remains the sole source, and there is still exactly one group and one broadcast per bucket. Actor workers on any other node are excluded (rank `-1`): they would have to pull a full copy over the fabric to contribute nothing. Relationship to open PRs: - **#7263** (`nccl_parallel`) targets the same issue by a different mechanism: a new backend where S actor ranks each send a stripe of the bucket sequence over S separate NCCL groups. This PR instead keeps the existing single-group broadcast and changes only group *membership*, as a flag on the existing `nccl` engine. The two are alternatives; see the discussion on #7167. - **#7107** (bucket-size fix) is a prerequisite. Relays size each broadcast to match the root, so they depend on the root sending `send_buf[:offset]`. This PR is rebased on / assumes #7107; it does not restate that change. ### Checklist Before Starting - [x] Search for similar PRs. Paste at least one query link here: https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+nccl+checkpoint+engine+broadcast and https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+7167 ### Test **GPU correctness and latency:** The following measurements include the mean over 10 iterations run on the same hardware as #7167. The measurements include the time taken for all-reduce + broadcast + copy the tensors on a non-blocking stream on the receiver side. All my measurements are based on a 4-trainer (1 node)+ 4-rollout (1 node) configuration. 1 warm-up iteration was run before starting with the measurements, and the iterations were run with `rebuild_group: False`. | Model | Mode | Mean (s) | Min–Max (s) | |---|---|---:|---:| | Qwen3.5-27B | single-sender | 5.026 | 5.02–5.04 | | Qwen3.5-27B | **multi-sender** | **1.791** | 1.78–1.86 | | Qwen3-30B-A3B | single-sender | 5.570 | 5.56–5.58 | | Qwen3-30B-A3B | **multi-sender** | **1.892** | 1.86–1.92 | Speedup: 2.81× on Qwen3.5-27B, 2.94× on Qwen3-30B-A3B. **Effects of scaling out** Scaling the rollout side out costs nothing measurable. A 4-trainer (1 node) + 12-rollout (3 nodes) configuration syncs Qwen3.5-27B in 1.82–1.88 s under multi-sender, within noise of the 1.78–1.86 s measured against 4 rollout workers on a single node. Scaling the *trainer* side across nodes does cost, but the cost is not in the broadcast. An 8-trainer (2 nodes) + 8-rollout (2 nodes) configuration syncs the same model in 2.60–2.62 s under multi-sender versus 5.85–5.90 s single-sender — still a 2.25× speedup. Both modes slow by nearly the same absolute amount relative to their single-trainer-node numbers (+0.82 s and +0.85 s), which points at a fixed cost incurred ahead of the broadcast rather than anything in the transfer itself: gathering full parameters now crosses the 25 GB/s CXI fabric instead of staying on NVLink. Only rank-0 nodes contribute relays by design, so the second trainer node adds gather cost without adding senders. ### API and Usage Example Off by default; existing behavior is unchanged. ```python CheckpointEngineConfig( backend="nccl", update_weights_bucket_megabytes=2048, engine_kwargs={"nccl": {"multi_sender": True}}, ) ``` The master's setting governs the whole group: it travels in `MasterMetadata`, so `build_topology` (a classmethod on the driver) resolves one mode for every worker. ### Design & Code Changes Rank 0 stays the only source of weights. `multi_sender` changes only who else is in the group. - **`WorkerMetadata`** — `prepare()` now reports each worker's Ray node id, since `build_topology` runs on the driver and cannot query worker placement itself. Only rank 0 fills the nested `master` field. Node id is the locality key: a node never spans more than one NVLink domain, so matching on it can under-select peers but never over-select them. - **`build_topology`** dispatches to `_single_sender_ranks` (upstream behaviour, `[0, -1, …]`) or `_multi_sender_ranks` (rank 0 plus every actor worker on its node, others `-1`), then assembles the kwargs once. Senders take the low ranks. - **`num_senders`** is passed to `init_process_group` and splits the group: ranks below it send or relay, ranks at or above it consume. It also fixes an incidental leak — only true consumers now subscribe to the zmq bucket-metadata topic, where previously every `rank > 0` did. - **`_relay_weights`** walks the same weight stream as the root via the new `get_weight_chunks_size` helper, which splits on identical boundaries to `split_weight_chunks` without materializing buffers, and issues a matching broadcast per bucket while dropping the payload. Boundaries must be derived locally rather than read off the wire: producing a weight is itself collective over the actor group, so blocking on the wire would deadlock against rank 0's gathers. - **Root wait placement** is now mode-dependent. Single sender waits just before enqueueing the next broadcast, preserving the fill/send overlap the `send_buf`/`recv_buf` swap exists for. With relays present the root waits as soon as the bucket is filled, so the root and the relays cannot issue the broadcast group's and the actor group's collectives in different orders, hence preventing a deadlock. AI assistance (claude code) was used to refactor this code and build the PR. **Why this makes sense after NCCL M2N** This PR proposes a small change to the existing legacy NCCL backend without introducing new collectives or a completely new backend. Moreover, M2N currently relies on GIN for inter-NVL domain communication, which is not supported on certain backends such as CXI.
5 of 8 tasks
12 tasks done
Author
|
Closing this. #7291 landed the multi-sender capability #7167 asked for with a much smaller change to the existing nccl engine, and carrying a second backend for the same goal is not worth the review cost. Thanks @wuxibin89. #7318 is independent of this PR and still stands on its own. If a striped-sender topology is ever wanted for a shape node-local relays cannot cover, the branch and the numbers in #7262 stay available. |
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.
What does this PR do?
Adds an additive checkpoint-engine backend,
nccl_parallel, that removes the single-sender bottleneck in cross-node weight sync. All S actor ranks send concurrently, each owning a stripe of the global bucket sequence over its own NCCL group; receivers join all S groups and reassemble in ascending global order. With S=1 it degenerates to a single-sender engine. RFC: #7262.This is the capability issue #7167 asks for, built with a different mechanism. The passive-participant topology measured there on 2026-08-03 came out ~19x slower end-to-end because the cross-group ordering it needed removed the export/broadcast overlap. This engine pins collective enqueue order with bulk-synchronous phase barriers and keeps execution overlapped through non-blocking comm streams and per-transfer CUDA events, so the export keeps running while broadcasts are in flight. Mechanism details are in the module docstring.
Measurements
2-node, 8-GPU FSDP2 trainer to 8-GPU vLLM rollout, 27B bf16 model, GEN_TP=4, bucket 2048MB, AWS EFA, n=20 refits: 2.51s ±0.27 per refit vs 3.01s ±0.25 for the stock
ncclengine. Plateau accuracy (steps 61-100 of 100-step GSM8K runs) matched stock within 0.0005. A same-window paired re-run (n=20 each, mean/median/p95) gave stock 3.005/2.911/3.533 s vs parallel 2.457/2.466/2.818 s. Both arms ran code behaviorally identical to this PR at S=8 (the 16-byte bucket rounding below is a no-op at 2048MB/8).4-node validation (same recipe, 10 refits per arm, one run per topology, phase=8): 1T+3R (S=8) 2.480 s, 2T+2R (S=16) 5.548 s, 3T+1R (S=24) 2.931 s. All completed with normal learning curves, and sync time did not move when the receiver fleet tripled. The 2T+2R split is slower at fixed phase and we don't have an explanation for that yet. A second pass of the matrix on the exact submitted tree reproduced the ordering (2.84/5.01/3.88 s, with 2T+2R again slowest); these runs are correctness evidence, not tuned benchmarks.
A note on bucket sizing: the per-group bucket is rounded down to a 16-byte boundary, with a CPU regression test. Without the rounding, 2048MB/24 senders is an odd number of bytes, and a bucket-sized cut through a large tensor leaves the tensors behind it at a dtype-misaligned offset (S=8 and S=16 divide 2048MB evenly and are unaffected). The rounding removes the bucket cut as a source of misalignment but is not a general alignment guarantee: back-to-back packing of a mixed-dtype stream can still misalign a later tensor at any bucket size. That behavior predates this engine, is shared with the stock
ncclengine, and does not arise on the homogeneous bf16 export path, so we think it belongs in its own fix rather than this PR.Duplicate-work check
nccl_parallelcheckpoint engine: aggregate sender NICs for cross-node weight sync #7262; our comment answering the open questions in that thread: [ckpt]: NCCL checkpoint engine broadcast optimization #7167 (comment). As of submission no PR implementing multi-sender transfer is open. If the announced PR appears, the designs differ architecturally (passive relay participants around one source there, active per-stripe senders with phase-pinned enqueue ordering here) and we would want to benchmark them against each other rather than treat either as a duplicate.ncclengine) is open. We are prepared to rebase and match its partial-bucket metadata behavior once it lands.kimi_ckpt_engineis inference-side (rotating owner, one bucket in flight, never concurrent with trainer collectives), so it solves a different problem.AI assistance disclosure
This contribution was developed with AI assistance (Claude, via internal agent tooling) covering code, tests, and documentation drafting. All changes were human-reviewed line by line and the branch was reviewed internally before submission; commits carry
Co-authored-by: Claudetrailers per AGENTS.md.Test
CPU (no GPU required; phase validation and bucket-size alignment):
pytest tests/checkpoint_engine/test_parallel_correctness_on_gpu.py \ -k "invalid_phase or bucket_size_alignment" -qGPU (single 8-GPU node; model defaults to ~/models/Qwen/Qwen3-8B-Base):
Results on our cluster (8x B200, 2026-08-04, the exact submitted branch content rebased on 1ae9455):
CI-cost note for maintainers: the correctness suite measured about 7 minutes wall clock on an 8x B200 node (six cases, each an 8B-model FSDP init plus 3 sync rounds; the (1, N-1) case stages the model in 7 rollout workers' host RAM), so budget more than that on L20x8. We can trim cases (for example, drop one phase<S variant) or move it to a
test_special_*name, though nothing runs those without dedicated workflow wiring. Whatever the maintainers prefer.