feat(weight-sync): bound in-flight weight-sync buckets - #2352
Open
0z5a wants to merge 2 commits into
Open
Conversation
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
marked this pull request as ready for review
September 7, 2026 09:36
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.
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
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 (0disables this dimension)--update-weight-max-inflight-bytes: maximum total logical bucket bytes in flight (0disables 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:
git diff --checkpassedPerformance 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.