Skip to content

feat(weight-sync): bound in-flight weight-sync buckets - #2352

Open
0z5a wants to merge 2 commits into
THUDM:mainfrom
0z5a:codex/weight-bucket-credit
Open

feat(weight-sync): bound in-flight weight-sync buckets#2352
0z5a wants to merge 2 commits into
THUDM:mainfrom
0z5a:codex/weight-bucket-credit

Conversation

@0z5a

@0z5a 0z5a commented Sep 4, 2026

Copy link
Copy Markdown

Draft update — 2026-09-05

Separates logical admission, transport completion, staging and consumer acknowledgements; failed versions remain unpublished and resources remain held until the appropriate lifecycle boundary. Real two-rank Gloo tests passed. Full-model NCCL and Ray/SGLang confirmation remain pending.

AI assistance: OpenAI Codex. Remains Draft; implementation, CPU/Gloo checks and production GPU validation are separate acceptance gates. No new GPU results were generated during publication.

The original submission description and validation history follow.


Summary

  • add a reusable weight-sync credit controller with independent bucket-count and byte limits
  • pipeline ordinary full-weight NCCL and colocated tensor buckets while releasing credit in FIFO completion order
  • preserve weight-version fences, reject impossible oversized buckets, and add focused CPU/Gloo tests plus English/Chinese documentation

Motivation

The existing full-weight path waits after each converted bucket. That is safe but prevents conversion/communication overlap. This change makes overlap opt-in and bounded, so operators can choose a memory/concurrency budget rather than relying on a rank-count-specific schedule.

Semantics and compatibility

The two independent controls are:

  • --update-weight-max-inflight-buckets: maximum logical buckets in flight (0 disables this dimension)
  • --update-weight-max-inflight-bytes: maximum total logical bucket bytes in flight (0 disables this dimension)

Both default to 0. When both are zero, slime keeps the legacy one-bucket blocking path. Bucket bytes are counted once regardless of rollout-engine fan-out. Colocated trainers use the maximum per-rank bucket size through their Gloo control group so every rank chooses identical flush boundaries. Reservations complete and release FIFO, and all reservations drain before the weight version commits. A bucket larger than a configured byte limit fails immediately instead of deadlocking.

The window applies only to ordinary full-weight NCCL/tensor transfers. Explicit expert-routed transfers remain serial because they reuse staging buffers, while still enforcing the byte ceiling. Disk and delta transports reject nonzero credit settings because they do not use this in-memory bucket data plane. There are no world-size or rank-count special cases.

Validation

Correctness validation completed before publication:

  • 47 relevant CPU tests passed, including controller invariants, CLI/default compatibility, distributed launch ordering, colocated admission, expert-path constraints, and documentation checks
  • a real two-rank asynchronous Gloo test reached the configured bucket/byte ceiling and verified received tensor contents
  • repository pre-commit hooks, Python compilation checks, and git diff --check passed
  • a single-card CUDA/NCCL smoke was limited to physical GPU 6 and observed peak credits of 2 buckets / 3072 bytes with checksum 6144; no other GPU was used

Performance validation is limited to proving that multiple asynchronous buckets can be admitted up to the configured credits. No production throughput or latency improvement is claimed from that smoke. A real Ray/SGLang end-to-end run was not performed. Before merge, this draft needs a deployment benchmark covering peak memory, weight-sync duration, and rollout latency at several bucket/byte windows versus the all-zero default.

Duplicate check

I searched current open issues and PRs for weight-bucket credits, in-flight bytes/buckets, backpressure, pipelined/asynchronous weight sync, and both exact option names. I found no substantive duplicate. Existing transport proposals such as #2146 and #2159 change how weights move; they do not provide this admission/credit controller.

AI assistance and required review

This implementation, tests, documentation, and PR text were prepared with OpenAI Codex assistance. A human contributor/maintainer must review every changed line, validate the credit and collective-ordering assumptions, and reproduce the relevant tests and deployment benchmark before this draft is considered ready to merge.

Keep logical credit, transport, staging, and consumer lifetimes distinct; propagate colocated load acknowledgement across each existing TP group; poison failed versions; and publish only after all consumers resume. Add fixed-cardinality metrics, failure coverage, and opt-in 2/4-rank NCCL validation.

Assisted-by: OpenAI Codex
Signed-off-by: jiachun <2609024259@qq.com>
@0z5a
0z5a marked this pull request as ready for review September 7, 2026 09:36
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.

2 participants