[ckpt]: nccl broadcast bucket size fix#7107
Conversation
![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 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.
2 participants
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.