fix(rollout): hold payload egress for a whole weight sync - #724
Open
karkuspeter wants to merge 1 commit into
Open
fix(rollout): hold payload egress for a whole weight sync#724karkuspeter wants to merge 1 commit into
karkuspeter wants to merge 1 commit 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.
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
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
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_sendsdrains what is already in flight, but nothing stops the packer claimingthe next payload the moment it returns — and a sync spends most of its wall time after that
drain, waiting on its peers. A payload claimed in that window lands in the middle of the
collective.
payload_egress_heldprefers a packer'shold_sendscontext manager, which holds for thewhole sync, and falls back to the drain for a packer that only has one. Both hooks stay
optional, so a packer that ships nothing over this device's NCCL is unaffected.
What changed
payload_egress_held(worker)inweight_sync.py, andholds_payload_egressas a decoratorfor the command handlers. Bracketing the handler rather than a line inside it means the hold
also covers a lazy engine initialization and the barrier wait, not just the transfer.
guarantee too — it had none before.
flush_pending_sendscalls are replaced by the bracket.Known gap, deliberately left out
The in-tree NCCL payload transport implements neither hook, so it is still exposed to this:
NCCLRolloutMixin's bounded sender pool accepts a request and launchesnccl_sendwheneverone arrives, whatever else is using the device. Giving it a
hold_sendsthat stops acceptingfor the duration is the natural follow-up, and wants its own multi-node validation rather than
riding along here.
Testing
tests/test_weight_sync_payload_egress.py, CPU-only. The packers record when they are held andreleased, so what is checked is the ordering against the broadcast: held for the whole handler
for a packer with
hold_sends, drained once for a packer with onlyflush_pending_sends, heldaround a command on the weight-sync thread, and nothing at all for a packer with neither. Added
to the pytest line in
tests/run_test.sh.Validation
python -m pytest -q tests/test_weight_sync_payload_egress.py tests/test_weight_sync.py tests/test_ranked_rollout_end_and_wst_fence.py— 52 passedpython -m pytest -q tests/test_discarded_rollout_accounting.py tests/test_terminal_drain_protocol.py tests/test_rollout_generation_mixin.py tests/contracts— passeduvx ruff@0.12.7 format --checkanduvx ruff@0.12.7 checkon the changed files — clean