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

[Pfannkuchensack]

Summary

Stacked on #9415. #9414 is now merged into main, so this branch carries #9415's Anima
changes plus its own. #9478 stacks on top. No logical dependency — the guard stands alone — but
all of them touch _should_use_fp8, _apply_fp8_to_nn_module, MainModelDefaultSettings.tsx
and the same test file, so merging in order avoids conflicts.

FP8 storage must never be applied to weights that are already quantized. The cast is not a no-op
on such a payload:

• GGUF raises ValueError: Operation changed the dtype of GGMLTensor unexpectedly.
• 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 — the model just produces garbage.

You can see the silent one directly, no model download needed (needs CUDA):

import torch, bitsandbytes as bnb
from invokeai.backend.model_manager.load.load_default import ModelLoader

l4 = bnb.nn.LinearNF4(256, 256, bias=False)
l4.weight = bnb.nn.Params4bit(torch.randn(256, 256), requires_grad=False, quant_type="nf4")
l4 = l4.cuda()
x = torch.randn(1, 256, device="cuda", dtype=torch.bfloat16)
ref = l4(x).clone()

ModelLoader._apply_fp8_to_nn_module(l4, torch.float8_e4m3fn, torch.bfloat16)
print("max abs diff:", (l4(x).float() - ref.float()).abs().max().item())
# before this PR: ~50   (no error, no warning)
# after:           0.0

This is defense-in-depth, not a live crash fix

An earlier revision of this description claimed that "every quantized-format loader reaches
_apply_fp8_layerwise_casting". That is not true, and the QA steps that followed from it did
not reproduce. Walking ModelLoaderRegistry and checking each class's MRO down to
ModelLoader shows that none of the 17 quantized-format loaders calls the helper —
FluxGGUFCheckpointModel, FluxBnbQuantizednf4bCheckpointModel, Flux2GGUFCheckpointModel,
QwenImageGGUFCheckpointModel, ZImageGGUFCheckpointModel, Krea2GGUFCheckpointModel,
WanGGUFCheckpointModel, the three SDNQ main loaders, and the T5 / Qwen3 / Gemma2 / Mistral
encoder loaders all build and return the model without it. (qwen_image.py does call it, but
from QwenImageCheckpointModel, not the GGUF one.) The same is true on main.

So enabling fp8_storage on a GGUF or bnb model today does nothing rather than breaking it. What
this PR fixes is that the door is standing open: _should_use_fp8 answers "yes" for these models,
the UI offers the toggle for them, and FP8 support is being wired into loaders one model at a time
(#9414 Z-Image, #9415 Anima). The next loader to be wired up would hit the crash / the silent
corruption above, with nothing in the way.

The guard

Two levels, because a format check alone is not enough — externally quantized weights can ship
under a plain diffusers format:

• _should_use_fp8 rejects gguf_quantized, bnb_quantized_nf4b, bnb_quantized_int8b and
  sdnq_quantized — every quantized member of ModelFormat.
• _apply_fp8_to_nn_module skips any module whose params are non-floating-point (bnb's packed
  uint8) or a torch.Tensor subclass (GGMLTensor, SDNQTensor), regardless of the declared
  format. This also covers a model that already carries a persisted fp8_storage=true from
  before this PR.

Frontend hides the toggle for the same four formats, so the UI stops offering a control the
backend refuses.

Known gap (follow-up, not addressed here)

_apply_fp8_layerwise_casting derives the compute dtype from next(model.parameters()).dtype
before the param-level check runs. For the case the backstop exists for — a plain-format
checkpoint with externally quantized weights — a quantized first parameter makes that
torch.uint8, so get_model_compute_dtype() returns uint8 and the pre-hook on the layers that
were cast does p.data.to(uint8). The quantized layers are correctly skipped, but the model is
still broken. The fix is to derive the compute dtype from the first floating-point parameter;
it is not reachable today (it needs a quantized model on a cast-calling loader) so it is left for
a follow-up rather than widening this PR.

Related Issues / Discussions

Follow-up to #8945 (FP8 storage). Same series as #9414 (Z-Image), #9415 (Anima) and #9478.

QA Instructions

The behaviour change is not observable through the UI for a GGUF/bnb/SDNQ model beyond the toggle
disappearing, because the cast never ran for those models in the first place. What is worth
checking:

1. UI — open a GGUF, bnb or SDNQ main model in Model Manager → Default Settings. The FP8
   Storage switch is gone. Open a non-quantized model (SDXL, FLUX checkpoint, Z-Image diffusers)
   → the switch is still there and still works.

2. No regression on non-quantized models — a model with FP8 enabled must still log
   FP8 layerwise casting enabled ... and load at roughly half its usual VRAM.

3. Legacy DB value — the backend must not rely on the UI. Force the old value back in:

   curl -X PATCH http://127.0.0.1:9090/api/v2/models/i/<gguf-model-key> \
     -H "Content-Type: application/json" \
     -d '{"default_settings":{"fp8_storage":true}}'
   

   Then generate with that model. Expect: model loads at its normal size, no
   FP8 layerwise casting enabled line in the log.

4. The direct repro above — the one path that does exercise the corruption. Needs CUDA and
   bitsandbytes.

Unit tests:

uv run --extra cuda --extra test pytest tests/backend/model_manager/load -q --no-cov

The new tests cover the format check (parametrized over all four quantized ModelFormat members,
not raw strings, so a drift in the enum values fails), a taxonomy check that pins every entry in
_QUANTIZED_MODEL_FORMATS to a real ModelFormat value, and the param-level guard (both signals:
non-floating-point payload and torch.Tensor subclass), plus a control assertion that an ordinary
layer in the same model is still cast.

Merge Plan

Merge after #9415 to avoid conflicts in load_default.py, MainModelDefaultSettings.tsx and
test_load_default_fp8.py. The guard itself has no logical dependency on the earlier PRs — it is
a self-contained change — but this branch carries #9415 merged in, so it is not a drop-in
replacement for it.

No DB schema, no redux slice, no API schema change. Existing fp8_storage=true values persisted
against quantized models stay in the DB and are simply ignored from now on; no migration needed.

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

[joshistoast]

This would allow me to run krea for once, please oh please put this in 6.14.0 👀 my dear @lstein

[lstein]

Adversarial review at d8b984162d, then merged main and pushed two follow-ups — now at b187ae46c2. Approving.

What I changed on the branch

1. Resolved the merge conflicts. #9414 landed on main as a squash (b06e57e63f) while this branch carried it as a merge of the original commits, so all three shared files conflicted. Every hunk was "ours adds, theirs is unchanged" — resolved by keeping the branch side and letting the rest auto-merge. The resulting tree against main is exactly #9415's Anima changes plus this PR's guard; nothing from #9414 is duplicated or lost.

2. Added sdnq_quantized to both lists (b187ae46c2). ModelFormat.SDNQQuantized was missing from _QUANTIZED_MODEL_FORMATS and from the frontend's isQuantized, even though it is first-class: 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 said yes for SDNQ main models and Model Manager still rendered the switch for them — the same dead control this PR removes for GGUF and bnb. Mutation-verified: dropping the entry fails the new [sdnq_quantized] case.

Alongside it, the format test now parametrizes 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), and test_quantized_format_set_matches_the_taxonomy pins every entry to a real ModelFormat value.

3. Corrected the description. It claimed "every quantized-format loader reaches _apply_fp8_layerwise_casting". Walking ModelLoaderRegistry and checking each class's MRO down to ModelLoader: none of the 17 quantized-format loaders calls it — not FluxGGUFCheckpointModel, not FluxBnbQuantizednf4bCheckpointModel, not the Qwen-Image / Z-Image / Krea-2 / Wan / FLUX.2 GGUF loaders, not the three SDNQ loaders, not the T5/Qwen3/Gemma2/Mistral encoders. (qwen_image.py:243 is in QwenImageCheckpointModel, not the GGUF one.) Same on main. So "reproduce on main: any GGUF main model fails at load" does not reproduce; only the direct _apply_fp8_to_nn_module snippet does.

That does not make the change less worth taking — it makes it defense-in-depth, which is the right framing given that FP8 is being wired into loaders one model at a time (#9414, #9415) and the next one to be wired would walk straight into the crash with _should_use_fp8 saying yes. The body, the code comment and the test docstring now say that instead.

Attacks that failed

• ModelFormat.X in frozenset[str] — my main suspicion, since Enum.__hash__ is hash(self._name_). It holds: the MRO is (ModelFormat, str, Enum, object), so str.__hash__ wins and hash(member) == hash(member.value).
• Real GGMLTensor in a Linear via load_state_dict(assign=True): type(p) and type(p.data) both stay GGMLTensor and is_floating_point() is False — caught by both signals, so either one alone would still do it.
• False positives — nn.Parameter(plain).data is torch.Tensor, not Parameter, so ordinary layers still cast. Zero params flagged across the real FLUX.1-dev transformer (780 params) and AnimaTransformer (601) built on meta. GGMLTensor and SDNQTensor are the only torch.Tensor subclasses in the repo, so the subclass signal has nothing benign to trip on.
• Mutation testing — deleting the format guard, the param guard, or either signal inside _is_quantized_param each fails exactly one test. Nothing here is decorative.
• Frontend — useMainModelDefaultSettings always supplies fp8Storage and RHF's shouldUnregister defaults to false, so data.fp8Storage.isEnabled in onSubmit still resolves for the hidden field. Same pattern as the existing isFluxFamily fields, and a persisted legacy true round-trips unchanged rather than being silently cleared.

One follow-up, not blocking

_apply_fp8_layerwise_casting derives the compute dtype from next(model.parameters()).dtype before the param-level check runs. For exactly the case the backstop advertises — a plain-format checkpoint whose weights were quantized externally — a quantized first parameter makes that torch.uint8:

get_model_compute_dtype -> torch.uint8
attn weight after cast:   torch.float8_e4m3fn   (the quantized layer is correctly skipped)
forward RAISED: RuntimeError: expected m1 and m2 to have the same dtype, but got: c10::BFloat16 != unsigned char
max abs weight drift vs original: 0.486

The quantized layers are skipped correctly, but the pre-hook on the layers that were cast does p.data.to(uint8) and destroys them. set_fp8_compute_dtype only rejects float8, not integer dtypes. One-line fix — derive from the first floating-point parameter — but it is not reachable today (it needs a quantized model on a cast-calling loader), so it does not need to widen this PR. Noted in the description as a follow-up.