[diffusion] feat: allow a component use retain its layerwise resident set

[mickqian]

Second of three, stacked on #40590 — its diff is included below it, so review this one for the second commit. #40590 named the question; this one gives the answer a home, and deliberately not a threshold.

Stack

1. [Diffusion] Separate a use-scoped layerwise release from release_all #40590 — separate release_after_use from release_all (no behaviour change)
2. this PR — declare the decision on the use (no behaviour change)
3. next — make the flag's help text and startup log describe what actually happens

What is being decided

A layerwise component's resident set is released when its use ends. For the DiT that is right: one use spans every denoise step, so the set has already earned the memory it holds. For a component used once per forward it means the set is prefetched at the start of the use and dropped at the end, every request — which is why --layerwise-resident-layers text_encoder=0.8 on Qwen-Image-2.1 logs resident=53/66 and moves neither memory nor latency.

Why the decision is not made here

I tried making it here, with a local headroom test, and measured both outcomes.
The same pipeline and the same flag value:

card	capacity	free once DiT+VAE are resident	the set needs	result
RTX 5090	31.4 GiB	~14.4 GiB	9.08 GiB	14.136 s → 13.737 s
RTX 4090	23.5 GiB	~6.9 GiB	9.08 GiB	server never starts

OutOfMemoryError: Tried to allocate 370.00 MiB.
GPU 0 has 23.52 GiB of which 121.81 MiB is free.
out of memory; retrying warmup at server warmup req (720x720, 2/40 steps)

Not a threshold that wants tuning. On the 4090 the right answer is simply "do not retain", and a test written at finish_use cannot reach it: what it sees is how much is free now; what it needs is how much a later stage will want. The OOM lands on a 720×720 warmup — a shape that decision point never saw.

Two attempts made it concrete, both caught by validation rather than review: measuring free memory during the use (the use's own transients are still allocated, so the patch changed nothing at all), then a fixed share of total
device memory (10% left at the decision point, 96% load peak on the 5090, and nothing boots on the 4090).

What this PR does

Declares the decision on the use, beside the placement intentions already there:

class ComponentUse:
    preferred_ready_after_request: bool = False
    keep_ready_after_warmup: bool = False
    retain_resident_layers: bool = False     # new

LayerwiseOffloadStrategy.finish_use passes it to
release_after_use(keep_resident=…), and skips park_non_layer_weights() when
it is set — parking pushes the component's non-layer weights to host, which
would undo the transfer just kept.

Nothing sets the flag in this PR, so behaviour is unchanged. #37918's
apply/validate/rollback layer is the intended setter — it already has the
per-phase device headroom #37917 enumerates candidates against, which is exactly
the resource this decision needs. The two cards above make a ready-made test
case: identical pipeline, identical resident set, opposite correct answers, and
the wrong answer on the 4090 is a boot failure rather than a slow request.

Tests

Two in test_layerwise_offload.py: the default still releases and still parks,
and the declaration reaches every manager and suppresses parking.

An earlier revision also added LayerwiseOffloadManager.retained_parameter_bytes()
to size what the declaration costs. It has been removed again (6ccefc6): the
caller that would weigh it is deliberately not in this PR, so the helper had no
production reference — only the two tests that exercised it.

Measured before that removal: test_layerwise_offload.py 81 failed / 55 passed
against #40590's 81 / 51 — the same pre-existing failures on a machine without
CUDA, plus the four tests of that revision. Dropping the helper drops two of
them; multimodal-gen-unit-test on the current head is the run that covers it,
as this file needs torch and the Mac it was edited on has none.
test_component_residency.py 26 passed, 9 skipped, unchanged: adding a defaulted
field to ComponentUse touches nothing there.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #35683690473
Latest PR Test (Extra): ✅ Run #35683690268
Latest PR Test (AMD ROCm 10): ❌ Run #35683690421

[mickqian]

/tag-and-rerun-ci extra

[mickqian]

CI status note (from the author's triage, so reviewers need not re-derive it).

Every red on this PR falls in one of three buckets, none of them from the diff:

1. NVIDIA multimodal-gen-test-2-gpu (0) — Load Latency (excluding warmup) perf validations only. All 15 validation failures in the latest run are that one metric; no correctness or e2e/denoise metric fails. Measured/expected load on runner h100-novita9-gpu-23: wan2_2_t2v_a14b_2gpu 1.48× on the first attempt and 1.26× on the in-job retry, wan2_1_i2v_14b_480P_2gpu 1.58× → 1.09×, wan2_2_i2v_a14b 1.24×, while the small checkpoints pass (flux2_modelopt_fp8_tp2 0.97×, qwen_image_t2i_2_gpus 1.12×). Second read of the same weights on the same machine is faster and the biggest checkpoints are worst: that is a cold page cache on the runner, not code. The control: [Diffusion] Correct the resident-layer help text to match its scope #40593, whose diff is help text and a doc, fails the same job with the same metric on the same runner pool.
2. AMD extra-a-*, multimodal-gen-test-*-amd — extra-a-test-1-gpu-small-amd fails identically on [Diffusion] Separate a use-scoped layerwise release from release_all #40590, [diffusion] feat: allow a component use retain its layerwise resident set #40592, [Diffusion] Correct the resident-layer help text to match its scope #40593 and [Diffusion] Add a permanent lifetime for layerwise resident layers #40599 (job logs are not uploaded by the runner); multimodal-gen-test-1-gpu-amd (0) is the whole unit-test suite timing out at the 30-minute budget with test_layerwise_offload.py fully PASSED before the cut. The sibling extra-a-test-1-gpu-large-amd passes on every PR.
3. NPU multimodal-gen-test-4-npu-a3 (0) — wan2_2_t2v_14b_w8a8_2npu: [consistency] GT image not found (404 on the pinned ci-data revision), with e2e 191197 ms vs baseline 191226 ms.

diffusion-coverage-check and the *-finish jobs are aggregators of the above (coverage saw 9 of 28 2-GPU cases because the cancelled partitions uploaded no report).

[mickqian]

CI note on multimodal-gen-test-2-gpu (1) at d5e5a79 (the main merge): 7 of 9 cases failed E2E Latency, with denoise steps at 1.8–3.3× their baselines (wan2_1_i2v_14b_480P_2gpu flat at ~2200 ms per step against 654 ms). That is the runner, not this branch and not the merged main:

partition (same head, same run, started 04:51–04:56 UTC)	runner	result
2-gpu (0)	h100-novita8-gpu-01	10/10 cases at 0.97–1.12× — incl. wan2_1_t2v_14b_2gpu, wan2_1_i2v_14b_720P_2gpu, wan2_2_t2v_a14b_2gpu, qwen_image_t2i_2_gpus_extra_high, minimax_h3_t2va_2gpu_h100
2-gpu (2)	h100-novita8-gpu-45	10/10 at 0.97–1.11× — incl. the LTX two-stage cases, wan2_1_i2v_14b_lora_2gpu, fsdp-inference
2-gpu (1)	h100-novita6-gpu-01	7/9 at 1.8–3.3×
1-gpu (0..3)	four other runners	all green, ratios normal

The same TP2 / Ulysses-2 / CFG-parallel / FSDP code paths ran at baseline on two other machines at the same minute; wan2_1_i2v_14b_720P (partition 0) sits at 1.01× while its 480P twin on partition 1 sits at 3.3×. Within the failing job the first case (04:57) and ltx_2_5_diffusion_decoder (05:21) passed at ~1.0×, so the machine was intermittently contended rather than uniformly throttled.

Also checked and excluded: the merge is clean (git diff origin/main under multimodal_gen is exactly this PR's hook, +6/+15−5/+54 tests); resolved dependency versions are identical to a passing run (torch 2.13.0, flashinfer 0.6.18, flash-attn-4 4.0.0b19, triton 3.7.1); #40646's parallel_state.py change is three context-binding fields, not a communication path; residency/parallel settings per case are identical to the last green run.

The rerun (attempt 2) was scheduled onto the same h100-novita6-gpu-01 at 06:31; if it fails the same way, the next rerun needs a different runner.

[mickqian]

Follow-up: the rerun (attempt 2) was scheduled onto the same h100-novita6-gpu-01 (06:31–07:06 UTC) and came back worse — 8 of 9 cases, wan2_1_i2v_14b_480P_2gpu at 4.69× its denoise baseline (3.37× an hour earlier), and ltx_2_5_diffusion_decoder_2gpus, which passed at 0.96× in attempt 1 on this very machine, now at 2.04× e2e / 1.38× denoise. Only the first case (flux2_modelopt_fp8_tp2, 1.02×) passes. A machine whose numbers drift like that between two runs of identical code is the variable; the two sibling partitions of this head on h100-novita8-gpu-01/-45 remain at baseline.

Re-queued the job once more; if it lands on h100-novita6-gpu-01 again I will cancel and retry. Worth taking that runner out of the pool until someone looks at it.

[mickqian]

Closing the loop on multimodal-gen-test-2-gpu (1): attempt 3 landed on a different machine, h100-novita6-gpu-23, and passed with every case at baseline — wan2_1_i2v_14b_480P_2gpu 0.99× e2e / 0.99× denoise (3.3× and 4.69× on gpu-01), flux_2_image_t2i_2_gpus 0.99× (2.74× / 2.89×), wan2_2_i2v_a14b_2gpu 0.97× / 1.00×, the rest 0.98–1.05×. Same head, same code, three runs: two on h100-novita6-gpu-01 at 2–5× and drifting, one elsewhere at 1.0×. The runner is the variable; h100-novita6-gpu-01 should be pulled from the pool until someone looks at it.