[OnePunchMonk]

Cuts startup time to the first train step without changing any component's interface. Draft.

The startup order stays as on main (engines boot, then train init). The changes make each step cheaper:

• Engine weights load with mmap instead of the RunAI streamer (copy-bound at about 2.4 GiB/s).
• sglang stays out of the processes that do not need it: the driver (sglang-d's CLI argument definitions are cached, keyed on the installed sglang source and primed at image build), the rollout manager and the engine actor. ServerArgs is built in the sglang server process, which imports sglang anyway. The train actor imports the weight-sync helpers where it builds the updater, and fsdp_utils resolves its exports lazily.
• All train ranks are created at once instead of after rank 0 has booted.
• Each train 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. Rank 0 spot-checks the filled slices against its loaded weights, ranks agree on the files first, and any failure falls back to the full broadcast. For MiniMax-H3 all 638 transformer params map 1:1. Rank 0 also no longer copies its own shard to the GPU just to have it overwritten by the broadcast.
• --rollout-prespawn-schedulers (opt-in): start the sglang-d scheduler workers before the server process imports sglang, so both imports overlap. Needs sglang's register_prespawned_schedulers ([diffusion] launch_server: accept pre-started scheduler workers OnePunchMonk/sglang#8, draft); without it the flag logs a warning and does nothing.
• The image bakes .pyc files.

Overlapping engine boot with train init gave the largest gain in the measurements, but it needs a hook in train init. That version is in OnePunchMonk#2 for reference. The full measurement history and experiment code are in OnePunchMonk#1.

[OnePunchMonk]

Summary (2x H200, full MiniMax-H3, H3 t2va recipe, time from driver start to the first weight sync):

• Default path in this PR: about 182 s. With --rollout-prespawn-schedulers plus the sglang hook: about 147 s. The original serial startup measured 462 s, on a slower box.
• What the PR contains is in the description above.
• The full measurement history and the experiment code (timing probes, benchmark hooks, per-step commits) are in Reference: startup time experiments for radixark/miles_diffusion#310 OnePunchMonk/miles_diffusion#1.
• Open item: the first weight sync and train step with the new default have not run on real hardware yet (Modal blocks CUDA IPC, so every run stopped at the sync).

@Rockdu could you review these changes, and also take a look at OnePunchMonk/sglang#8? If the approach looks fine to you, I can open it as an upstream PR in sgl-project/sglang.

[OnePunchMonk]

Pushed 2ffa788:

• The train backend no longer takes the rollout manager during init. The driver passes an optional opaque gpu_ready handle (wrapped so Ray does not resolve it before the task starts) and decides what it means: engines initialized when --startup-engine-reserve-gib is set (engines stay resident), engines offloaded otherwise. The H3 recipe sets the reserve.
• Train-side startup: sharded fill of the FSDP shards straight from the safetensors checkpoint (each rank reads its own slice), rank 0 no longer copies its shard to the GPU before the broadcast, sglang is no longer imported by the driver or at train actor import, and sglang-d's CLI args are cached (primed at image build).

CPU checks (benchmark image):

• tests/fast: 445 passed.
• Parsing the H3 recipe's args gives identical values on main and on this branch, apart from the three new flags.
• parse_args with the cache: 44 s to 3 s, without importing sglang. Train actor module import: 31 s to about 12 s.
• Sharded fill: two-rank FSDP2 test matches rank 0's load exactly, including the fallback for adapters, renamed keys and a converted tensor caught by the spot check.

GPU check (2x H200, full MiniMax-H3, one cold run in a fresh container):

• Ran end to end up to the first weight sync (which fails on Modal as before, CUDA IPC is blocked). No gate errors, no fill mismatches or failures, no OOM; the engines stayed resident and the scheduler bound as soon as the server started.
• Cold timeline: job submitted 20:00:59, train init done 20:05:35, engine healthy 20:05:57, first weight sync at 20:06:00 (about 301 s). The run was fully cold (no CLI cache in that image, weights read from cold disk, bytecode compiled on first import), and the engine was the critical path, about 20 s behind train init.
• The same-box comparison against the 147 s version did not finish within the run's budget, so the speed of this commit is not measured yet.

[OnePunchMonk]

Pushed 13306b7. Startup overlap now lives in the driver only, so no component takes another's handle.

• --startup-overlap {off,spawn,init}: the driver starts the rollout manager without waiting, creates the train ranks (and with init also runs their init) while the engines boot, then waits for both.
• --startup-keep-engines-resident: skips the engine offload after boot and the matching onload before the first sync. init needs it when colocated and --num-rollout set, otherwise it falls back to spawn.
• Fixed the sharded fill: under LoRA every key missed the checkpoint index (it kept the base_model.model. prefix), so it filled 0 of 1262 params. It now fills 638.
• Added the startup timing probes and a 0.5 s engine health poll.

Numbers (2x H200, full MiniMax-H3, H3 t2va recipe, warm page cache, driver start to the first weight sync):

run	engine boot	train init after boot	total
off	249 s	122 s	about 387 s
init + resident	163 s	0.03 s	about 182 s

Train init (122 s) is now fully hidden behind engine boot. The two runs were on different hosts and engine boot differed by 86 s, so the raw gap overstates it. Swapping in the faster boot for the serial run gives about 285 s vs 179 s, so about 106 s saved.

One run each. The sync still fails on Modal (pidfd_getfd: Operation not permitted, CUDA IPC), so the first train step has not run. CPU tests: 491 passed, 9 failed and 3 files did not collect, all in H3 SFT encoder and weight bridge tests this change does not touch.

[OnePunchMonk]

Pushed an opt-in --startup-skip-p2p-check. It sets SGLANG_SKIP_P2P_CHECK=1 on the rollout engines only, so sglang trusts the driver's P2P report instead of timing a copy between every GPU pair. Off by default. No interface changes.

2x H200, full MiniMax-H3, H3 t2va recipe, warm page cache, driver start to the first weight sync, --startup-overlap init --startup-keep-engines-resident:

run	engine boot	total
without the flag	163 s	176 s
with the flag	143 s	158 s

About 18 s saved. The runs were on different hosts, so treat it as indicative. The P2P step it removes measured 21 to 28 s in the engine logs. If a machine reports P2P but it does not work, custom all-reduce could break, so it is worth one check on the target hardware.

Things I tried that gave nothing:

• --rollout-prespawn-schedulers does nothing on the current sglang-miles branch. Workers now start through bootstrap_scheduler_process(spec), which avoids the heavy imports until activation, and the sglang hook does not apply to that file.
• Warm import times (actor 6.9 s, rollout manager 2.8 s, engine actor 0.6 s) leave nothing worth moving. Train init (about 107 s) is hidden behind engine boot.

Still not run: the first weight sync and train step (Modal blocks CUDA IPC).

[OnePunchMonk]

Pushed a cleanup so the start-up changes stay inside their components.

• actor_group.py and train_actor.py are back to main (no set_master_addr_and_port). Train ranks are created one after another again, which only matters outside --startup-overlap init, where train init is hidden behind engine boot anyway.
• The order of start-up steps now lives in one helper, start_rollout_and_train in placement_group.py. create_rollout_manager and create_training_models are unchanged, so train_sft.py and train_diffusion_async.py are untouched.
• The engine hands its server args to the child process as two plain dicts instead of a filtered args namespace. _compute_server_args still exists and behaves the same.
• Removed --rollout-prespawn-schedulers and its module. It does nothing on the current sglang-miles branch.

CPU tests: 496 passed, same 9 failures and 3 collection errors as before (H3 SFT encoder and weight bridge, not touched here). 2x H200 run with --startup-overlap init --startup-keep-engines-resident --startup-skip-p2p-check: 172 s to the first weight sync, fill filled 638 params, train init done at 105 s and the engine at 172 s. That run was on a different host than the earlier 158 s one, and the difference was in engine weight loading. The sync still fails on Modal (CUDA IPC), so the first train step has not run.