feat: add per-model FP8 layerwise casting for VRAM reduction - #8945
Merged
lstein merged 21 commits intoMay 12, 2026
Conversation
Add fp8_storage option to model default settings that enables diffusers' enable_layerwise_casting() to store weights in FP8 (float8_e4m3fn) while casting to fp16/bf16 during inference. This reduces VRAM usage by ~50% per model with minimal quality loss. Supported: SD1/SD2/SDXL/SD3, Flux, Flux2, CogView4, Z-Image, VAE (diffusers-based), ControlNet, T2IAdapter. Not applicable: Text Encoders, LoRA, GGUF, BnB, custom classes
Add per-model FP8 storage toggle in Model Manager default settings for both main models and control adapter models. When enabled, model weights are stored in FP8 format in VRAM (~50% savings) and cast layer-by-layer to compute precision during inference via diffusers' enable_layerwise_casting(). Backend: add fp8_storage field to MainModelDefaultSettings and ControlAdapterDefaultSettings, apply FP8 layerwise casting in all relevant model loaders (SD, SDXL, FLUX, CogView4, Z-Image, ControlNet, T2IAdapter, VAE). Gracefully skips non-ModelMixin models (custom checkpoint loaders, GGUF, BnB). Frontend: add FP8 Storage switch to model default settings panels with InformationalPopover, translation keys, and proper form handling.
JPPhoto
requested changes
Mar 21, 2026
JPPhoto
left a comment
Collaborator
There was a problem hiding this comment.
In my quantized Krea dev setup, your code was never called - is this by design or an overlooked class?
I'd also like the UI to be tweaked so the fp8 setting appears as a single slider under Settings like CPU-only for text encoders rather than as a dual-slider in the model defaults section.
FluxCheckpointModel and Flux2CheckpointModel were missing the _apply_fp8_layerwise_casting call. Additionally, the FP8 casting only worked for diffusers ModelMixin models. Add manual layerwise casting via forward hooks for plain nn.Module (custom Flux class). Also simplify FP8 UI toggle from dual-slider to single switch, matching the CPU-only toggle pattern per review feedback on invoke-ai#8945.
Z-Image's transformer has dtype mismatches with diffusers' enable_layerwise_casting: skipped modules (t_embedder, cap_embedder) stay in bf16 while hooked modules cast to fp16, causing crashes in attention layers. Also hide the FP8 toggle in the UI for Z-Image models.
Models like Flux are loaded in bf16 but the global torch dtype is fp16, causing dtype mismatches during FP8 layerwise casting. Detect the model's actual parameter dtype and use it as compute_dtype for both diffusers ModelMixin and plain nn.Module models.
Resolve merge conflict in vae.py by keeping upstream's Anima/QwenImage VAE loader paths and dropping the FP8 call from the AutoencoderKL checkpoint path. Exclude VAEs from FP8 layerwise casting in _should_use_fp8 (both standalone ModelType.VAE and the VAE/VAEDecoder/VAEEncoder submodel types of Main models). FP8 storage causes noticeable quality degradation on VAE decode.
Member
Author
|
[2026-05-09 23:23:46,331]::[ModelLoadService]::INFO --> FP8 layerwise casting enabled for CogView4-6B (storage=float8_e4m3fn, compute=torch.bfloat16, param_size=6768MB) [2026-05-09 23:25:10,617]::[ModelManagerService]::INFO --> [MODEL CACHE] Loaded model '38d890c8-e623-46f2-8299-6f1a6f97dfee:transformer' (CogView4Transformer2DModel) onto cuda device in 4.11s. Total model size: 12148.13MB, VRAM: 12148.13MB (100.0%) [2026-05-09 23:28:06,929]::[ModelLoadService]::INFO --> FP8 layerwise casting enabled for sdxl-base-1.0 (storage=float8_e4m3fn, compute=torch.float16, param_size=2449MB) [2026-05-09 23:28:53,733]::[ModelManagerService]::INFO --> [MODEL CACHE] Loaded model 'ddab47ce-8e0a-49e6-8d67-3d1d48b2848c:unet' (UNet2DConditionModel) onto cuda device in 0.83s. Total model size: 4897.05MB, VRAM: 4897.05MB (100.0%) [2026-05-09 23:29:54,400]::[ModelLoadService]::INFO --> FP8 layerwise casting enabled for dreamshaper-8 (storage=float8_e4m3fn, compute=torch.float16, param_size=820MB) [2026-05-09 23:29:28,779]::[ModelManagerService]::INFO --> [MODEL CACHE] Loaded model '06e7f29b-74ff-4f66-87ce-9b7158e55636:unet' (UNet2DConditionModel) onto cuda device in 1.69s. Total model size: 1639.41MB, VRAM: 1639.41MB (100.0%) [2026-05-09 23:33:02,516]::[ModelLoadService]::INFO --> FP8 layerwise casting enabled for controlnet-canny-sdxl-1.0 (storage=float8_e4m3fn, compute=torch.float16, param_size=1193MB) [2026-05-09 23:33:42,713]::[ModelManagerService]::INFO --> [MODEL CACHE] Loaded model 'd9f95803-6f69-4d43-abc8-c70e1b4a2c2d' (ControlNetModel) onto cuda device in 0.85s. Total model size: 2386.12MB, VRAM: 2386.12MB (100.0%) |
Member
Author
|
[2026-05-09 23:39:48,573]::[ModelLoadService]::INFO --> FP8 layerwise casting enabled for FLUX.2-klein-9B (storage=float8_e4m3fn, compute=torch.bfloat16, param_size=8691MB) [2026-05-09 23:40:20,500]::[ModelManagerService]::INFO --> [MODEL CACHE] Loaded model '8dbb1dfe-7a08-4eae-8846-2f482e9320f8:transformer' (Flux2Transformer2DModel) onto cuda device in 4.34s. Total model size: 17316.02MB, VRAM: 17316.02MB (100.0%) |
Pfannkuchensack
marked this pull request as ready for review
May 9, 2026 21:46
Pfannkuchensack
requested review from
blessedcoolant,
dunkeroni and
lstein
as code owners
May 9, 2026 21:46
Member
Author
|
[2026-05-10 00:20:23,041]::[ModelLoadService]::INFO --> FP8 layerwise casting enabled for FLUX.1 dev (storage=float8_e4m3fn, compute=torch.bfloat16, param_size=11350MB) [2026-05-10 00:21:06,488]::[ModelManagerService]::INFO --> [MODEL CACHE] Loaded model 'fa37dbd1-9681-4d00-a5b2-dc5699237d7f:transformer' (Flux) onto cuda device in 6.37s. Total model size: 22700.13MB, VRAM: 17660.13MB (77.8%) |
Collaborator
|
…le fallback, hide ControlLoRA toggle - Add ModelCache.drop_model() and call it from update_model_record when fp8_storage or cpu_only change. These settings are baked into the loaded nn.Module at load time, so toggling them was silently a no-op until the cache entry was evicted by other means. - Replace the pre-hook/post-hook pair in _apply_fp8_to_nn_module with a forward wrapper using try/finally. register_forward_hook only fires on successful forward, so an exception left params in compute dtype and defeated the FP8 storage savings. - Hide the FP8 toggle in the UI for ControlLoRA and exclude LoRA/ControlLoRA in _should_use_fp8. LoRAs are patched into base models rather than run as a standalone forward pass, so layerwise-casting hooks would never fire. - Add tests for drop_model, the exception-safe FP8 wrapper, the ControlLoRA/LoRA exclusion, and the _load_settings_changed predicate.
Collaborator
|
Progress! Here are some more issues or things that need tests to validate current behavior:
|
…es, skip precision-sensitive layers - _wrap_forward_with_fp8_cast now dispatches via type(module).forward at call time instead of capturing the bound method. ModelCache.put() swaps nn.Linear.__class__ to CustomLinear (sharing __dict__), which would otherwise leave our instance forward shadowing CustomLinear.forward and silently bypass LoRA/ControlLoRA patch dispatch on FP8 checkpoints. - drop_model() now marks locked entries is_stale instead of skipping them silently; unlock() evicts stale entries once the last lock releases. Without this, a setting toggled during an in-flight generation survived on the locked entry and the next generation reused the pre-change module. - _apply_fp8_to_nn_module mirrors diffusers' apply_layerwise_casting: only the supported layer classes (Linear/Conv*/Embedding) get cast, and module paths matching norm/pos_embed/patch_embed/proj_in/proj_out are skipped. FLUX RMSNorm.scale and similar precision-sensitive scalars are no longer crushed to FP8. - drop_model() and the unlock-stale path now update stats.cleared and fire on_cache_models_cleared callbacks, matching _make_room_internal so the UI stats panel and observers don't miss invalidations. - Add 14 tests: class-swap dispatch, norm/pos_embed/proj_in_out skip, unsupported-type skip, stale-marking, multi-lock release, stats and callback firing for both paths, no-op silence.
Collaborator
|
Sorry for missing this one...
|
…tch survives apply_custom_layers_to_model Previous fix was wrong. `apply_custom_layers_to_model` does not do `module.__class__ = CustomLinear` — `wrap_custom_layer` constructs a NEW CustomLinear via __new__ and shares the original Linear's __dict__, then setattr installs the new object on the parent. The new object has type() == CustomLinear, but our wrapped forward closed over the original Linear instance, so `type(module).forward(module, ...)` resolved to Linear.forward on the captured old object and silently bypassed CustomLinear.forward — breaking LoRA/ControlLoRA patch dispatch for FP8 checkpoint models. Reproduced on a fresh worktree. Replace the instance-forward override with register_forward_pre_hook + register_forward_hook(always_call=True). Hooks are dispatched by nn.Module._call_impl with the actual called instance, so they fire on the new CustomLinear and self.forward resolves normally via class lookup — reaching CustomLinear.forward and its patch-aware branch. always_call=True keeps the exception-safety guarantee (post-hook fires even when forward raises). Replace the simulated __class__-swap test with one that runs real apply_custom_layers_to_model, attaches a sentinel _patches_and_weights, and asserts the patch-aware branch in CustomLinear.forward is reached. Verified the test fails under the old instance-forward implementation with the reviewer-described symptom and passes under the hook fix.
…annkuchensack/InvokeAI into feature/fp8-layerwise-casting
JPPhoto
self-requested a review
May 11, 2026 23:13
JPPhoto
approved these changes
May 11, 2026
JPPhoto
left a comment
Collaborator
There was a problem hiding this comment.
Looks good, add documentation and then it can go in the next RC so the community can try it out.
lstein
approved these changes
May 12, 2026
lstein
left a comment
Collaborator
There was a problem hiding this comment.
Looking forward to trying this.
lstein
enabled auto-merge (squash)
May 12, 2026 02:32
5 tasks
This was referenced Jul 31, 2026
5 tasks
lstein
pushed a commit
that referenced
this pull request
Aug 24, 2026
* 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>
5 tasks
lstein
added a commit
that referenced
this pull request
Aug 25, 2026
* 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>
lstein
added a commit
that referenced
this pull request
Aug 25, 2026
* 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>
JPPhoto
pushed a commit
to JPPhoto/InvokeAI
that referenced
this pull request
Sep 10, 2026
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
FP8 Layerwise Casting - Implementation
Summary
Add per-model
fp8_storageoption to model default settings that enables diffusers'enable_layerwise_casting()to store weights in FP8 (float8_e4m3fn) while casting to fp16/bf16 during inference. This reduces VRAM usage by ~50% per model with minimal quality loss.Supported: SD1/SD2/SDXL/SD3, Flux, Flux2, CogView4, Z-Image, ControlNet, T2IAdapter.
Not applicable: Text Encoders, LoRA, GGUF, BnB, custom classes.
Related Issues / Discussions
enable_layerwise_casting()(available in diffusers 0.36.0)QA Instructions
fp8_storage: truein a model'sdefault_settings(via API or Model Manager UI)Test Matrix
fp8_storage=true- load and generatefp8_storage=true- load and generatefp8_storage=true- load and generatefp8_storage=true- load and generatefp8_storage=true- load and generatefp8_storage=true- load and generatefp8_storage=true- load and generatefp8_storage=true- load and generatefp8_storage=true- load and generatefp8_storagefp8_storageis silently ignoredfp8_storage=true- load and generate # it does not workChecklist
What's Newcopy (if doing a release after this PR)