Skip to content

feat(payload): bound NCCL receive memory with explicit consumer leases - #754

Open
shuangwu wants to merge 14 commits into
nvidia-cosmos:mainfrom
shuangwu:feat/bounded-nccl-receive
Open

shuangwu wants to merge 14 commits into
nvidia-cosmos:mainfrom
shuangwu:feat/bounded-nccl-receive

Conversation

@shuangwu

@shuangwu shuangwu commented Sep 18, 2026

Copy link
Copy Markdown
Collaborator

Summary

Depends on #753, including its #747 watchdog dependency. Merge those first, then rebase this PR onto the merged implementation.

Add an opt-in NCCL receiver payload-storage budget, with concurrent receive windows sized by available bytes and explicit consumer leases. A zero custom.nccl_receive_budget_bytes preserves existing reception.

Before sender rendezvous, reserve all schema-derived decoded storage plus the largest wire buffer. Use remaining available capacity for a larger wire workspace without waiting while holding a partial reservation. Each window posts matching receives immediately after rendezvous, then completes reception and decoding before proceeding. Independent NCCL communicators are not grouped. Tight budgets produce single-payload windows; more headroom permits concurrent transfers.

Decoded storage remains charged until the consumer calls release_prefetch(streams=...) after its final readers and drops external aliases. The lease retains backing storage independently of payload-dictionary mutations and waits for reader streams before returning capacity.

Prepared prefetch retains the same lease through CPU preparation. Collecting its future transfers the lease to the ordinary consumer cache; it does not infer final use or release capacity. Unconsumed leases are cleaned up after shutdown joins the worker. Both submission paths use the shared independent watchdog, including when results are never consumed.

Boundaries

  • This budgets transport payload storage, not model/optimizer memory, allocator overhead, or memory created by preparation/training. Schema sizes predict transport decoding, not expanded training-sample memory.
  • Oversized batches fail before rendezvous. Admission waits are bounded and shutdown wakes them. There is no disk or host-payload spill.
  • Consumer integration is explicit: release the old batch before collecting the next, and do not retain prepared results/futures across shutdown. Completing CPU preparation alone does not prove final payload use.
  • UCXX does not implement this byte budget. The shared terminal-failure policy is transport-neutral.

Validation

  • Final revision 30c8862: two-GPU Redis/NCCL acceptance run passed 49 tests with zero skips, including a six-payload concurrent window and prepared CPU-copy lease handoff/release. Both processes exited cleanly.
  • Two-rank CUDA/NCCL canary passed dynamic and fixed schedules, including empty/nonfinite/recoverable-preparation cases and prefetch-on/off numerical parity for parameters, momentum and scheduler state.
  • Local regression run passed 119 tests; its distributed test passed separately with local networking enabled (the sandbox blocked loopback in the first run). CUDA cases were skipped on the CPU host. The final focused ownership/preparation regression run passed 88 tests. Ruff and whitespace checks passed.
  • Full GitHub CI remains required before merge.

The portable suites are tests/test_receive_memory.py, tests/test_receive_memory_cuda.py, and tests/test_nccl_prefetch_failfast.py; receive-memory tests are included in CI. Coverage includes byte-limited windows, admission pressure, ownership through preparation/consumption, cleanup, native-hang watchdogs, and delayed reader streams.

The prior serial-reception benchmark is historical only; its memory savings and fetch-time overhead do not describe this concurrent revision. No application-throughput or workload-specific OOM-remedy claim is made.

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