Repository navigation
Keep Hybrid ownership coherent across lending, activation and publication - #20
Merged
Merged
Conversation
Drain admitted requests before installation and reopen only after physical TP worker receipts and prefix-cache reset. Give each transfer a fresh snapshot identity and actual optimizer revision, while renewing NCCL sessions only when receiver membership changes. Keep training ranks in a common export failure protocol, fence NCCL stream dependencies, and ignore stale startup commands. Omit the unused KV reservation and migration APIs. Validation: 59 distinct targeted CPU cases passed across publication, serving, NCCL lifecycle, and two-rank collective export tests.
Partition one fixed global batch into complete GRPO groups, carry actual sample counts through Hybrid tasks and gradient packets, and weight Core and Hybrid contributions by those counts. Preserve empty Core lanes in the collective and give asynchronous CPU gradient packets independent storage. Use the mean of each trajectory's valid-token loss, followed by a sample mean, for the shared Core/Hybrid objective and PG/KL statistics. Zero-mask trajectories contribute zero while remaining in the batch. This intentionally changes the previous objective for variable-length microbatches larger than one; clipping, the KL coefficient and Adam remain unchanged. Include the partition, nonzero-gradient, sample-mean, worker-task and transport regressions with the implementation.
Release logits after recompute copies or loss construction, and release loss/scaled_loss after backward so unused references do not overlap the next allocation peak. Include lifetime and numerical regression coverage. Release unused CUDA cache once before the first Hybrid optimizer update, retaining the mitigation introduced after an observed NCCL allocation failure. Register its regression test with the file that implements it.
…tion Publish to the current serving receivers while preserving borrowed Hybrid replicas. Return resources through planner or sampling-supply transitions, confirm old process exit, and install current weights before reopening. Lend complete same-host replicas using the upstream TP/EP/PP/CP width and physical rollout reserve. Expand existing loans without changing their slots or generation; roll back only the new allocation after partial failure. Hand off cold startup after the first installation ACK. Ready joins wait for a Core boundary without a separate queue deadline, while worker startup and warm installation retain their existing timeout budgets. Reconcile activation and cancellation under the existing lifecycle lock. Release unused failed-join snapshots at a shared Core boundary after all loaders finish. Wait for assigned tasks during normal return, retire stale tasks after confirmed abnormal exit, and aggregate only consumed Hybrid gradient statistics. Include the lifecycle and boundary regression tests.
This was referenced Oct 22, 2026
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.
Part 3 of 3; depends on #19, which depends on #18. Merge in order: #18 → #19 → #20. This PR targets
main, so its current diff includes the two prerequisite parts until they are merged.Review only this part: EHP resource lifecycle changes.
When a rollout GPU is loaned to Hybrid training, its stopped serving process must leave the weight-publication receiver set. This change publishes to the actual serving receivers while preserving Hybrid ownership across ordinary publication. GPU return remains an explicit planner or sampling-supply operation; restored services install current weights before admission reopens.
Lending treats the requested GPU count as a target total. It selects complete same-host replicas using the existing TP × EP × PP × CP width, counts the rollout reserve in physical serving GPUs, and preserves existing Hybrid replicas, slot IDs and loan generation during expansion. Each newly selected owner must confirm exit before its GPUs enter training. Failed expansion restores and republishes only the new allocations. Runtime candidates respect the existing fixed-Core rank set and the actual rollout GPU slots; overlapping process replacements stop the previous owner before starting the next one.
All Core ranks participate in state capture and release at a shared training boundary. Workers install matching model, master parameters and Adam state through CPU/TCP live-state transfer before becoming eligible for a fixed batch. TP1 workers use the same distributed launch environment. Cold startup keeps its 180-second budget, and warm installation ACKs keep their 60-second budget. After the initial ACK, a ready Hybrid waits for a Core boundary or cancellation without a separate queue deadline. Boundary errors still propagate to the join operation.
Normal resource return and training completion wait for assigned Hybrid tasks to finish. Cancellation and failure stop the owned process, then retire unfinished inputs as cancelled after exit is confirmed. Unused snapshots are released collectively once pending loaders and activations have ended; manifests are removed after successful release. Core incorporates the sample statistics of the Hybrid gradients it actually consumed.
Validation covers complete-replica selection, incremental lending and partial rollback, process ownership, failed snapshot cleanup, task replay and activation timing. The source tree passed FYP job 1856870 on Qwen3-4B with four GPUs and preflight skipped: 50 updates, 200 unique complete GRPO groups, 800 trajectories, 12 successful Hybrid training tasks, and the same 20 held-out tasks at revisions 0, 20, 40 and 50. Final publication matched the actual TP2 + TP1 receiver set, all three physical-worker ACKs and live attestations, and the sender's measured 398 tensors. Audit, cleanup and Slurm exit checks passed. The consolidated commit has exactly the same file tree as the GPU-tested source commit 896a5ba. Its follow-up CI correction changes only test exception references: the asynchronous and Future timeout tests now use asyncio.TimeoutError and concurrent.futures.TimeoutError on Python 3.10. Production code is unchanged; all 69 tests in the four affected files pass on both Python 3.10.22 and Python 3.12.
The GPU run covers this single-node four-GPU workload. Incremental expansion and partial rollback have CPU coverage. Ordinary publication was observed with the Hybrid GPU still borrowed; subsequent training on both sides of that publication for the same active Hybrid was not exercised.
This PR builds on the fixed-batch/optimizer PR and preserves the existing fixed-Core execution model.
This three-part series supersedes #16 and #17.