fix(fp8): never apply FP8 storage to already-quantized weights

invokeai/backend/anima/anima_transformer.py

@@ -985,6 +985,27 @@ class AnimaTransformer(MiniTrainDIT):
     text embeddings before they are fed to the DiT cross-attention layers.
     """
 
+    # Modules that must keep their compute dtype when FP8 storage is enabled. Read by
+    # `ModelLoader._apply_fp8_layerwise_casting`, which uses the same attribute name diffusers
+    # models use, so no loader-side special-casing is needed.
+    #
+    # The generic skip patterns don't reach these: they match `norm`, `pos_embed`, `patch_embed`
+    # and `proj_in/out`, but this architecture names the equivalent modules differently.
+    #
+    # `t_embedder` is the one that matters, and it is not a rounding-quality nicety: with it cast
+    # to FP8 the model renders a heavily dithered image with no fine detail at all (verified
+    # against a bf16 run at the same seed/steps/CFG). It feeds `adaln_lora` into every block, so
+    # its error is applied to every token everywhere. Measured, same seed each time: casting
+    # nothing = broken; `t_embedder` alone = clean, and adding either of the two below changes
+    # nothing further. They are kept as ~2MB of margin on the I/O layers, matching what diffusers
+    # skips by default for comparable DiTs. `adaln_modulation` was also tested and is deliberately
+    # NOT listed — it costs 168MB and made no difference.
+    _skip_layerwise_casting_patterns = [
+        "t_embedder",  # timestep embedding MLP -> adaln_lora for every block
+        "x_embedder",  # patch embedding (named `patch_embed` in diffusers models)
+        "final_layer",  # output projection
+    ]
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.llm_adapter = LLMAdapter()

invokeai/backend/model_manager/load/load_default.py

@@ -58,6 +58,32 @@
     r"^proj_out$",
 )
 
+# Model formats whose weights are already quantized. FP8 storage is meaningless for them (the
+# payload is packed integers, not values we may re-encode) and actively harmful — see
+# `_should_use_fp8`. Declared as strings to keep this module free of a taxonomy import at module
+# scope; compared against `config.format`, which is a `ModelFormat` str-enum.
+_QUANTIZED_MODEL_FORMATS: frozenset[str] = frozenset(
+    {
+        "gguf_quantized",
+        "bnb_quantized_nf4b",
+        "bnb_quantized_int8b",
+    }
+)
+
+
+def _is_quantized_param(param: torch.nn.Parameter) -> bool:
+    """Whether `param` holds a quantized payload that must not be re-encoded as FP8.
+
+    Two signals, both observed in practice:
+
+    - Not floating point. bnb's NF4/INT8 weights are packed `uint8` (and `bnb.nn.LinearNF4`
+      subclasses `nn.Linear`, so a class check alone does not catch them). Casting those to float8
+      succeeds silently and the layer then returns finite garbage.
+    - A `torch.Tensor` *subclass*, e.g. `GGMLTensor`, which keeps its quantized payload plus
+      metadata and rejects dtype changes outright.
+    """
+    return not param.data.is_floating_point() or type(param.data) is not torch.Tensor
+
 
 # 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
@@ -234,11 +260,16 @@ def _should_use_fp8(self, config: AnyModelConfig, submodel_type: Optional[SubMod
         if self._torch_device.type != "cuda":
             return False
 
-        # Z-Image has dtype mismatch issues with diffusers' layerwise casting
-        # (skipped modules produce bf16, hooked modules expect fp16).
-        from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType
+        from invokeai.backend.model_manager.taxonomy import ModelType
 
-        if hasattr(config, "base") and config.base == BaseModelType.ZImage:
+        # Already-quantized models are excluded. Their weights are packed integer payloads, not
+        # values we may re-encode, and every quantized-format loader reaches this helper. Casting
+        # them is not a no-op:
+        #   - GGUF raises `Operation changed the dtype of GGMLTensor unexpectedly` at load time.
+        #   - bnb NF4 corrupts *silently* — `bnb.nn.LinearNF4` subclasses `nn.Linear`, so the packed
+        #     uint8 payload is cast to float8, inference still returns finite numbers, and the model
+        #     just produces garbage.
+        if hasattr(config, "format") and config.format in _QUANTIZED_MODEL_FORMATS:
             return False
 
         # VAEs are excluded — fp8 storage causes noticeable quality degradation in decode.
@@ -310,7 +341,19 @@ 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):
-            self._apply_fp8_to_nn_module(model, storage_dtype=storage_dtype, compute_dtype=compute_dtype)
+            # 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 ()),
+            )
         else:
             return model
 
@@ -323,23 +366,40 @@ def _apply_fp8_layerwise_casting(
         return model
 
     @staticmethod
-    def _apply_fp8_to_nn_module(model: torch.nn.Module, storage_dtype: torch.dtype, compute_dtype: torch.dtype) -> None:
+    def _apply_fp8_to_nn_module(
+        model: torch.nn.Module,
+        storage_dtype: torch.dtype,
+        compute_dtype: torch.dtype,
+        extra_skip_patterns: tuple[str, ...] = (),
+    ) -> None:
         """Apply FP8 layerwise casting to a plain nn.Module.
 
         Mirrors diffusers' `apply_layerwise_casting` semantics: only the layer classes in
         `_FP8_SUPPORTED_PYTORCH_LAYERS` are cast, and modules whose dotted path matches any of
         `_FP8_DEFAULT_SKIP_PATTERNS` (norm, pos_embed, patch_embed, proj_in/out) are skipped.
         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
+        layer types or generic name patterns.
+
+        Modules holding already-quantized weights are skipped regardless of their class. This is a
+        backstop behind the format check in `_should_use_fp8`, which cannot see quantization that
+        is not reflected in the model's format (e.g. a `diffusers`-format checkpoint whose weights
+        were quantized by an external tool).
         """
+        skip_patterns = _FP8_DEFAULT_SKIP_PATTERNS + tuple(extra_skip_patterns)
         for module_name, module in model.named_modules():
             if not isinstance(module, _FP8_SUPPORTED_PYTORCH_LAYERS):
                 continue
-            if any(re.search(pattern, module_name) for pattern in _FP8_DEFAULT_SKIP_PATTERNS):
+            if any(re.search(pattern, module_name) for pattern in skip_patterns):
                 continue
             params = list(module.parameters(recurse=False))
             if not params:
                 continue
+            if any(_is_quantized_param(p) for p in params):
+                continue
 
             for param in params:
                 param.data = param.data.to(storage_dtype)

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

@@ -181,6 +181,12 @@ def _load_from_singlefile(
                 f"Checkpoint is missing {len(load_result.missing_keys)} keys "
                 f"(expected for inv_freq buffers). First 5: {load_result.missing_keys[:5]}"
             )
+
+        # Without this the `fp8_storage` toggle is shown for Anima models but does nothing. The
+        # state dict was cast to a single `model_dtype` above, so the layerwise cast has one
+        # unambiguous compute dtype to restore to. AnimaTransformer is a plain nn.Module, so this
+        # takes the hook-based path in `_apply_fp8_to_nn_module`.
+        model = self._apply_fp8_layerwise_casting(model, config, SubModelType.Transformer)
         return model
 
 

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

@@ -289,6 +289,11 @@ 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.
+        model = self._apply_fp8_layerwise_casting(model, config, SubModelType.Transformer)
         return model
 
 

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings.tsx

@@ -56,8 +56,10 @@ export const MainModelDefaultSettings = memo(({ modelConfig }: Props) => {
     return ['flux', 'flux2'].includes(modelConfig.base);
   }, [modelConfig]);
 
-  const isZImage = useMemo(() => {
-    return modelConfig.base === 'z-image';
+  // Already-quantized weights cannot also be stored as FP8 — the backend refuses it (see
+  // `_should_use_fp8`), so offering the switch would be a control that silently does nothing.
+  const isQuantized = useMemo(() => {
+    return ['gguf_quantized', 'bnb_quantized_nf4b', 'bnb_quantized_int8b'].includes(modelConfig.format);
   }, [modelConfig]);
 
   const defaultSettingsDefaults = useMainModelDefaultSettings(modelConfig);
@@ -148,7 +150,7 @@ export const MainModelDefaultSettings = memo(({ modelConfig }: Props) => {
         {!isFluxFamily && <DefaultCfgRescaleMultiplier control={control} name="cfgRescaleMultiplier" />}
         <DefaultWidth control={control} optimalDimension={optimalDimension} />
         <DefaultHeight control={control} optimalDimension={optimalDimension} />
-        {!isZImage && <DefaultFp8Storage control={control} name="fp8Storage" />}
+        {!isQuantized && <DefaultFp8Storage control={control} name="fp8Storage" />}
       </SimpleGrid>
     </>
   );

tests/backend/model_manager/load/test_load_default_fp8.py

@@ -272,6 +272,153 @@ def __init__(self):
     assert model.rms.scale.dtype == compute_dtype
 
 
+def test_apply_fp8_to_nn_module_honors_extra_skip_patterns():
+    """A model's own `_skip_layerwise_casting_patterns` must be applied on top of our defaults."""
+
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.t_embedder = torch.nn.Linear(4, 4)
+            self.attn = torch.nn.Linear(4, 4)
+
+    storage_dtype = torch.float16
+    compute_dtype = torch.float32
+    model = _Model()
+    for p in model.parameters():
+        p.data = p.data.to(compute_dtype)
+
+    ModelLoader._apply_fp8_to_nn_module(
+        model, storage_dtype, compute_dtype, extra_skip_patterns=("t_embedder", "cap_embedder")
+    )
+
+    assert model.attn.weight.dtype == storage_dtype
+    assert model.t_embedder.weight.dtype == compute_dtype
+
+
+def test_apply_fp8_layerwise_casting_passes_model_declared_skip_patterns():
+    """Regression test for Z-Image + fp8 crashing with
+    `RuntimeError: "addmm_cuda" not implemented for 'Float8_e4m3fn'`.
+
+    Diffusers models declare precision-sensitive modules in `_skip_layerwise_casting_patterns`, and
+    `enable_layerwise_casting()` honors them. Our hook-based replacement must read that list too —
+    it is not redundant with `_FP8_DEFAULT_SKIP_PATTERNS`. `ZImageTransformer2DModel` declares
+    `['t_embedder', 'cap_embedder']` because `TimestepEmbedder.forward` reads
+    `self.mlp[0].weight.dtype` and casts its *input* to it: with an fp8 weight the input becomes
+    float8 before the pre-hook restores the weight, and `F.linear` has no float8 kernel.
+    """
+
+    class _FakeZImage(torch.nn.Module):
+        _skip_layerwise_casting_patterns = ["t_embedder", "cap_embedder"]
+
+        def __init__(self):
+            super().__init__()
+            self.t_embedder = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.Linear(4, 4))
+            self.cap_embedder = torch.nn.Linear(4, 4)
+            self.layers = torch.nn.Linear(4, 4)
+
+    loader = _make_loader(device="cuda")
+    model = _FakeZImage().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, base=BaseModelType.ZImage))
+
+    # The declared modules keep their compute dtype...
+    assert model.t_embedder[0].weight.dtype == torch.bfloat16
+    assert model.cap_embedder.weight.dtype == torch.bfloat16
+    # ...while everything else is stored in fp8, so the toggle still saves VRAM.
+    assert model.layers.weight.dtype == torch.float8_e4m3fn
+
+
+def test_anima_transformer_declares_t_embedder_skip():
+    """Regression guard for Anima + FP8 rendering a heavily dithered image.
+
+    `AnimaTransformer.t_embedder` produces the `adaln_lora` conditioning consumed by every block,
+    so casting it to FP8 corrupts every token of every block — verified against a bf16 run at the
+    same seed/steps/CFG. None of the generic `_FP8_DEFAULT_SKIP_PATTERNS` match it (this
+    architecture doesn't use diffusers' module names), so the model has to declare it itself.
+    """
+    from invokeai.backend.anima.anima_transformer import AnimaTransformer
+
+    assert "t_embedder" in AnimaTransformer._skip_layerwise_casting_patterns
+
+    # And the declared patterns actually reach the cast, matched against dotted module paths.
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.t_embedder = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.Linear(4, 4))
+            self.blocks = torch.nn.Linear(4, 4)
+
+    model = _Model().to(torch.float32)
+    ModelLoader._apply_fp8_to_nn_module(
+        model,
+        storage_dtype=torch.float16,
+        compute_dtype=torch.float32,
+        extra_skip_patterns=tuple(AnimaTransformer._skip_layerwise_casting_patterns),
+    )
+
+    assert model.t_embedder[0].weight.dtype == torch.float32
+    assert model.blocks.weight.dtype == torch.float16
+
+
+@pytest.mark.parametrize("fmt", ["gguf_quantized", "bnb_quantized_nf4b", "bnb_quantized_int8b"])
+def test_should_use_fp8_excludes_quantized_formats(fmt: str):
+    """Already-quantized weights must never be re-encoded as FP8.
+
+    Every quantized-format loader reaches `_apply_fp8_layerwise_casting`, and casting there is not
+    a no-op: GGUF raises `Operation changed the dtype of GGMLTensor unexpectedly` at load time, and
+    bnb NF4 corrupts silently (`bnb.nn.LinearNF4` subclasses `nn.Linear`, so its packed uint8
+    payload is cast to float8 and inference then returns finite garbage).
+    """
+    loader = _make_loader(device="cuda")
+    config = _make_config(ModelType.Main, fp8=True)
+    config.format = fmt
+    assert loader._should_use_fp8(config) is False
+
+
+def test_apply_fp8_skips_quantized_params_regardless_of_format():
+    """Backstop behind the format check, for quantization the model's format does not reveal
+    (e.g. a `diffusers`-format checkpoint quantized by an external tool).
+
+    Both signals are covered: a non-floating-point payload (bnb's packed uint8) and a
+    `torch.Tensor` subclass (GGUF's `GGMLTensor`).
+    """
+
+    class _FakeQuantTensor(torch.Tensor):
+        """Stands in for GGMLTensor: a Tensor subclass carrying a quantized payload."""
+
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.packed = torch.nn.Linear(4, 4, bias=False)  # bnb-style uint8 payload
+            self.subclassed = torch.nn.Linear(4, 4, bias=False)  # GGUF-style tensor subclass
+            # NB: not `normal` — that would be caught by the `norm` skip pattern.
+            self.attn = torch.nn.Linear(4, 4, bias=False)
+
+    model = _Model().to(torch.bfloat16)
+    model.packed.weight = torch.nn.Parameter(torch.zeros(8, 1, dtype=torch.uint8), requires_grad=False)
+    model.subclassed.weight = torch.nn.Parameter(
+        torch.zeros(4, 4, dtype=torch.bfloat16).as_subclass(_FakeQuantTensor), requires_grad=False
+    )
+
+    ModelLoader._apply_fp8_to_nn_module(model, torch.float8_e4m3fn, torch.bfloat16)
+
+    assert model.packed.weight.dtype == torch.uint8
+    assert not model.packed._forward_pre_hooks, "a quantized layer must not get cast hooks either"
+    assert model.subclassed.weight.dtype == torch.bfloat16
+    assert not model.subclassed._forward_pre_hooks
+    # Control: an ordinary layer in the same model is still cast.
+    assert model.attn.weight.dtype == torch.float8_e4m3fn
+
+
+def test_should_use_fp8_allows_z_image():
+    """Z-Image was excluded while we used diffusers' `enable_layerwise_casting()` with the global
+    torch dtype (fp16) as compute dtype, which clashed with the model's bf16 weights. The compute
+    dtype now comes from the model itself, so the exclusion is obsolete.
+    """
+    loader = _make_loader(device="cuda")
+    assert loader._should_use_fp8(_make_config(ModelType.Main, fp8=True, base=BaseModelType.ZImage)) is True
+
+
 def test_wrap_forward_reaches_custom_linear_after_apply_custom_layers():
     """Production order: `_load_model` applies FP8 wrapping, THEN `ModelCache.put()` calls
     `apply_custom_layers_to_model` which constructs a NEW `CustomLinear` object via