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

docs/src/content/docs/configuration/fp8-storage.mdx

@@ -62,18 +62,32 @@ The setting takes effect on the next load. If the model is already in the cache,
 
 FP8 Storage is **only** applied to layers where the precision trade-off is acceptable:
 
-| Model type                    | FP8 applied?                           |
-| ----------------------------- | -------------------------------------- |
-| Main models (SD1, SD2, SDXL)  | Yes                                    |
-| FLUX.1 / FLUX.2 Klein         | Yes                                    |
-| ControlNet, T2I-Adapter       | Yes                                    |
-| VAE                           | No — visible decode-quality regression |
-| Text encoders, tokenizers     | No — small models, no benefit          |
-| Z-Image (any variant)         | No — dtype mismatch with skipped layers|
-| LoRA, ControlLoRA             | No — patched into base, not run alone  |
+| Model type                              | FP8 applied?                           |
+| --------------------------------------- | -------------------------------------- |
+| Main models (SD1, SD2, SDXL)            | Yes                                    |
+| FLUX.1 / FLUX.2 Klein                   | Yes                                    |
+| Z-Image, Anima, Krea-2, Qwen-Image, Wan | Yes                                    |
+| ControlNet, T2I-Adapter                 | Yes                                    |
+| VAE                                     | No — visible decode-quality regression |
+| Text encoders, tokenizers               | No — small models, no benefit          |
+| Anima ControlNet-LLLite                 | No — adapter is only tens of MB        |
+| LoRA, ControlLoRA                       | No — patched into base, not run alone  |
 
 Within a supported model, **norm layers, position/patch embeddings, and `proj_in`/`proj_out` are skipped** so precision-sensitive tiny learned scalars (e.g. FLUX `RMSNorm.scale`) aren't crushed to FP8. This mirrors the diffusers default skip list.
 
+**On top of that, each model's own declared exclusions are honored.** Architectures list their precision-sensitive modules themselves (diffusers calls these `_skip_layerwise_casting_patterns` and `_keep_in_fp32_modules`), and those layers stay at compute precision too. This matters where the generic patterns don't fit the architecture's naming, and it is not a nicety: Z-Image and Anima both name their timestep-embedding MLP `t_embedder`, which no generic pattern matches. Cast to FP8, Z-Image fails outright and Anima renders a heavily dithered image with no fine detail.
+
+The cost is a slightly smaller saving on the models that declare a lot. Measured against the generic defaults alone:
+
+| Model              | Weights kept at compute precision                     | Saving given up |
+| ------------------ | ----------------------------------------------------- | --------------- |
+| Wan 14B            | 232 M (`condition_embedder`, `patch_embedding`)       | ~221 MiB        |
+| Krea-2             | 39 M (`time_embed`)                                   | ~38 MiB         |
+| Anima              | 19 M (`t_embedder`, `x_embedder`, `final_layer`)      | ~18 MiB         |
+| FLUX.1, Qwen-Image | 0                                                     | 0               |
+
+If you are budgeting VRAM to the last few hundred MB on Wan, account for that ~220 MiB.
+
 ## Quality trade-offs
 
 FP8 Storage is **near-lossless** for most workloads because:
@@ -102,7 +116,7 @@ The cache eviction is immediate for idle models, but **deferred until the next u
 
 If VRAM still hasn't dropped:
 
-- Check the InvokeAI log for `FP8 layerwise casting enabled for <model name>`. If the line isn't there, the model is on the exclusion list (VAE, text encoder, Z-Image, LoRA — see table above).
+- Check the InvokeAI log for `FP8 layerwise casting enabled for <model name>`. If the line isn't there, the model is on the exclusion list (VAE, text encoder, LoRA — see table above).
 - Confirm you are on CUDA. FP8 Storage is silently disabled on CPU and MPS.
 
 ### Quality regression on a specific model

invokeai/app/invocations/z_image_denoise.py

@@ -36,6 +36,7 @@
 from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState
 from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ZImageConditioningInfo
 from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.fp8 import get_model_compute_dtype
 from invokeai.backend.z_image.extensions.regional_prompting_extension import ZImageRegionalPromptingExtension
 from invokeai.backend.z_image.text_conditioning import ZImageTextConditioning
 from invokeai.backend.z_image.z_image_control_adapter import ZImageControlAdapter
@@ -586,7 +587,9 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                     timestep = torch.tensor([model_t], device=device, dtype=inference_dtype).expand(latents.shape[0])
 
                     # Run transformer for positive prediction
-                    latent_model_input = latents.to(transformer.dtype)
+                    # `transformer.dtype` is the float8 *storage* dtype once FP8 storage is on, and
+                    # torch has no arithmetic kernels for it (see `get_model_compute_dtype`).
+                    latent_model_input = latents.to(get_model_compute_dtype(transformer))
                     latent_model_input = latent_model_input.unsqueeze(2)  # Add frame dimension
                     latent_model_input_list = list(latent_model_input.unbind(dim=0))
 
@@ -700,7 +703,8 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor:
                     # Run transformer for positive prediction
                     # Z-Image transformer expects: x as list of [C, 1, H, W] tensors, t, cap_feats as list
                     # Prepare latent input: [B, C, H, W] -> [B, C, 1, H, W] -> list of [C, 1, H, W]
-                    latent_model_input = latents.to(transformer.dtype)
+                    # See above: never build tensors from the float8 storage dtype.
+                    latent_model_input = latents.to(get_model_compute_dtype(transformer))
                     latent_model_input = latent_model_input.unsqueeze(2)  # Add frame dimension
                     latent_model_input_list = list(latent_model_input.unbind(dim=0))
 

invokeai/backend/model_manager/load/load_default.py

@@ -186,6 +186,34 @@ def _is_quantized_param(param: torch.nn.Parameter) -> bool:
     return not param.data.is_floating_point() or type(param.data) is not torch.Tensor
 
 
+def _model_declared_skip_patterns(model: torch.nn.Module) -> tuple[str, ...]:
+    """The precision-sensitive modules a model declares for itself, as skip patterns.
+
+    Diffusers' `enable_layerwise_casting()` unions two class attributes before casting:
+    `_skip_layerwise_casting_patterns` and `_keep_in_fp32_modules`. We no longer call it (see
+    `_apply_fp8_layerwise_casting`), so we have to read both ourselves — this is not cosmetic.
+    Z-Image's `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 can restore the weight, and
+    `F.linear` dies with `"addmm_cuda" not implemented for 'Float8_e4m3fn'`. Hence
+    `['t_embedder', 'cap_embedder']` for that model.
+
+    `_keep_in_fp32_modules` protects nothing extra on any model we currently load — verified on
+    Krea-2, Wan 14B, Z-Image and FLUX.1. 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; and Krea-2's entries are all `norm*`, already covered by
+    `_FP8_DEFAULT_SKIP_PATTERNS`. It is read anyway so the next model to declare one does not lose
+    it silently.
+    """
+    patterns: list[str] = []
+    for attr in ("_skip_layerwise_casting_patterns", "_keep_in_fp32_modules"):
+        declared = getattr(model, attr, None) or ()
+        # Diffusers stores these as lists of strings, but a subclass could set a bare string.
+        if isinstance(declared, str):
+            declared = (declared,)
+        patterns.extend(p for p in declared if isinstance(p, str) and p not in patterns)
+    return tuple(patterns)
+
+
 # The construction path is not thread-safe on its own; it monkey-patches process-global torch state
 # (see MODEL_LOAD_LOCK). Concurrent callers must hold the MODEL_LOAD_LOCK write lock (see
 # _load_and_cache).
@@ -466,18 +494,11 @@ def _apply_fp8_layerwise_casting(
         # `register_forward_hook` path fires around `nn.Module._call_impl` without replacing
         # `forward`, so `CustomLinear.forward` is still reached.
         if isinstance(model, torch.nn.Module):
-            # Diffusers models declare their own precision-sensitive modules in
-            # `_skip_layerwise_casting_patterns`, and `enable_layerwise_casting()` honors them. Since
-            # we no longer call it, we have to apply that list ourselves — it is not cosmetic. Z-Image's
-            # `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 can restore the weight,
-            # and `F.linear` dies with `"addmm_cuda" not implemented for 'Float8_e4m3fn'`. Hence
-            # `['t_embedder', 'cap_embedder']` for that model.
             self._apply_fp8_to_nn_module(
                 model,
                 storage_dtype=storage_dtype,
                 compute_dtype=compute_dtype,
-                extra_skip_patterns=tuple(getattr(model, "_skip_layerwise_casting_patterns", None) or ()),
+                extra_skip_patterns=_model_declared_skip_patterns(model),
             )
         else:
             return model
@@ -505,8 +526,8 @@ def _apply_fp8_to_nn_module(
         Without the skip list, precision-sensitive tiny learned scalars (e.g. FLUX RMSNorm.scale)
         get crushed to FP8 and quality degrades noticeably.
 
-        `extra_skip_patterns` carries the model's own declared exclusions (diffusers'
-        `_skip_layerwise_casting_patterns`), which are model-specific and cannot be inferred from
+        `extra_skip_patterns` carries the model's own declared exclusions (see
+        `_model_declared_skip_patterns`), which are model-specific and cannot be inferred from
         layer types or generic name patterns.
 
         Modules holding already-quantized weights are skipped regardless of their class. This is a

invokeai/backend/model_manager/load/model_loaders/krea2.py

@@ -370,6 +370,10 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel:
 
         model.load_state_dict(sd, assign=True, strict=False)
         _reject_incomplete_load(model, what="Krea-2 single-file checkpoint")
+        # `assign=True` aliases every param to its `sd` tensor. Drop the dict's references before
+        # the FP8 cast, or each param's `model_dtype` original stays reachable while its fp8 copy is
+        # allocated, overshooting the `make_room()` reservation above by ~50%.
+        sd.clear()
         # Honor the fp8-storage setting (re-quantizes the dequantized weights to fp8-resident on CUDA).
         model = self._apply_fp8_layerwise_casting(model, config, SubModelType.Transformer)
         return model

invokeai/backend/model_manager/load/model_loaders/z_image.py

@@ -482,10 +482,20 @@ def _load_from_singlefile(
             sd[k] = sd[k].to(model_dtype)
 
         model.load_state_dict(sd, assign=True)
-
-        # Every param is uniform `model_dtype` at this point (the loop above casts the whole state
-        # dict, including ComfyUI fp8 checkpoints, whose scale metadata was filtered out above), so
-        # the layerwise cast has a single unambiguous compute dtype to restore to.
+        # `assign=True` aliases every param to its `sd` tensor, so the dict keeps the whole model
+        # alive a second time. The FP8 cast below allocates the fp8 copy per param while the
+        # `model_dtype` original is still reachable through `sd`, pushing peak RAM to ~1.5x what
+        # `make_room()` reserved above (~17.4GB actual vs ~11.5GB reserved for Z-Image). Dropping
+        # the dict's references lets each original free as soon as its param is cast.
+        sd.clear()
+
+        # Every param is uniform `model_dtype` at this point, so the layerwise cast has a single
+        # unambiguous compute dtype to restore to.
+        #
+        # Caveat, pre-existing and not addressed here: for a ComfyUI *scaled*-fp8 checkpoint the
+        # filter above drops `.scale_weight` / `scaled_fp8` without folding them in, so the raw fp8
+        # codes are cast to `model_dtype` unscaled and the model loads with wrong weights. That is a
+        # separate bug in the key filtering, not something this cast makes safe.
         model = self._apply_fp8_layerwise_casting(model, config, SubModelType.Transformer)
         return model
 

tests/backend/model_manager/load/test_load_default_fp8.py

@@ -26,6 +26,7 @@
     _QUANTIZED_MODEL_FORMATS,
     ModelLoader,
     _device_supports_fp8_storage,
+    _model_declared_skip_patterns,
 )
 from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_linear import (
     CustomLinear,
@@ -692,3 +693,58 @@ def test_device_supports_fp8_storage_is_cached_per_device():
     # xpu:1 has not been probed, so a failing probe there must be observed, not short-circuited.
     with patch("torch.zeros", side_effect=RuntimeError("no float8 on this device")):
         assert _device_supports_fp8_storage(torch.device("xpu", 1)) is False
+
+
+def test_model_declared_skip_patterns_unions_both_diffusers_attributes():
+    """`enable_layerwise_casting()` unions `_skip_layerwise_casting_patterns` with
+    `_keep_in_fp32_modules`. We replaced that call with our own hook-based path, so we have to read
+    both — otherwise a model that declares only the latter loses its exclusions silently.
+    """
+
+    class _Model(torch.nn.Module):
+        _skip_layerwise_casting_patterns = ["t_embedder"]
+        _keep_in_fp32_modules = ["time_embedder"]
+
+    assert _model_declared_skip_patterns(_Model()) == ("t_embedder", "time_embedder")
+
+
+def test_model_declared_skip_patterns_tolerates_missing_and_odd_declarations():
+    """Most models declare neither attribute; diffusers sets them to `None` on some. A bare string
+    is accepted too, so a subclass that writes one instead of a list isn't silently expanded into
+    per-character patterns."""
+
+    class _Bare(torch.nn.Module):
+        pass
+
+    class _Nulls(torch.nn.Module):
+        _skip_layerwise_casting_patterns = None
+        _keep_in_fp32_modules = None
+
+    class _Strings(torch.nn.Module):
+        _keep_in_fp32_modules = "time_embedder"
+
+    assert _model_declared_skip_patterns(_Bare()) == ()
+    assert _model_declared_skip_patterns(_Nulls()) == ()
+    assert _model_declared_skip_patterns(_Strings()) == ("time_embedder",)
+
+
+def test_keep_in_fp32_modules_are_not_cast():
+    """End-to-end through the cast: a module named only by `_keep_in_fp32_modules` keeps its
+    compute dtype."""
+
+    class _Model(torch.nn.Module):
+        _keep_in_fp32_modules = ["time_embedder"]
+
+        def __init__(self):
+            super().__init__()
+            self.time_embedder = torch.nn.Linear(4, 4)
+            self.attn = torch.nn.Linear(4, 4)
+
+    loader = _make_loader(device="cuda")
+    model = _Model().to(torch.bfloat16)
+
+    with patch.object(ModelLoader, "_should_use_fp8", return_value=True):
+        loader._apply_fp8_layerwise_casting(model, _make_config(ModelType.Main, fp8=True))
+
+    assert model.time_embedder.weight.dtype == torch.bfloat16
+    assert model.attn.weight.dtype == torch.float8_e4m3fn

tests/backend/model_manager/load/test_z_image_fp8_wiring.py

@@ -0,0 +1,106 @@
+"""Guards for Z-Image's FP8-storage wiring in the single-file loader.
+
+`_load_from_singlefile` is the only place Z-Image checkpoints reach the layerwise cast, and nothing
+else in the suite observes it: deleting the `_apply_fp8_layerwise_casting` call leaves the whole
+`tests/backend/model_manager` suite green, so the toggle can go back to being rendered-and-inert
+with CI passing.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import diffusers
+import safetensors.torch
+import torch
+
+from invokeai.backend.model_manager.configs.main import Main_Checkpoint_ZImage_Config
+from invokeai.backend.model_manager.load.model_loaders.z_image import ZImageCheckpointModel
+from invokeai.backend.model_manager.taxonomy import SubModelType, ZImageVariantType
+
+
+class _TinyZImageTransformer(torch.nn.Module):
+    def __init__(self, **_kwargs) -> None:
+        super().__init__()
+        # `x_pad_token` is one of the loader's `valid_exact` keys, so it survives key filtering.
+        self.x_pad_token = torch.nn.Parameter(torch.empty(2, 2))
+
+
+def _prepare_loader(monkeypatch, tmp_path, state_dict: dict[str, torch.Tensor]):
+    checkpoint_path = tmp_path / "z_image.safetensors"
+    checkpoint_path.touch()
+    config = Main_Checkpoint_ZImage_Config.model_construct(
+        path=str(checkpoint_path), variant=ZImageVariantType.Turbo, fp8_storage=True
+    )
+
+    monkeypatch.setattr(diffusers, "ZImageTransformer2DModel", _TinyZImageTransformer, raising=False)
+    monkeypatch.setattr(safetensors.torch, "load_file", lambda _path: state_dict)
+    monkeypatch.setattr(
+        "invokeai.backend.model_manager.load.model_loaders.z_image.TorchDevice.choose_torch_device",
+        lambda: torch.device("cpu"),
+    )
+    monkeypatch.setattr(
+        "invokeai.backend.model_manager.load.model_loaders.z_image.TorchDevice.choose_bfloat16_safe_dtype",
+        lambda _device: torch.float32,
+    )
+
+    loader = object.__new__(ZImageCheckpointModel)
+    loader._ram_cache = SimpleNamespace(make_room=MagicMock())
+    return loader, config
+
+
+def test_single_file_loader_applies_fp8_layerwise_casting(monkeypatch, tmp_path) -> None:
+    """The `fp8_storage` toggle has to reach the cast for Z-Image checkpoints at all."""
+    cast_calls: list[tuple[object, object]] = []
+
+    loader, config = _prepare_loader(monkeypatch, tmp_path, {"x_pad_token": torch.ones(2, 2)})
+    loader._apply_fp8_layerwise_casting = lambda model, cfg, submodel: (
+        cast_calls.append((cfg, submodel)),
+        model,
+    )[1]
+
+    model = loader._load_from_singlefile(config)
+
+    assert isinstance(model, _TinyZImageTransformer)
+    assert torch.equal(model.x_pad_token, torch.ones(2, 2))
+    assert cast_calls == [(config, SubModelType.Transformer)]
+
+
+def test_state_dict_is_released_before_the_fp8_cast(monkeypatch, tmp_path) -> None:
+    """Peak RAM must not overshoot the `make_room()` reservation.
+
+    `load_state_dict(..., assign=True)` aliases every param to its state-dict tensor. If the dict is
+    still holding those references when the FP8 cast runs, `param.data.to(float8)` allocates the fp8
+    copy while the original is still reachable, so the model is briefly resident ~1.5x over — about
+    17.4GB actual against an ~11.5GB reservation for Z-Image. Nothing in the loader reads `sd` after
+    the load, so the dict must be empty by the time the cast starts.
+    """
+    state_dict = {"x_pad_token": torch.ones(2, 2)}
+    observed_sd_len: list[int] = []
+
+    loader, config = _prepare_loader(monkeypatch, tmp_path, state_dict)
+    loader._apply_fp8_layerwise_casting = lambda model, _cfg, _submodel: (
+        observed_sd_len.append(len(state_dict)),
+        model,
+    )[1]
+
+    loader._load_from_singlefile(config)
+
+    assert observed_sd_len == [0]
+
+
+def test_assign_true_really_aliases_the_state_dict() -> None:
+    """The premise of the test above: without clearing, the originals stay alive through `sd`.
+
+    If a future torch release stopped aliasing under `assign=True`, `sd.clear()` would become dead
+    weight rather than a fix, and this is the test that says so.
+    """
+    model = _TinyZImageTransformer()
+    sd = {"x_pad_token": torch.ones(2, 2)}
+
+    model.load_state_dict(sd, assign=True)
+
+    assert model.x_pad_token.data_ptr() == sd["x_pad_token"].data_ptr()
+
+    # And the cast leaves the state dict's copy behind at the original dtype.
+    model.x_pad_token.data = model.x_pad_token.data.to(torch.float16)
+    assert sd["x_pad_token"].dtype == torch.float32