Skip to content

fix(rollout): hold payload egress for a whole weight sync - #724

Open
karkuspeter wants to merge 1 commit into
nvidia-cosmos:mainfrom
karkuspeter:pkarkus/hold-payload-egress-for-weight-sync
Open

fix(rollout): hold payload egress for a whole weight sync#724
karkuspeter wants to merge 1 commit into
nvidia-cosmos:mainfrom
karkuspeter:pkarkus/hold-payload-egress-for-weight-sync

Conversation

@karkuspeter

Copy link
Copy Markdown

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_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 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 hooks stay
optional, so a packer that ships nothing over this device's NCCL is unaffected.

What changed

  • payload_egress_held(worker) in weight_sync.py, and holds_payload_egress as a decorator
    for 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.
  • The weight-sync thread takes the same bracket in its run loop, which brings P2R under the
    guarantee too — it had none before.
  • The two inline flush_pending_sends calls 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 launches nccl_send whenever
one arrives, whatever else is using the device. Giving it a hold_sends that stops accepting
for 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 and
released, 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 only flush_pending_sends, held
around 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 passed
  • python -m pytest -q tests/test_discarded_rollout_accounting.py tests/test_terminal_drain_protocol.py tests/test_rollout_generation_mixin.py tests/contracts — passed
  • uvx ruff@0.12.7 format --check and uvx ruff@0.12.7 check on the changed files — clean

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant