Skip to content

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

Merged
lstein merged 7 commits into
invoke-ai:mainfrom
Pfannkuchensack:fix/fp8-review-followups
Aug 29, 2026
Merged

fix(fp8): honor every declared skip list, and stop overshooting the RAM reservation#9538
lstein merged 7 commits into
invoke-ai:mainfrom
Pfannkuchensack:fix/fp8-review-followups

Conversation

@Pfannkuchensack

Copy link
Copy Markdown
Member

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

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

…AM reservation

Review follow-ups for invoke-ai#9414 and invoke-ai#9415. Depends on invoke-ai#9415: the docs below
describe Anima's FP8 support, which lands there.

Read `_keep_in_fp32_modules` alongside `_skip_layerwise_casting_patterns`.
Diffusers' `enable_layerwise_casting()` unions both; we replaced that call
with our own hook-based path and were reading only the first, so a model
declaring the second would lose its exclusions silently. Verified to
protect nothing extra today - on Krea-2, Wan 14B, Z-Image and FLUX.1 - so
this changes nothing now and stops being a trap later.

Release the state dict before the FP8 cast in the Z-Image and Krea-2
single-file loaders. `load_state_dict(..., assign=True)` aliases every
param to its `sd` tensor, so the compute-dtype originals stayed reachable
while `param.data.to(float8)` allocated the fp8 copies, putting peak RAM
~50% over what `make_room()` reserved (~17.4GB actual against ~11.5GB
reserved for Z-Image). Nothing reads `sd` after the load.

Add `test_z_image_fp8_wiring.py`. Deleting the cast call from the Z-Image
single-file loader previously left the whole model_manager suite green.
The new tests fail on that, on removing `sd.clear()`, and on the aliasing
premise itself, should torch ever stop assigning by reference.

Use `get_model_compute_dtype()` in the Z-Image denoise loop instead of
`transformer.dtype`. It is correct today only because `x_pad_token` happens
to be parameter zero and is never cast; move the pad tokens under a
submodule and the loop starts feeding float8 into `F.linear`.

Reword the comment above the Z-Image cast. Dropping `.scale_weight` /
`scaled_fp8` is not "filtering out metadata" - for a ComfyUI scaled-fp8
checkpoint it loads unscaled weights. That bug is pre-existing and out of
scope here, but the comment read as though the cast made it safe.

Update the FP8 docs, which still said Z-Image was excluded for a dtype
mismatch and listed it in the troubleshooting exclusion list - the opposite
of what the code has done since invoke-ai#9414. Add Anima and its LLLite adapters,
and document that a model's own declared exclusions are honored on top of
the generic skip list, with the measured cost: Wan 14B gives up ~221 MiB of
savings, Krea-2 ~38 MiB, Anima ~18 MiB, FLUX.1 and Qwen-Image nothing.
@github-actions github-actions Bot added python PRs that change python files invocations PRs that change invocations backend PRs that change backend files python-tests PRs that change python tests docs PRs that change docs labels Aug 25, 2026
@lstein lstein added the 6.14.1 label Aug 25, 2026
@lstein lstein moved this to 6.14.1: Bug fixes to 6.14.0 in Invoke - Community Roadmap Aug 25, 2026
Pfannkuchensack and others added 4 commits August 25, 2026 15:33
The "What FP8 Storage applies to" table listed Wan under "Yes", and the
skip-list cost table used Wan as its headline example, closing with advice
to budget ~220 MiB of lost saving. None of the three Wan loaders reaches
the cast:

- `WanDiffusersModel._load_model` fully overrides
  `GenericDiffusersLoader._load_model` and returns without calling
  `_apply_fp8_layerwise_casting`.
- `WanCheckpointModel._load_from_singlefile` never calls it.
- `WanGGUFCheckpointModel` never calls it (and GGUF is excluded anyway).

`MainModelDefaultSettings.tsx` renders the switch for any non-quantized
main model, so a Wan user sees the toggle, sets it, and gets nothing —
the same rendered-and-inert shape this PR's point 6 exists to stop
documenting as working.

Wan now has its own row saying the switch has no effect yet, the cost
table keeps the measurement but labels it as not applied, and the
troubleshooting bullet names Wan alongside the other cases where the
`FP8 layerwise casting enabled` log line is absent.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TfA5DYQDv45CncruoQgjAe

@lstein lstein left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adversarial review at 671077e9fa, plus one docs fix pushed as 350ec60e72. Approving.

All three code changes are correct, all six claims in the PR body check out, and the new guards are load-bearing. The only real defect I found was in the docs, and I've fixed it on the branch rather than bouncing the PR.


Fixed on the branch: the docs claimed Wan gets FP8 Storage

The new table row | Z-Image, Anima, Krea-2, Qwen-Image, Wan | Yes |, the cost table's headline Wan row, and the closing "budget for that ~220 MiB on Wan" were all false. No Wan loader reaches the cast:

  • WanDiffusersModel subclasses GenericDiffusersLoader but fully overrides _load_model (wan.py:56-103) and returns result directly — it never reaches the parent's _apply_fp8_layerwise_casting at generic_diffusers.py:50.
  • WanCheckpointModel._load_from_singlefile — no call (the PR body notes this one).
  • WanGGUFCheckpointModel — no call.

grep -n _apply_fp8_layerwise_casting invokeai/backend/model_manager/load/model_loaders/wan.py returns nothing, on this branch and on main. And the toggle is rendered: MainModelDefaultSettings.tsx:156 shows DefaultFp8Storage for any non-quantized main model — so a Wan user flips the switch, reads "Yes" in the docs, and gets nothing. That's the same rendered-and-inert shape point 6 exists to stop documenting as working, so I didn't want to ship it in the commit that fixes it elsewhere.

350ec60e72 gives Wan its own row ("Not yet — the switch is shown but has no effect"), keeps the 232 M / ~221 MiB measurement but labels it not-applied, replaces the budgeting advice, and adds Wan to the troubleshooting bullet listing the cases where the FP8 layerwise casting enabled log line is absent. Docs only; no code touched.

Worth noting the body's "deliberately not fixed here" paragraph understates the scope — it's all three Wan loaders, not just the single-file one. A follow-up PR wiring Wan up is very much worth having.


Verified

Point 1 — _keep_in_fp32_modules. Union semantics confirmed at diffusers modeling_utils.py:493-494, and re.search matching mirrors diffusers' own _apply_layerwise_casting. I scanned all of diffusers 0.39.0 independently: of the classes we load, only Wan (["time_embedder","scale_shift_table","norm1","norm2","norm3"], every entry already covered by its own ["patch_embedding","condition_embedder","norm"]) and Krea-2 (["norm","norm1","norm2","norm_q","norm_k"], all covered by the generic norm) declare one. SD1/2/XL UNet, FLUX, Qwen-Image, CogView4 declare none. Zero live delta, exactly as claimed, and the direction is conservative — it can only ever over-skip.

Point 2 — sd.clear(). Read both loaders end to end: nothing reads sd after the load. Prefix-strip and format-conversion rebind sd to a new dict object, so clear() always targets the live one; krea2's _reject_incomplete_load runs before the clear; strict=False unexpected keys are simply freed earlier than they were. The aliasing premise holds (data_ptr equality) and test_assign_true_really_aliases_the_state_dict pins it. Arithmetic checks out: ~6B params, bf16 11.5 GB + fp8 5.75 GB ≈ 17.25 GB.

Point 3 — the guards bite. Mutation matrix reproduces your QA table exactly: drop sd.clear() → 1 failure; drop the cast call → 2; drop _keep_in_fp32_modules from the loop → 3.

Point 4 — the denoise sites. Confirmed a true no-op today and correct going forward: inference_dtype = TorchDevice.choose_bfloat16_safe_dtype(device) and the loader's model_dtype come from the same function, and compute_dtype = first_param.dtype = x_pad_token.dtype = model_dtype, so old and new expressions agree with FP8 both on and off. No other model-dtype read anywhere in the Z-Image path — invokeai/backend/z_image/** is clean and the invocation's remaining .dtype reads are all on latents. Bonus: this is cheaper than transformer.dtype when FP8 is on, because the marker short-circuits the full named_modules() _diffusers_hook walk in diffusers' get_parameter_dtype (modeling_utils.py:162-169).

Points 5 and 6 — comment and docs read correctly now.

Checks: tests/backend/model_manager → 1209 passed, 142 skipped. ruff check and ruff format --check clean on every changed file.


Non-blocking follow-ups

1. fp8_storage=True is silently discarded in the new test's config. test_z_image_fp8_wiring.py::_prepare_loader does Main_Checkpoint_ZImage_Config.model_construct(..., fp8_storage=True). It isn't a field and there's no extra="allow", so default_settings stays None and _should_use_fp8(config, Transformer) returns False — I checked. This is the exact trap test_anima_fp8_wiring.py carries a comment warning against (lines 116-117), where it uses default_settings=MainModelDefaultSettings(fp8_storage=True) instead.

Harmless today, since both tests replace _apply_fp8_layerwise_casting with a lambda and the gate is never consulted — they still bite, as the mutation matrix shows. But the file's premise is "the toggle has to reach the cast", and the config it uses has the toggle off; whoever later un-mocks it to assert real fp8 weights gets a vacuous pass. Worth mirroring anima.

2. The krea2 sd.clear() has no guard. Deleting it leaves tests/backend/model_manager/load/ fully green (572 passed). z_image got a regression guard for the identical fix; krea2 got the same fix with none. test_krea2_loader_boundaries.py::test_single_file_loader_constructs_and_materializes_model is the natural home.

3. The merge-after-#9415 blocker is stale. main already has AnimaCheckpointModel._apply_fp8_layerwise_casting (anima.py:194), AnimaTransformer._skip_layerwise_casting_patterns (anima_transformer.py:1013) and test_anima_fp8_wiring.py — they came through #9416's squash. AnimaControlNetLLLiteModel has never called the cast, so the Anima ControlNet-LLLite | No row is already true. Every Anima row is accurate against main today; this can merge whenever you like.

4. Nit. The dedup in _model_declared_skip_patternspatterns.extend(p for p in declared if ... p not in patterns) — works only because list.extend appends incrementally while consuming the generator. It is correct (I verified with a Krea-2-shaped declaration → ('time_embed', 'norm', 'norm1', 'norm_q')), but a plain if p not in patterns: patterns.append(p) loop would be obviously correct rather than subtly correct.


Attacks that failed

  • sd lifecycle across the prefix-strip, GGUF-conversion and key-filter paths — every rebind targets a fresh dict; no aliasing survivor reads it after the load.
  • Clearing before an exception in the cast — leaves only a dead local.
  • Regex hostility in declared patterns — every entry across diffusers 0.39.0 is a valid regex; None, bare-string and non-string-element declarations are all handled.
  • A per-step cost regression from moving to get_model_compute_dtype in the denoise loop — it's neutral or better, never worse.
  • An over-skip regression from the _keep_in_fp32_modules union on any architecture we currently load — none exists.

@lstein
lstein enabled auto-merge (squash) August 29, 2026 17:49
@lstein
lstein merged commit f3a8e95 into invoke-ai:main Aug 29, 2026
17 checks passed
@Pfannkuchensack
Pfannkuchensack deleted the fix/fp8-review-followups branch August 29, 2026 20:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

6.14.1 backend PRs that change backend files docs PRs that change docs invocations PRs that change invocations python PRs that change python files python-tests PRs that change python tests

Projects

Status: 6.14.1: Bug fixes to 6.14.0

Development

Successfully merging this pull request may close these issues.

2 participants