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

@@ -156,6 +156,35 @@ def _device_supports_fp8_storage(device: torch.device, logger: Optional[Logger]
     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. Must list every
+# quantized member of `ModelFormat`; `test_quantized_format_set_matches_the_taxonomy` pins the
+# strings to the enum so a rename cannot silently disable the check.
+_QUANTIZED_MODEL_FORMATS: frozenset[str] = frozenset(
+    {
+        "gguf_quantized",
+        "bnb_quantized_nf4b",
+        "bnb_quantized_int8b",
+        "sdnq_quantized",
+    }
+)
+
+
+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
@@ -346,6 +375,18 @@ def _should_use_fp8(self, config: AnyModelConfig, submodel_type: Optional[SubMod
         """Check if FP8 layerwise casting should be applied to a model."""
         from invokeai.backend.model_manager.taxonomy import ModelType
 
+        # Already-quantized models are excluded. Their weights are packed integer payloads, not
+        # values we may re-encode, and casting them is not a no-op:
+        #   - GGUF raises `Operation changed the dtype of GGMLTensor unexpectedly`.
+        #   - 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.
+        # No quantized-format loader calls `_apply_fp8_layerwise_casting` today, so this is a guard
+        # against the next loader that gets wired up (they are being added one model at a time)
+        # rather than a fix for a live crash.
+        if hasattr(config, "format") and config.format in _QUANTIZED_MODEL_FORMATS:
+            return False
+
         # VAEs are excluded — fp8 storage causes noticeable quality degradation in decode.
         if hasattr(config, "type") and config.type == ModelType.VAE:
             return False
@@ -468,6 +509,11 @@ def _apply_fp8_to_nn_module(
         `_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).
+
         Records the compute dtype on the model. After the cast, `model.dtype` reports the float8
         storage dtype, which must never be used to create or cast tensors — torch has no arithmetic
         kernels for it (see `get_model_compute_dtype`). The marker is set here rather than at the
@@ -484,6 +530,8 @@ def _apply_fp8_to_nn_module(
             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/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/MainModelDefaultSettings/MainModelDefaultSettings.tsx

@@ -56,6 +56,15 @@ export const MainModelDefaultSettings = memo(({ modelConfig }: Props) => {
     return ['flux', 'flux2'].includes(modelConfig.base);
   }, [modelConfig]);
 
+  // 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.
+  // Keep in sync with `_QUANTIZED_MODEL_FORMATS` in `load_default.py`.
+  const isQuantized = useMemo(() => {
+    return ['gguf_quantized', 'bnb_quantized_nf4b', 'bnb_quantized_int8b', 'sdnq_quantized'].includes(
+      modelConfig.format
+    );
+  }, [modelConfig]);
+
   const defaultSettingsDefaults = useMainModelDefaultSettings(modelConfig);
   const optimalDimension = useMemo(() => {
     const modelBase = modelConfig?.base;
@@ -144,7 +153,7 @@ export const MainModelDefaultSettings = memo(({ modelConfig }: Props) => {
         {!isFluxFamily && <DefaultCfgRescaleMultiplier control={control} name="cfgRescaleMultiplier" />}
         <DefaultWidth control={control} optimalDimension={optimalDimension} />
         <DefaultHeight control={control} optimalDimension={optimalDimension} />
-        <DefaultFp8Storage control={control} name="fp8Storage" />
+        {!isQuantized && <DefaultFp8Storage control={control} name="fp8Storage" />}
       </SimpleGrid>
     </>
   );

tests/backend/model_manager/load/test_load_default_fp8.py

@@ -23,6 +23,7 @@
 from invokeai.backend.model_manager.load.load_default import (
     _FP8_PROBE_FAILURE_REPORTED,
     _FP8_STORAGE_SUPPORTED,
+    _QUANTIZED_MODEL_FORMATS,
     ModelLoader,
     _device_supports_fp8_storage,
 )
@@ -32,7 +33,7 @@
 from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import (
     apply_custom_layers_to_model,
 )
-from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType, SubModelType
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType
 
 
 def _make_loader(device: str = "cuda") -> ModelLoader:
@@ -57,13 +58,22 @@ def _make_config(model_type: ModelType, fp8: bool, base: BaseModelType = BaseMod
     )
 
 
+def _make_quantized_config(fmt: ModelFormat = ModelFormat.GGUFQuantized):
+    """A config carrying a quantized `format`, which `_make_config` deliberately omits."""
+    config = _make_config(ModelType.Main, fp8=True)
+    config.format = fmt
+    return config
+
+
 @pytest.mark.parametrize(
     "config,submodel",
     [
         (_make_config(ModelType.VAE, fp8=True), None),
         (_make_config(ModelType.LoRA, fp8=True), None),
         # Z-Image used to be listed here. It is no longer excluded — see
         # `test_should_use_fp8_allows_z_image` for why the exclusion became obsolete.
+        # A quantized model takes its place: its guard must also sit ahead of the device probe.
+        (_make_quantized_config(), None),
         (_make_config(ModelType.Main, fp8=True), SubModelType.Tokenizer),
         (_make_config(ModelType.Main, fp8=False), None),
     ],
@@ -371,6 +381,107 @@ def __init__(self):
     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",
+    [
+        ModelFormat.GGUFQuantized,
+        ModelFormat.BnbQuantizednf4b,
+        ModelFormat.BnbQuantizedLlmInt8b,
+        ModelFormat.SDNQQuantized,
+    ],
+)
+def test_should_use_fp8_excludes_quantized_formats(fmt: ModelFormat):
+    """Already-quantized weights must never be re-encoded as FP8.
+
+    Casting them is not a no-op: GGUF raises `Operation changed the dtype of GGMLTensor
+    unexpectedly`, 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).
+
+    Parametrized over `ModelFormat` members rather than raw strings: `_QUANTIZED_MODEL_FORMATS`
+    holds strings, so testing it with strings would pass even if the enum values drifted.
+    """
+    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_quantized_format_set_matches_the_taxonomy():
+    """Every entry in `_QUANTIZED_MODEL_FORMATS` must still name a real `ModelFormat` value.
+
+    The set is declared as raw strings to keep `load_default` free of a taxonomy import at module
+    scope, so nothing else stops a rename in `ModelFormat` from silently disabling the check —
+    `config.format` would simply never match again, and FP8 would be re-enabled for that format.
+    """
+    assert _QUANTIZED_MODEL_FORMATS <= {fmt.value for fmt in ModelFormat}
+
+
+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