feat(rollout): give the rollout backend a pause for weight sync - #725
Open
karkuspeter wants to merge 2 commits into
Open
karkuspeter wants to merge 2 commits into
karkuspeter wants to merge 2 commits into
Conversation
NCCL gives no guarantee for two communicators at once on one device, so a data packer that ships payloads over NCCL must have no send in flight while weight sync uses the device. `flush_pending_sends` drains what is already in flight, but nothing stops the packer claiming the next payload the moment it returns, and a sync spends most of its wall time after that drain, waiting on the barrier for its peers. A payload claimed in that window lands in the middle of the collective. `payload_egress_held` prefers a packer's `hold_sends` context manager, which holds for the whole sync, and falls back to the drain for a packer that only has one. Both stay optional, so a packer that ships nothing over this device's NCCL is unaffected. The two command handlers take it as a decorator, so it covers a lazy engine initialization and the barrier wait as well as the transfer, and the weight-sync thread takes it in its run loop, which also brings P2R under the same guarantee. The in-tree NCCL payload transport implements neither hook today, so it is still exposed to this: its bounded sender pool accepts a request and launches `nccl_send` whenever one arrives. Giving `NCCLRolloutMixin` a `hold_sends` that stops accepting for the duration is the natural follow-up, and needs its own multi-node validation.
Weight sync overwrites the served model's tensors in place, which is only safe while no forward is running. Nothing enforces that today: the main loop dequeues a command whenever it comes round, and whether generation happens to be idle at that moment is a property of the backend rather than anything the sync checks. For an engine whose generation call blocks the main loop it is idle by construction, which is why `RolloutBase.paused` defaults to doing nothing. Two cases are not like that. Stream generation feeds a scheduler that keeps working while the main loop dequeues commands, so a P2R receive or an R2R broadcast can land on tensors an active task is reading. And a backend that serves on a thread of its own, as an out-of-tree simulator-driven one does, never stops on its own account. `paused` is where such a backend parks, and the sync brackets itself with it. The async modes are exempt: they write a buffer clone and `install_inference_sync` swaps it in at its own safe point. `RolloutTaskScheduler` already has the `paused` this needs, added for weight sync and never called from anywhere. Wiring it up is the obvious next step and is left out here deliberately: `paused` waits out active tasks, and whether that can strand a task whose results the main loop collects needs a stream-generation run to answer.
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
1 participant
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.
Summary
Weight sync overwrites the served model's tensors in place, which is only safe while no forward
is running. Nothing enforces that today: the main loop dequeues a command whenever it comes
round, and whether generation happens to be idle at that moment is a property of the backend
rather than anything the sync checks.
RolloutBase.paused()is where a backend that is not idle parks, and the sync brackets itselfwith it. The default does nothing, which is correct for an engine whose generation call blocks
the main loop — every in-tree backend today.
Why it matters
Two cases are not idle when a command is dequeued:
main_loopcallsconsume_commandon every iteration whileRolloutTaskSchedulerkeeps working, so a P2R receive or an R2R broadcast can land on tensorsan active task is reading.
generation is a long-running episode rather than a call — never stops on its own account.
The async modes are exempt: they write a buffer clone and
install_inference_syncswaps it in atits own safe point, so the served tensors are never written under a forward.
Open question for a maintainer
RolloutTaskScheduleralready has thepaused()this needs — added for weight sync, with adocstring that says so, and never called from anywhere. Wiring it up is the obvious next step and
is left out here deliberately:
paused(wait_for_active_tasks=True)waits out active tasks, andwhether that can strand a task whose results
main_loopis responsible for collecting needs astream-generation run to answer. Happy to add it here if you can tell me it is safe, or to leave
it for someone who can run that mode.
Testing
tests/test_weight_sync_generation_pause.py, CPU-only. A backend that records when it parksshows the R2R broadcast and the P2R receive both happening inside the park, an async-mode sync
leaving generation running, and the base class's default doing nothing. Added to the pytest line
in
tests/run_test.sh.Validation
python -m pytest -q tests/test_weight_sync_generation_pause.py tests/test_weight_sync_payload_egress.py tests/test_weight_sync.py tests/test_ranked_rollout_end_and_wst_fence.py— 56 passeduvx ruff@0.12.7 format --checkanduvx ruff@0.12.7 checkon the changed files — clean