Repository navigation
Startup: overlap engine boot and train init with a GPU-ready gate (invasive version) - #2
Draft
OnePunchMonk wants to merge 3 commits into
Draft
OnePunchMonk wants to merge 3 commits into
OnePunchMonk wants to merge 3 commits into
Conversation
- Overlap engine boot with train actor init: the rollout manager no longer waits for engine init, and the startup offload/onload round trip is skipped. - Colocated train ranks gate their first large GPU allocation on engine init plus a memory budget (one shard + the largest tensor + margin). If that does not fit next to the resident engines, the engines are offloaded once and onloaded again before the first weight sync. - --startup-engine-reserve-gib lets the train GPU load overlap engine boot too, when an upper bound on engine memory is known. - --serial-startup restores the old order. - Load engine weights with mmap instead of the RunAI streamer. - Keep sglang out of the rollout manager and engine actor processes: ServerArgs is built in the sglang server process, and module-level sglang imports are deferred. - Create all train ranks at once instead of waiting for rank 0 to boot. - Bake .pyc files into the image.
…mports sglang --rollout-prespawn-schedulers starts the scheduler workers in the sglang server process before its own sglang import, so the workers' import of the same stack runs in parallel instead of after it. Needs sglang's launch_server.register_prespawned_schedulers (OnePunchMonk/sglang#8); with an sglang that lacks it the waiting workers are told to exit and startup falls back to the old path. On H3 t2va (2x H200) engine boot went from about 138 s to 104 s.
…tartup The train backend no longer gets the rollout manager during init. The driver hands colocated train ranks an optional opaque gpu_ready handle to wait on before their first large GPU allocation, and decides what it stands for: with --startup-engine-reserve-gib the engines stay resident and ranks whose budget fits next to the reserve go right away (otherwise they wait for engine init); without it the engines are offloaded during train init and reloaded before the first weight sync. The H3 recipe sets the reserve. Train-side startup: - Each rank fills its FSDP shard straight from the safetensors checkpoint (its own dim-0 slice) for parameters that map 1:1 onto it; the rest (adapters, converted tensors) is broadcast from rank 0 as before, and rank 0 spot-checks the filled slices against its loaded weights. - Rank 0 no longer copies its own shard to the GPU just to have it overwritten by the broadcast. - The weight-sync helpers (which import sglang) are imported where the updater is built, and fsdp_utils resolves its exports lazily, so the driver no longer imports sglang for them. - sglang-d's CLI argument definitions are cached (primed at image build), so parsing miles args no longer imports sglang.
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
1 participant
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. The invasive part of the startup-time work, split out of radixark#310 so that PR can ship only changes that keep each component's interface as it is.
This branch is radixark#310 as it was at 2ffa788, which adds on top of the non-invasive changes:
gpu_readyhandle (throughRayTrainGroup.async_init,FSDPTrainRayActor.initandload_fsdp_models) to wait on before their first large GPU allocation. With--startup-engine-reserve-gibthe engines stay resident through train init; without it they are offloaded and reloaded before the first weight sync.--serial-startupto restore the old order.This is where most of the startup gain came from (about 190 s on the slower box in the ladder), but it changes the train init contract, which the maintainers asked to avoid. Kept here as a reference for a design discussion on where such a hook should live.
Status: GPU-validated up to the first weight sync on 2x H200 (one cold run, no gate errors or OOM). tests/fast: 445 passed. Not meant to be merged as is.