Repository navigation
[ckpt]: nccl broadcast bucket size fix - #7107
Conversation
There was a problem hiding this comment.
Code Review
This pull request updates the NCCL checkpoint engine to handle variable-sized buckets during weight transmission. Specifically, it slices the broadcast buffers to the actual offset length and includes this length in the metadata. It also updates the byte-tracking logic to accumulate the actual transmitted length instead of the fixed bucket size. There are no review comments to address.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| group_name=self.group_name, | ||
| bucket=send_buf, | ||
| metadata={"bucket_meta": bucket_meta, "is_last": False}, | ||
| bucket=send_buf[:offset], |
There was a problem hiding this comment.
recv_buf should be also restricted by offset: recv_buf[:offset]?
There was a problem hiding this comment.
The recv_buf restriction happens in BroadcastOperation._run() before calling the final broadcast collective:
https://github.com/verl-project/verl/pull/7107/files#diff-2cdb965b958cf2a60dd21504e41bd098a583b4d2b82291a612dc4c5dbae85b77R88
The recv_buf cannot be restricted when using it to create the BroadcastOperation in the receive_weights() function itself, because the offset is not known until that point.
|
@wuxibin89 do you see any other issues with this? |
…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.
### What does this PR do? The NCCL checkpoint engine always broadcasts the full bucket_size buffer per bucket, even when the bucket was only partially filled (which is the common case for the last bucket of every send_weights call, and any time the parameter stream doesn't divide evenly into bucket_size). This wasted NCCL bandwidth is proportional to the unused tail of each bucket. This PR fixes it by passing the bucket length as metadata, ensuring both sides agree on the broadcast size before the broadcast begins.
What does this PR do?
The NCCL checkpoint engine always broadcasts the full bucket_size buffer per bucket, even when the bucket was only partially filled (which is the common case for the last bucket of every send_weights call, and any time the parameter stream doesn't divide evenly into bucket_size). This wasted NCCL bandwidth is proportional to the unused tail of each bucket.
This PR fixes it by passing the bucket length as metadata, ensuring both sides agree on the broadcast size before the broadcast begins.