Skip to content

Keep Hybrid ownership coherent across lending, activation and publication - #20

Merged
kaiwenKevinn merged 9 commits into
NetX-lab:mainfrom
tcyjohn:codex/final-ehp-20261022
Oct 23, 2026
Merged

kaiwenKevinn merged 9 commits into
NetX-lab:mainfrom
tcyjohn:codex/final-ehp-20261022

Conversation

@tcyjohn

@tcyjohn tcyjohn commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

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.

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.
@kaiwenKevinn
kaiwenKevinn merged commit 7a966c7 into NetX-lab:main Oct 23, 2026
3 checks passed
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