fix(fp8): honor every declared skip list, and stop overshooting the RAM reservation

[Pfannkuchensack]

Summary

Merge after #9415. This is a follow-up to the reviews on #9414 (merged) and #9415 (open). The docs here describe Anima's FP8 support, which lands in #9415 — merging this first would ship documentation for behaviour that isn't in main yet. No code conflict either way; the ordering is about the docs being true when they land.

Six review follow-ups, none of which changes what FP8 storage does on a correctly-behaving model today. Two are real defects, three are guards and correctness hardening, one is user-facing documentation that currently says the opposite of what the code does.

1. _keep_in_fp32_modules was not honored. Diffusers' enable_layerwise_casting() unions two class attributes before casting — _skip_layerwise_casting_patterns and _keep_in_fp32_modules. We replaced that call with our own hook-based path in #9231 and only ever read the first, so a model declaring the second would have lost its exclusions silently. Both are now read through _model_declared_skip_patterns().

This is inert today, and I checked rather than assumed — on Krea-2, Wan 14B, Z-Image and FLUX.1 it protects zero additional modules. Wan's time_embedder sits under condition_embedder, which its _skip_layerwise_casting_patterns already names; scale_shift_table is a bare Parameter, not a castable layer; Krea-2's entries are all norm*, already covered by _FP8_DEFAULT_SKIP_PATTERNS. So: no behaviour change now, and it stops being a trap for the next architecture we add.

2. Peak RAM overshot the make_room() reservation. load_state_dict(sd, assign=True) aliases every parameter to its state-dict tensor, so the dict keeps the whole model alive a second time. The FP8 cast then allocated each fp8 copy while the compute-dtype original was still reachable through sd, putting peak RAM roughly 50% over what was reserved — about 17.4 GB actual against an ~11.5 GB reservation for Z-Image. sd.clear() before the cast lets each original free as soon as its parameter is cast. Nothing reads sd after the load, in either loader. Same shape existed in the Krea-2 single-file loader; fixed there too.

3. The single-file FP8 wiring had no regression guard. Deleting the _apply_fp8_layerwise_casting call from the Z-Image single-file loader left the entire tests/backend/model_manager suite green — the dead toggle #9414 fixed could come straight back with CI passing. test_z_image_fp8_wiring.py closes that, and also pins the sd.clear() fix and the aliasing premise it rests on. (#9415 gets the equivalent guard for Anima.)

4. transformer.dtype in the Z-Image denoise loop was dormant, not safe. z_image_denoise.py built latent_model_input from transformer.dtype, which reports the float8 storage dtype once FP8 storage is on. It happens to work only because get_parameter_dtype returns the first floating-point parameter in named_parameters() order, which is the root-level x_pad_token — not a Linear, so never cast. Move the pad tokens under a submodule, or put a Linear ahead of them, and the loop starts feeding float8 into F.linear. Both sites now use get_model_compute_dtype(), which is what backend/util/fp8.py exists for, and the coupling to diffusers' parameter ordering is gone.

5. A comment blessed a pre-existing bug. The comment above the Z-Image cast described dropping .scale_weight / scaled_fp8 as filtering out metadata. It isn't: for a ComfyUI scaled-fp8 checkpoint those keys are the scales, and discarding them loads the raw fp8 codes unscaled — i.e. wrong weights. That bug is pre-existing and out of scope here, but the comment read as though the cast made it safe. Reworded to say plainly what happens.

6. The FP8 docs contradicted the code. fp8-storage.mdx still carried the row | Z-Image (any variant) | No — dtype mismatch with skipped layers | and named Z-Image in the troubleshooting exclusion list a user is told to check when VRAM doesn't drop — the exact exclusion #9414 deleted. Both corrected, Anima and its LLLite adapters added, and the skip-list description now says that a model's own declared exclusions are honored on top of the generic list.

That last point has a measurable cost worth documenting, since it was understated in #9414. Measured on meta-device builds with real configs, counting only what the declared lists protect beyond the generic defaults:

Model	Weights kept at compute precision	Saving given up
Wan 14B	232.0 M (condition_embedder, patch_embedding)	~221 MiB
Krea-2	39.3 M (time_embed)	~38 MiB
Anima	18.6 M (t_embedder, x_embedder, final_layer)	~18 MiB
FLUX.1, Qwen-Image	0	0

Wan users on a tight budget will see ~220 MiB more usage than the generic defaults alone would give. It is the right direction — it is what diffusers intends — but it should be written down.

One correction to #9415 while I was here: that PR's comment describes x_embedder + final_layer as "~2MB of margin", which is right, but the total Anima delta is 17.8 MiB, because t_embedder alone is 16.8 M parameters. That reconciles exactly with the size table in #9415 (2012.0 − 1994.2 MB = 17.8 MB).

Related Issues / Discussions

• Follow-up to the review on feat(fp8): enable FP8 storage for Z-Image #9414 (notes 1–4, 6, 7) and the review on feat(fp8): enable FP8 storage for Anima #9415 (the Anima documentation).
• Builds on feat: add per-model FP8 layerwise casting for VRAM reduction #8945 (FP8 storage), fix(fp8): route ModelMixin through hook-based path to survive partialload #9231 (hook-based casting), feat(fp8): enable FP8 storage for Z-Image #9414 (Z-Image + extra_skip_patterns).

Noted, deliberately not fixed here: WanCheckpointModel._load_from_singlefile never calls _apply_fp8_layerwise_casting, so the FP8 toggle is rendered and inert for single-file Wan checkpoints — the same dead-toggle shape #9414 and #9415 fixed for Z-Image and Anima, but on a 14B model where wiring it is clearly worth more than hiding it. Out of scope for a review follow-up; happy to open a separate PR.

Also unchanged: the ComfyUI scaled-fp8 key filtering in the Z-Image loader (point 5). Same issue raised on #9478.

QA Instructions

No CUDA GPU needed for the automated checks; the peak-RAM check needs a real Z-Image single-file checkpoint.

Unit tests

uv run --extra cuda --extra test pytest tests/backend/model_manager -q --no-cov

Expect 1192 passed, 147 skipped, 1 xfailed.

The new guards are load-bearing — verify they bite. Each of these must fail:

Mutation	Result
Delete sd.clear() from z_image.py	1 failure — test_state_dict_is_released_before_the_fp8_cast
Delete the _apply_fp8_layerwise_casting(...) line from z_image.py	2 failures — test_single_file_loader_applies_fp8_layerwise_casting and test_state_dict_is_released_before_the_fp8_cast
Drop _keep_in_fp32_modules from the loop in _model_declared_skip_patterns	3 failures — test_model_declared_skip_patterns_unions_both_diffusers_attributes, ..._tolerates_missing_and_odd_declarations, test_keep_in_fp32_modules_are_not_cast

Peak RAM (point 2). Needs a Z-Image single-file checkpoint and FP8 Storage enabled on it. Watch the InvokeAI process's RSS across the load — on main it peaks around 17.4 GB against the ~11.5 GB the loader reserved; here it should track the reservation. The log line is unchanged:

FP8 layerwise casting enabled for <model> (storage=float8_e4m3fn, compute=torch.bfloat16, param_size=...)

Regression (points 1 and 4). Generate on Z-Image with FP8 Storage on and off, fixed seed. Output must be unchanged from before this PR in both cases — points 1 and 4 are hardening, not behaviour changes. Disable the invocation cache first (PUT /api/v1/app/invocation_cache/disable), or the second run just replays the first. Worth one FLUX.1 and one Wan generation too, since _model_declared_skip_patterns is on the shared path: FLUX.1 should be bit-identical, Wan unchanged (its _keep_in_fp32_modules adds nothing, as measured above).

Docs. Render docs/src/content/docs/configuration/fp8-storage.mdx and check the "What FP8 Storage applies to" section — no Z-Image exclusion row, no Z-Image in the troubleshooting list, and the two tables format correctly.

Merge Plan

Merge after #9415. The Anima rows in fp8-storage.mdx document support that lands there; merging this first would ship docs ahead of the code. There is no textual conflict between the two branches — #9415 touches only its three Anima files, and the tests added here are appended at the end of test_load_default_fp8.py.

#9416 and #9478 stack above #9415 and both touch _apply_fp8_to_nn_module. This PR touches its caller and its docstring, not its signature, so it should pass through cleanly — but update the stack downward as usual rather than merging each branch against main separately.

No DB schema, no redux slice, no API schema change.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration — n/a, no frontend changes
• Documentation added / updated (if applicable) — fp8-storage.mdx; this is point 6 above
• Updated What's New copy (if doing a release after this PR) — n/a, no user-visible feature change