perf(startup): overlap engine boot with train init + mmap loading

[Rockdu]

What

• Overlap engine boot with train-actor init: RolloutManager.__init__ no longer blocks on ray.get(engine.init...) (the long-standing TODO at init_rollout_engines); the driver's first offload becomes fire-and-forget, and every later engine touchpoint (offload/onload/generate/eval/get_rollout_engines_and_lock) syncs through _ensure_engines_ready(). --debug-rollout-only keeps the old synchronous path. VRAM safety under colocate: only the CPU-side part of train init overlaps engine boot — the first large train-side GPU allocation (FSDP wrap + rank-0 broadcast) gates on boot_offload, which releases the engines' boot-time weights itself, exactly once. Ray only orders actor tasks per caller, so the gate must not assume the driver's own submission ran first; making the release idempotent and self-triggering is what keeps rollout and training residency sequential for any model size. The gate is skipped when rollout keeps its weights resident (disaggregated placement).
• Default sglang-d engines to mmap weight loading (SGLANG_USE_RUNAI_MODEL_STREAMER=false in the engine env): safetensors safe_open replaces the RunAI streamer's read-into-private-buffer + per-tensor clone pipeline. Operators can re-enable the streamer by exporting the var in the launcher env (still the right choice for cold remote FS).
• Startup timing probes: every process on the startup path (driver, RolloutManager, each engine actor, each train rank) emits parseable STARTUP_TIMING lines; scripts/parse_startup_timing.py rebuilds a full cross-process timeline from a single ray job log.

Why

Startup (launch → first update_weights done) took 4–8+ min per run. Probing showed the two dominant serial blocks: engine boot (~130s: subprocess re-imports + weight load + post-boot offload) ran strictly before train-actor init (~62s) due to one ray.get, and the streamer capped weight reads at ~2.4 GiB/s even from a warm page cache (it is copy-pipeline-bound, not disk-bound, on local checkpoints).

Validation

2×2 ablation on 4× H200 (single node, MiniMax-H3 t2va GRPO recipe: tp=2 → 2 engines, 4 colocated train ranks, LoRA-IPC, weights on local /scratch, warm page cache). Two independent runs per cell, spread ≤ ±2s; metric is driver main → startup_done (first weight sync complete).

variant	total startup	vs base
C0 base	242.9s	—
C1 mmap only	216.0s	−26.9s (−11%)
C2 overlap only	188.2s	−54.7s (−23%)
C3 both	153.0s	−89.9s (−37%)

Mechanism checks: with mmap, per-engine health-wait drops 88s → 63s (DiT 61.7 GiB load 37s → 14s); independently reproduced on a second machine (2 engines, 96s → 62s). With overlap, the driver's create_rollout_manager block drops 145s → 0.1s and train init hides entirely inside engine boot. Effects compose slightly super-additively. The ablation numbers above were measured before the VRAM gate (engine weights coexisted with the broadcast peak at ~94% VRAM on H3 — fits, but only by luck; larger fp32 configs would OOM, hence the gate). With the gate, the GPU-heavy ~45s of train init serializes after engine offload, so the overlap saving shrinks to an estimated ~25s warm (more on cold cache, where the overlapped CPU-side checkpoint read dominates); gated re-measurement pending.

Caveats found during validation:

• Current safetensors get_tensor copies on CPU (measured Rss≈Pss across engines), so mmap gives no resident-RAM dedup — the read side shares the page cache and peak RAM drops ~53 GB (no streamer buffers/clones), but per-engine resident copies remain.

Files

• miles/ray/rollout.py — deferred engine init + _ensure_engines_ready(); mmap env default; probes
• miles/ray/placement_group.py — fire-and-forget first offload; probes
• miles/ray/train_actor.py, miles/backends/fsdp_utils/actor.py — probes (NCCL init, per-component load/FSDP/broadcast, first weight sync)
• miles/backends/sglang_diffusion_utils/sglang_diffusion_engine.py — probes (spawn, health wait, router registration)
• train_diffusion.py — driver phase probes
• miles/utils/startup_timing.py — probe module (import-light, print-based, cross-process)
• scripts/parse_startup_timing.py — log → timeline/rollup report

Remaining work (why draft)

• Re-run the ablation with the VRAM gate to confirm the retained saving
• Confirm on hardware that a colocate startup never has engine weights and the FSDP broadcast peak resident together (the new test pins the call ordering; peak memory itself still needs a GPU run)
• Run a full multi-rollout training to exercise generate/fault-tolerance/recover_rollout_engines under deferred init.
• Decide whether the mmap default should auto-detect local vs remote checkpoint paths instead of an env override.

Checklist

• pre-commit run --all-files passes
• Added/updated tests for new behaviour — tests/fast/rollout/test_startup_offload_ordering.py (stage-a-cpu) covers the colocate ordering invariant; plus 8 instrumented end-to-end startup runs + 2 corroboration runs on H200
• pytest -x is green — not run locally: the dev machine has no ray/torch/sglang (all fast tests in this package import them). The new test's CI registration parses, its logic was verified against an equivalent replica, and CI runs the real thing.
• If launch flags changed, python3 train.py --help still parses — no CLI flags changed; arg parsing exercised by all 8 validation runs
• If a public flag was added, it appears in the CLI reference docs — n/a, no new flags
• If an example was added, it has a real walkthrough — n/a