[diffusion] Plan reversible residency before warmup

[mickqian]

Summary

• inventory the selected checkpoint components before loading and seed a conservative resident set that the post-load planner can fully replace
• plan complete resident, component-offload, and partial layerwise target states across DiTs, encoders, VAEs, lazy components, and two-stage pipelines, including independent resident counts and implicit policies per repeated layer group
• use temporary manager-level warmup counters and denoising-stage attribution so repeated layer groups, including non-standard component keys, receive independent residency choices while one-shot or unused groups do not consume stage-resident VRAM
• use one exact selection vector for every rank/phase VRAM constraint, placement-transition scratch, replica-pooled HostPin allowance, and host repacking headroom; assign per-rank pin quotas only after the joint solve
• account the full persistent host-store footprint when an ordinary resident component first becomes layerwise; reuse existing CPU weights and charge only incremental repacking scratch from component offload
• charge complete apply-time device deltas for coarse/layerwise transitions, and model prefetch with the same weight footprint as active use for all current strategies
• run the same startup sequence for server and offline entrypoints: static placement, one representative warmup, and at most one measured adjustment; resident-only verification is one full-shape step
• preserve the current HostPin layout during static weight-only planning and construct the complete pin frontier only after the representative warmup can score it
• keep loader-required placement fixed for LTX2.3 original mode, retain allocator reserve after the realized placement runs, and let numerically non-invariant pipelines decline automatic residency changes

Design boundary

The preload decision is a replaceable load-feasibility seed, not an independent serving optimizer. Its load VRAM budget is a separate lifecycle feasibility constraint; it is never composed after the fact as an independently optimal serving decision. The serving solve uses one selection vector for every runtime phase VRAM, placement transition, and node HostPin constraint. Once modules have exact sizes, static planning may demote or replace every seeded component before the first forward; warmup then validates and refines the complete serving frontier with measured workload phases and latency utility.

Structural loading choices such as explicit FSDP remain outside this planner because changing them after calibration requires reloading weights. Explicit/online quantization and direct GPU loading likewise fix only their affected components; unrelated components in the same pipeline remain eligible for measured placement. The exactness guarantee is therefore relative to the supplied replaceable placement frontier, resource estimates, and target workload, rather than a claim about startup-plus-serving lifecycle cost.

Within one request replica, multi-process HostPin is pooled into the same placement solve, so TP/SP/PP/CFG ranks may receive asymmetric quotas while their sum remains bounded. Different DP replicas retain disjoint local-worker shares rather than coordinating one cross-replica solve; homogeneous startup remains symmetric, while heterogeneous DP workloads are outside the exactness claim. Quotas are applied transactionally and rollback restores the prior allocation.

Validation

B200: Wan2.2 A14B NVFP4, 640x384x17, 12 steps

The load-safe seed selected both DiTs resident. The text encoder and VAE retained their native layerwise loading semantics; the 4-step full-shape probe then selected both resident. This preserves checkpoint/model initialization behavior while moving serving placement out of model-specific thresholds. The post-adjustment pass validates the complete serving shape; resident-only changes use one denoise step because they remove hooks/transfers without changing the compute path.

metric	old-head CI result
startup to ready	80.4 s
E2E	2740.56 ms
average denoise	209.78 ms
load / runtime peak	25182 / 52378 MiB
consistency	passed

Conservative preload boundary

A 29.7 GiB Wan2.1 1.3B experiment showed why checkpoint bytes alone cannot promote auxiliary modules safely: a weight-only plan exhausted activation headroom near 29.1 GiB and then paid for layerwise recovery. LTX2.3 showed the same issue on a larger request. The production preload seed therefore covers native DiTs only; encoders and VAEs retain their configured loading path until the bounded full-shape probe supplies activation evidence.

The old-head NVIDIA run passed unit, BCG, component-accuracy, B200, RTX 5090, four of five H100 1-GPU partitions, and one completed H100 2-GPU partition before the three known long-running partitions were cancelled to avoid seven retries against stale VRAM guards. The current-head run validates workload-faithful CI sampling, resident-only one-step validation, and explicit before/after calibration timing.

Focused tests cover initial placement, the exact joint solver, per-rank checkpoint loading, HostPin accounting, layerwise transitions, and warmup/rollback orchestration. Changed-file and merge-conflict pre-commit passed on 93f1464c2af64. The current head preserves node-pooled HostPin and per-group policy choices, avoids speculative pin-frontier construction before calibration, and collapses a measured pin frontier only when every rank-local share and the exact per-component latency cap prove it non-binding.

Completed H100 evidence from parent run 33183409022 additionally showed why both planner layers are needed: LTX2.3 HQ selected a live-allocation-feasible layout whose mapped allocator pool left no request reserve, while SANA-WM changed numerics after runtime promotion. The stack now combines the child branch allocation-gap correction during candidate planning with the parent post-placement peak-reserved guard, fixes both LTX DiTs in original mode, and keeps SANA-WM on its proven static placement. No additional warmup or retry round was added.

Stack

Depends on #35335.

---

CI States

Latest PR Test (Base): ❌ Run #34829333187
Latest PR Test (Extra): ❌ Run #34829333044
Latest PR Test (AMD ROCm 10): ❌ Run #34829333150

Baseline note (2026-09-05)

• 5090.json turbo_wan2_1_t2v_1.3b load peak 928 → 3810 MiB: the load-safe seed keeps the 1.3B DiT resident from load on a 32 GiB card (same effect already recorded for flux_2_klein_base_image_t2i and zimage_image_t2i).

[mickqian]

/tag-and-rerun-ci

[mickqian]

/tag-and-rerun-ci