fix(fp8): route ModelMixin through hook-based path to survive partialload by Pfannkuchensack · Pull Request #9231 · invoke-ai/InvokeAI

Pfannkuchensack · 2026-05-24T16:59:52Z

Summary

Diffusers' enable_layerwise_casting() installs a LayerwiseCastingHook that (a) only casts dtype in pre_forward, not device, and (b) replaces Linear.forward with an instance-level wrapper that calls the original Linear.forward captured before the hook was installed. ModelCache.put() later runs apply_custom_layers_to_model, which constructs a new CustomLinear sharing the original Linear's __dict__ — so the diffusers wrapper carries over and routes calls to the captured original forward, silently bypassing CustomLinear.forward and its cast_to_device autocast.

With partial loading (e.g. FLUX.2 Klein 9B on a constrained GPU), some Linear weights stay on CPU. The diffusers pre_forward only casts dtype, so F.linear then sees input on cuda:0 and weight on cpu and raises "Expected all tensors to be on the same device".

Route every nn.Module — including ModelMixin — through _apply_fp8_to_nn_module, which uses register_forward_pre_hook / register_forward_hook(always_call=True). nn.Module._call_impl dispatches these around forward without replacing it, so CustomLinear.forward is still reached and cast_to_device moves the weight to the input device. Lose diffusers' _disable_peft_input_autocast in the process, which is irrelevant — InvokeAI patches LoRAs through CustomLinear's _patches_and_weights, not PEFT BaseTunerLayer.

Add regression test that asserts the ModelMixin branch calls _apply_fp8_to_nn_module and not enable_layerwise_casting.

Related Issues / Discussions

https://discord.com/channels/1020123559063990373/1508132779164962850

Reported on Discord: FP8 storage on FLUX.2 Klein 9B crashes with
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument mat2 in method wrapper_CUDA_mm)
at Flux2FeedForward.linear_out inside ff_context.

Stack trace points to the diffusers LayerwiseCastingHook wrapper (diffusers/hooks/hooks.py:189 → torch/nn/modules/linear.py:125).

QA Instructions

Repro (pre-fix):

Install FLUX.2 Klein 9B (Diffusers format) + matching Qwen3 8B encoder.
In Model Manager → FLUX.2 Klein 9B → enable FP8 Storage.
Constrain VRAM so partial loading kicks in (e.g. set max VRAM well below 9 GB, or run on a 12 GB GPU with other models cached).
Generate an image with the Flux2 Klein workflow.
Pre-fix: crash with the device-mismatch RuntimeError at linear_out in the first transformer block.
Post-fix: generation completes normally.

Regression coverage:

pytest tests/backend/model_manager/load/test_load_default_fp8.py — 13 tests, all green. The new test_apply_fp8_layerwise_casting_uses_hook_path_for_model_mixin fails on the pre-fix code (it would observe enable_layerwise_casting being called) and passes on the fix.

Also verify the existing FP8 paths still work:

FLUX.1 + FP8 on Diffusers format → should still cast and infer correctly.
FLUX.1 + FP8 on single-file checkpoint → already used _apply_fp8_to_nn_module, behavior unchanged.
SDXL / SD1 + FP8 → should still work.
LoRA on top of an FP8 base model → CustomLinear._autocast_forward_with_patches branch should fire (covered by test_wrap_forward_reaches_custom_linear_after_apply_custom_layers).

Merge Plan

Straight merge. No DB or schema changes. No frontend changes. Cache invalidation on the FP8 toggle already exists (drop_model on settings change), so a user toggling FP8 off/on after pulling this PR will get the fixed loader on next load.

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
Documentation added / updated (if applicable)
Updated What's New copy (if doing a release after this PR)

… load Diffusers' enable_layerwise_casting() installs a LayerwiseCastingHook that (a) only casts dtype in pre_forward, not device, and (b) replaces Linear.forward with an instance-level wrapper that calls the original Linear.forward captured before the hook was installed. ModelCache.put() later runs apply_custom_layers_to_model, which constructs a new CustomLinear sharing the original Linear's __dict__ — so the diffusers wrapper carries over and routes calls to the captured original forward, silently bypassing CustomLinear.forward and its cast_to_device autocast. With partial loading (e.g. FLUX.2 Klein 9B on a constrained GPU), some Linear weights stay on CPU. The diffusers pre_forward only casts dtype, so F.linear then sees input on cuda:0 and weight on cpu and raises "Expected all tensors to be on the same device". Route every nn.Module — including ModelMixin — through _apply_fp8_to_nn_module, which uses register_forward_pre_hook / register_forward_hook(always_call=True). nn.Module._call_impl dispatches these around forward without replacing it, so CustomLinear.forward is still reached and cast_to_device moves the weight to the input device. Lose diffusers' _disable_peft_input_autocast in the process, which is irrelevant — InvokeAI patches LoRAs through CustomLinear's _patches_and_weights, not PEFT BaseTunerLayer. Add regression test that asserts the ModelMixin branch calls _apply_fp8_to_nn_module and not enable_layerwise_casting.

JPPhoto · 2026-05-25T19:13:44Z

It looks like this needs a documentation update.

invokeai/docs/src/content/docs/configuration/fp8-storage.mdx:26 and :34 still say InvokeAI's FP8 path uses Diffusers enable_layerwise_casting, but PR 9231 removes that path for ModelMixin and routes every torch.nn.Module through InvokeAI's hook-based _apply_fp8_to_nn_module path at invokeai/backend/model_manager/load/load_default.py:221-233. This leaves the user-facing FP8 Storage docs stale for anyone diagnosing FP8 behavior or hardware impact.

lstein

Works as advertised.

JPPhoto · 2026-05-27T01:40:14Z

@lstein @Pfannkuchensack Will the documentation changes be a part of another PR?

…9241) PR #9231 routed every nn.Module — including diffusers ModelMixin — through InvokeAI's `register_forward_pre_hook` / `register_forward_hook` path, but the FP8 Storage docs still described the old `enable_layerwise_casting` implementation. Also corrects two unrelated inaccuracies the rewrite surfaced: pre-Ampere CUDA cards are not a no-op (the FP8 path gates only on `device.type == "cuda"`, and `float8_e4m3fn` is a pure storage dtype that works on any CUDA device), and the UI does not grey out the toggle based on hardware. Tell users what to include when reporting an FP8 problem so triage isn't blocked on follow-up questions: repro steps, exact model + variant, LoRA stack, partner toggles (low-VRAM, cpu_only), GPU + VRAM, OS, and the relevant log lines.

* fix(lora): sidecar-patch fp8 weights to avoid float8 add crash A full-precision Diffusers transformer loaded with fp8_storage enabled (PR #9231) keeps its Linear/Conv weights in float8_e4m3fn between forward passes. Applying a LoRA via direct patching does an in-place add on the model weight, and CUDA has no add kernel for float8, so it crashes with "ufunc_add_CUDA not implemented for Float8_e4m3fn". GGUF/BnB models avoid this because they are flagged quantized and use the sidecar path. Detect fp8 weights at the patching layer (analogous to the existing on-CPU check) and force sidecar patching for any module with float8 parameters. This takes precedence over force_direct_patching, since direct patching is not possible on fp8 weights, and fixes every model architecture at once (Flux.1/2, SD3, Qwen, SD1/SDXL UNet, etc.). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * tests: use getpixel instead of Pillow-12-only get_flattened_data Image.get_flattened_data() does not exist in Pillow 11.3.0, which is what uv.lock pins. The test only passed before because CI ignored the lock and installed the latest Pillow; the first locked test run (after the uv run fix) exposed it. getpixel() is stable across Pillow versions and already used elsewhere in this test class. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Alexander Eichhorn <alex@eichhorn.dev>

…nvoke-ai#9241) PR invoke-ai#9231 routed every nn.Module — including diffusers ModelMixin — through InvokeAI's `register_forward_pre_hook` / `register_forward_hook` path, but the FP8 Storage docs still described the old `enable_layerwise_casting` implementation. Also corrects two unrelated inaccuracies the rewrite surfaced: pre-Ampere CUDA cards are not a no-op (the FP8 path gates only on `device.type == "cuda"`, and `float8_e4m3fn` is a pure storage dtype that works on any CUDA device), and the UI does not grey out the toggle based on hardware. Tell users what to include when reporting an FP8 problem so triage isn't blocked on follow-up questions: repro steps, exact model + variant, LoRA stack, partner toggles (low-VRAM, cpu_only), GPU + VRAM, OS, and the relevant log lines.

…e-ai#9246) * fix(lora): sidecar-patch fp8 weights to avoid float8 add crash A full-precision Diffusers transformer loaded with fp8_storage enabled (PR invoke-ai#9231) keeps its Linear/Conv weights in float8_e4m3fn between forward passes. Applying a LoRA via direct patching does an in-place add on the model weight, and CUDA has no add kernel for float8, so it crashes with "ufunc_add_CUDA not implemented for Float8_e4m3fn". GGUF/BnB models avoid this because they are flagged quantized and use the sidecar path. Detect fp8 weights at the patching layer (analogous to the existing on-CPU check) and force sidecar patching for any module with float8 parameters. This takes precedence over force_direct_patching, since direct patching is not possible on fp8 weights, and fixes every model architecture at once (Flux.1/2, SD3, Qwen, SD1/SDXL UNet, etc.). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * tests: use getpixel instead of Pillow-12-only get_flattened_data Image.get_flattened_data() does not exist in Pillow 11.3.0, which is what uv.lock pins. The test only passed before because CI ignored the lock and installed the latest Pillow; the first locked test run (after the uv run fix) exposed it. getpixel() is stable across Pillow versions and already used elsewhere in this test class. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Alexander Eichhorn <alex@eichhorn.dev>

* feat(fp8): enable FP8 storage for Z-Image Z-Image was excluded from FP8 storage in #8945 because diffusers' enable_layerwise_casting() was called with the global torch dtype (fp16) while Z-Image loads in bf16: skipped modules stayed bf16, hooked ones produced fp16, and attention crashed. That root cause was fixed later in the same PR — the compute dtype now comes from the model's own parameters — so the exclusion is obsolete. Removing it alone is not enough. Our hook-based cast (#9231) dropped one thing diffusers' enable_layerwise_casting() did: honoring the model's declared _skip_layerwise_casting_patterns. Z-Image needs it, and not for quality — TimestepEmbedder.forward reads self.mlp[0].weight.dtype and casts its *input* to it. With an fp8 weight the input becomes float8 before our pre-hook restores the weight, and F.linear dies with: RuntimeError: "addmm_cuda" not implemented for 'Float8_e4m3fn' which is why ZImageTransformer2DModel declares ['t_embedder', 'cap_embedder']. _apply_fp8_to_nn_module now takes extra_skip_patterns and the caller passes the model's list. For other models this is a strict superset of our defaults (FLUX/SD3 pos_embed+norm, UNet norm, CogView4 also proj_out), so it only ever skips more. Also wire the cast into ZImageCheckpointModel: only the diffusers loader called it, so the toggle was a silent no-op for single-file Z-Image models even though both paths build the same ZImageTransformer2DModel. Tested end to end on CUDA: transformer resident VRAM drops from ~11.5GB to 5880MB for both Z-Image-Turbo (diffusers) and Z-Image-Turbo (checkpoint, 14.37GB file), with clean output images in both cases. * Chore openapi * test(fp8): drop the Z-Image entry from the exclusion parametrize main added a device-probe parametrize listing Z-Image as an excluded model. This branch removes that exclusion, so the entry contradicts `test_should_use_fp8_allows_z_image` and the case now returns the probe's value instead of False. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

* feat(fp8): enable FP8 storage for Z-Image Z-Image was excluded from FP8 storage in #8945 because diffusers' enable_layerwise_casting() was called with the global torch dtype (fp16) while Z-Image loads in bf16: skipped modules stayed bf16, hooked ones produced fp16, and attention crashed. That root cause was fixed later in the same PR — the compute dtype now comes from the model's own parameters — so the exclusion is obsolete. Removing it alone is not enough. Our hook-based cast (#9231) dropped one thing diffusers' enable_layerwise_casting() did: honoring the model's declared _skip_layerwise_casting_patterns. Z-Image needs it, and not for quality — TimestepEmbedder.forward reads self.mlp[0].weight.dtype and casts its *input* to it. With an fp8 weight the input becomes float8 before our pre-hook restores the weight, and F.linear dies with: RuntimeError: "addmm_cuda" not implemented for 'Float8_e4m3fn' which is why ZImageTransformer2DModel declares ['t_embedder', 'cap_embedder']. _apply_fp8_to_nn_module now takes extra_skip_patterns and the caller passes the model's list. For other models this is a strict superset of our defaults (FLUX/SD3 pos_embed+norm, UNet norm, CogView4 also proj_out), so it only ever skips more. Also wire the cast into ZImageCheckpointModel: only the diffusers loader called it, so the toggle was a silent no-op for single-file Z-Image models even though both paths build the same ZImageTransformer2DModel. Tested end to end on CUDA: transformer resident VRAM drops from ~11.5GB to 5880MB for both Z-Image-Turbo (diffusers) and Z-Image-Turbo (checkpoint, 14.37GB file), with clean output images in both cases. * Chore openapi * feat(fp8): enable FP8 storage for Anima The fp8_storage toggle was shown for Anima main models but did nothing: AnimaCheckpointModel never called _apply_fp8_layerwise_casting. Wire it in — the state dict is cast to a single model_dtype before load_state_dict, so the layerwise cast has one unambiguous compute dtype to restore to. Wiring alone renders a heavily dithered image with no fine detail. The cause is t_embedder: it produces the adaln_lora conditioning consumed by every block, so casting it to FP8 corrupts every token everywhere. None of the generic skip patterns match it — they target diffusers' module names (norm, pos_embed, patch_embed, proj_in/out) and this architecture names things differently. AnimaTransformer now declares _skip_layerwise_casting_patterns, the same attribute diffusers models use, so the loader needs no special-casing. Measured on CUDA, same seed/steps/CFG each run: casting nothing = broken at 1994MB; t_embedder alone = clean at 2010MB; adding x_embedder and final_layer changes nothing further (2012MB) and is kept as margin on the I/O layers; adaln_modulation was tested too and is deliberately not listed — it costs 168MB and made no difference. Against a bf16 reference (3988MB) the FP8 result keeps the same composition and loses only a little micro-detail. * fix(fp8): never apply FP8 storage to already-quantized weights Every quantized-format loader reaches _apply_fp8_layerwise_casting, and the cast there is not a no-op. Verified on real layers: - GGUF raises "Operation changed the dtype of GGMLTensor unexpectedly" at load. - bnb NF4 corrupts silently: bnb.nn.LinearNF4 subclasses nn.Linear, so the isinstance check passes and the packed uint8 payload is cast to float8. Inference still returns finite numbers and the model just produces garbage (max abs deviation 50.4 against a reference forward pass). Both are reachable today by enabling the fp8_storage toggle, which the UI offered for these models. Guard on two levels, because a format check alone is not enough — an externally quantized checkpoint can carry a plain `diffusers` format (e.g. SDNQ): - _should_use_fp8 rejects gguf_quantized and both bnb formats. - _apply_fp8_to_nn_module skips any module whose params are non-floating-point or a torch.Tensor subclass, regardless of the model's declared format. Frontend hides the toggle for quantized formats, so the control is not shown for something the backend refuses. Verified end to end: with fp8_storage forced true in the DB (the legacy case the UI no longer offers), a GGUF Z-Image model now loads cleanly with no FP8 casting and no GGMLTensor error, while non-quantized models still show the toggle and still get cast. * Merge branch 'main' into feat/fp8_quantized_guard Resolves a conflict in `_should_use_fp8` where both sides restructured the same guard chain. Upstream moved device support probing to the end of the chain (#9401, XPU), so it runs only for a model that actually wants FP8. This branch still had the older `_torch_device.type != "cuda"` precondition at the top, which would have short-circuited every non-CUDA device and undone that. Dropped it; the trailing `_device_supports_fp8_storage` call covers the same ground. Also reconciles a semantic conflict git merged cleanly: upstream added a parametrised case pinning Z-Image as excluded, while this stack removes that exclusion (#9414 gives Z-Image checkpoints a uniform model dtype, and `test_should_use_fp8_allows_z_image` documents why the exclusion is obsolete). Replaced the Z-Image case with a quantized one, which pins the property this branch is actually about: the quantized-format guard runs ahead of the device probe. The now-unused `BaseModelType` import is gone. * test(fp8): drop the Z-Image entry from the exclusion parametrize main added a device-probe parametrize listing Z-Image as an excluded model. This branch removes that exclusion, so the entry contradicts `test_should_use_fp8_allows_z_image` and the case now returns the probe's value instead of False. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(fp8): add sdnq_quantized to the quantized-format guard `ModelFormat.SDNQQuantized` was missing from both `_QUANTIZED_MODEL_FORMATS` and the frontend's `isQuantized` list, even though it is a first-class quantized format — `Main_SDNQ_{FLUX,Flux2,ZImage}_Config` and `Main_SDNQ_Diffusers_{FLUX,Flux2,ZImage}_Config` all declare `format: Literal[ModelFormat.SDNQQuantized]`. So `_should_use_fp8` still returned True for SDNQ main models and Model Manager still rendered the FP8 switch for them: exactly the dead control this change removes for GGUF and bnb. Not a corruption path today — no SDNQ loader calls `_apply_fp8_layerwise_casting` — but neither is any other quantized format, which is the point of guarding all four. Also: - Parametrize the format test over `ModelFormat` members instead of raw strings. The set under test holds strings, so a string-only test passes even if the enum values drift. - Add `test_quantized_format_set_matches_the_taxonomy`, which pins every entry to a real `ModelFormat` value so a rename cannot silently re-enable FP8 for that format. - Drop the claim that "every quantized-format loader reaches this helper" from the code comment and the test docstring. Walking `ModelLoaderRegistry` shows that none of the 17 quantized-format loaders calls `_apply_fp8_layerwise_casting` anywhere in its MRO. The guard is defense-in-depth for the next loader that gets wired up, not a live crash fix. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01561rwg8YRLzkjUu3kjc73R --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>

* feat(fp8): enable FP8 storage for Z-Image Z-Image was excluded from FP8 storage in #8945 because diffusers' enable_layerwise_casting() was called with the global torch dtype (fp16) while Z-Image loads in bf16: skipped modules stayed bf16, hooked ones produced fp16, and attention crashed. That root cause was fixed later in the same PR — the compute dtype now comes from the model's own parameters — so the exclusion is obsolete. Removing it alone is not enough. Our hook-based cast (#9231) dropped one thing diffusers' enable_layerwise_casting() did: honoring the model's declared _skip_layerwise_casting_patterns. Z-Image needs it, and not for quality — TimestepEmbedder.forward reads self.mlp[0].weight.dtype and casts its *input* to it. With an fp8 weight the input becomes float8 before our pre-hook restores the weight, and F.linear dies with: RuntimeError: "addmm_cuda" not implemented for 'Float8_e4m3fn' which is why ZImageTransformer2DModel declares ['t_embedder', 'cap_embedder']. _apply_fp8_to_nn_module now takes extra_skip_patterns and the caller passes the model's list. For other models this is a strict superset of our defaults (FLUX/SD3 pos_embed+norm, UNet norm, CogView4 also proj_out), so it only ever skips more. Also wire the cast into ZImageCheckpointModel: only the diffusers loader called it, so the toggle was a silent no-op for single-file Z-Image models even though both paths build the same ZImageTransformer2DModel. Tested end to end on CUDA: transformer resident VRAM drops from ~11.5GB to 5880MB for both Z-Image-Turbo (diffusers) and Z-Image-Turbo (checkpoint, 14.37GB file), with clean output images in both cases. * Chore openapi * feat(fp8): enable FP8 storage for Anima The fp8_storage toggle was shown for Anima main models but did nothing: AnimaCheckpointModel never called _apply_fp8_layerwise_casting. Wire it in — the state dict is cast to a single model_dtype before load_state_dict, so the layerwise cast has one unambiguous compute dtype to restore to. Wiring alone renders a heavily dithered image with no fine detail. The cause is t_embedder: it produces the adaln_lora conditioning consumed by every block, so casting it to FP8 corrupts every token everywhere. None of the generic skip patterns match it — they target diffusers' module names (norm, pos_embed, patch_embed, proj_in/out) and this architecture names things differently. AnimaTransformer now declares _skip_layerwise_casting_patterns, the same attribute diffusers models use, so the loader needs no special-casing. Measured on CUDA, same seed/steps/CFG each run: casting nothing = broken at 1994MB; t_embedder alone = clean at 2010MB; adding x_embedder and final_layer changes nothing further (2012MB) and is kept as margin on the I/O layers; adaln_modulation was tested too and is deliberately not listed — it costs 168MB and made no difference. Against a bf16 reference (3988MB) the FP8 result keeps the same composition and loses only a little micro-detail. * test(fp8): drop the Z-Image entry from the exclusion parametrize main added a device-probe parametrize listing Z-Image as an excluded model. This branch removes that exclusion, so the entry contradicts `test_should_use_fp8_allows_z_image` and the case now returns the probe's value instead of False. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix(fp8): pin Anima's skip patterns to real modules and guard the wiring Review follow-ups for #9415. Add `tests/.../test_anima_fp8_wiring.py`. Deleting the `_apply_fp8_layerwise_casting` call from the Anima single-file loader previously left the whole model_manager and anima suites green, so the dead `fp8_storage` toggle this PR fixes could come straight back with CI passing. The new boundary test fails on that mutation. The pattern test now instantiates the real `AnimaTransformer` under `accelerate.init_empty_weights()` and pins all three declared patterns to actual dotted module paths, instead of asserting a string is in a list against a hand-built stand-in. A second test records that `_FP8_DEFAULT_SKIP_PATTERNS` covers zero modules in this architecture, so the declared list is demonstrably not redundant. Lift the transformer kwargs to `ANIMA_TRANSFORMER_CONFIG` so tests build the real graph without duplicating them, mirroring `KREA2_TRANSFORMER_CONFIG`. Correct the skip-list comment. `adaln_modulation` "made no difference" was not supported by measurement: relative L2 against bf16 on a single forward goes 0.134 -> 0.091 when it is skipped, making it the largest remaining error source. The 168MB call still stands, but it rests on a 35-step A/B showing no visible difference, and the comment now says so. Also note that most of what the `final_layer` entry shields is `final_layer.adaln_modulation.*` (1.57 of 1.70M params). Stop offering FP8 storage for Anima LLLite ControlNets in the model manager. `AnimaControlNetLLLiteModel` never calls the layerwise cast, so the toggle was rendered and inert; at 16-63MB per adapter, hiding it beats wiring it. * fix(fp8): stop the hidden Anima ControlNet fp8 toggle from re-persisting Two fixes from an adversarial review of the merge: - `ControlAdapterModelDefaultSettings` hid the FP8 storage control for Anima LLLite adapters but kept sending its value. react-hook-form keeps unrendered fields in `defaultValues` (`shouldUnregister` defaults to false), so a value persisted before the control was hidden was re-sent verbatim on every save, with no UI left to clear it. Null it out wherever the control is hidden. - `test_single_file_loader_applies_fp8_layerwise_casting` passed `fp8_storage` as a top-level kwarg to `model_construct`. It is not a field of `Main_Checkpoint_Anima_Config` and the model has no `extra="allow"`, so pydantic silently discarded it and `default_settings` stayed `None` -- the toggle was off in the test that exists to prove the toggle is wired up. Build a real `MainModelDefaultSettings(fp8_storage=True)` instead. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>

Z-Image was excluded from FP8 storage in invoke-ai#8945 because diffusers' enable_layerwise_casting() was called with the global torch dtype (fp16) while Z-Image loads in bf16: skipped modules stayed bf16, hooked ones produced fp16, and attention crashed. That root cause was fixed later in the same PR — the compute dtype now comes from the model's own parameters — so the exclusion is obsolete. Removing it alone is not enough. Our hook-based cast (invoke-ai#9231) dropped one thing diffusers' enable_layerwise_casting() did: honoring the model's declared _skip_layerwise_casting_patterns. Z-Image needs it, and not for quality — TimestepEmbedder.forward reads self.mlp[0].weight.dtype and casts its *input* to it. With an fp8 weight the input becomes float8 before our pre-hook restores the weight, and F.linear dies with: RuntimeError: "addmm_cuda" not implemented for 'Float8_e4m3fn' which is why ZImageTransformer2DModel declares ['t_embedder', 'cap_embedder']. _apply_fp8_to_nn_module now takes extra_skip_patterns and the caller passes the model's list. For other models this is a strict superset of our defaults (FLUX/SD3 pos_embed+norm, UNet norm, CogView4 also proj_out), so it only ever skips more. Also wire the cast into ZImageCheckpointModel: only the diffusers loader called it, so the toggle was a silent no-op for single-file Z-Image models even though both paths build the same ZImageTransformer2DModel. Tested end to end on CUDA: transformer resident VRAM drops from ~11.5GB to 5880MB for both Z-Image-Turbo (diffusers) and Z-Image-Turbo (checkpoint, 14.37GB file), with clean output images in both cases.

Pfannkuchensack requested review from JPPhoto, blessedcoolant, dunkeroni and lstein as code owners May 24, 2026 16:59

github-actions Bot added python PRs that change python files backend PRs that change backend files python-tests PRs that change python tests labels May 24, 2026

lstein added the v6.13.x label May 25, 2026

lstein added this to Invoke - Community Roadmap May 25, 2026

lstein moved this to 6.13.x Theme: MODELS in Invoke - Community Roadmap May 25, 2026

lstein assigned JPPhoto May 25, 2026

lstein approved these changes May 26, 2026

View reviewed changes

Merge branch 'main' into fix/fp8-klein9b

592b3a5

lstein enabled auto-merge (squash) May 26, 2026 20:19

lstein merged commit 811103e into invoke-ai:main May 26, 2026
14 checks passed

Pfannkuchensack deleted the fix/fp8-klein9b branch May 26, 2026 20:57

Pfannkuchensack mentioned this pull request May 27, 2026

docs(fp8): update FP8 Storage docs #9241

Merged

5 tasks

lstein mentioned this pull request May 29, 2026

fix(lora): sidecar-patch fp8 weights to avoid float8 add crash #9246

Merged

4 tasks

This was referenced Jul 31, 2026

feat(fp8): enable FP8 storage for Z-Image #9414

Merged

feat(fp8): enable FP8 storage for Anima #9415

Merged

Pfannkuchensack mentioned this pull request Aug 19, 2026

feat(fp8): run scaled and raw fp8 checkpoints on the fp8 tensor cores #9478

Open

5 tasks

Pfannkuchensack mentioned this pull request Aug 25, 2026

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

Merged

5 tasks

Navigation Menu

Uh oh!

fix(fp8): route ModelMixin through hook-based path to survive partialload - #9231

fix(fp8): route ModelMixin through hook-based path to survive partialload#9231
lstein merged 2 commits into
invoke-ai:mainfrom
Pfannkuchensack:fix/fp8-klein9b

Pfannkuchensack commented May 24, 2026

Uh oh!

JPPhoto commented May 25, 2026

Uh oh!

lstein left a comment

Uh oh!

Uh oh!

JPPhoto commented May 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Pfannkuchensack commented May 24, 2026

Summary

Related Issues / Discussions

QA Instructions

Merge Plan

Checklist

Uh oh!

JPPhoto commented May 25, 2026

Uh oh!

lstein left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

JPPhoto commented May 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants