From 5d61920191447af7b7977fbcf750600800c0a121 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Tue, 25 Aug 2026 02:42:08 +0200 Subject: [PATCH] fix(fp8): honor every declared skip list, and stop overshooting the RAM reservation Review follow-ups for #9414 and #9415. Depends on #9415: the docs below describe Anima's FP8 support, which lands there. Read `_keep_in_fp32_modules` alongside `_skip_layerwise_casting_patterns`. Diffusers' `enable_layerwise_casting()` unions both; we replaced that call with our own hook-based path and were reading only the first, so a model declaring the second would lose its exclusions silently. Verified to protect nothing extra today - on Krea-2, Wan 14B, Z-Image and FLUX.1 - so this changes nothing now and stops being a trap later. Release the state dict before the FP8 cast in the Z-Image and Krea-2 single-file loaders. `load_state_dict(..., assign=True)` aliases every param to its `sd` tensor, so the compute-dtype originals stayed reachable while `param.data.to(float8)` allocated the fp8 copies, putting peak RAM ~50% over what `make_room()` reserved (~17.4GB actual against ~11.5GB reserved for Z-Image). Nothing reads `sd` after the load. Add `test_z_image_fp8_wiring.py`. Deleting the cast call from the Z-Image single-file loader previously left the whole model_manager suite green. The new tests fail on that, on removing `sd.clear()`, and on the aliasing premise itself, should torch ever stop assigning by reference. Use `get_model_compute_dtype()` in the Z-Image denoise loop instead of `transformer.dtype`. It is correct today only because `x_pad_token` happens to be parameter zero and is never cast; move the pad tokens under a submodule and the loop starts feeding float8 into `F.linear`. Reword the comment above the Z-Image cast. Dropping `.scale_weight` / `scaled_fp8` is not "filtering out metadata" - for a ComfyUI scaled-fp8 checkpoint it loads unscaled weights. That bug is pre-existing and out of scope here, but the comment read as though the cast made it safe. Update the FP8 docs, which still said Z-Image was excluded for a dtype mismatch and listed it in the troubleshooting exclusion list - the opposite of what the code has done since #9414. Add Anima and its LLLite adapters, and document that a model's own declared exclusions are honored on top of the generic skip list, with the measured cost: Wan 14B gives up ~221 MiB of savings, Krea-2 ~38 MiB, Anima ~18 MiB, FLUX.1 and Qwen-Image nothing. --- .../docs/configuration/fp8-storage.mdx | 34 ++++-- invokeai/app/invocations/z_image_denoise.py | 8 +- .../model_manager/load/load_default.py | 41 +++++-- .../model_manager/load/model_loaders/krea2.py | 4 + .../load/model_loaders/z_image.py | 18 ++- .../load/test_load_default_fp8.py | 56 +++++++++ .../load/test_z_image_fp8_wiring.py | 106 ++++++++++++++++++ 7 files changed, 241 insertions(+), 26 deletions(-) create mode 100644 tests/backend/model_manager/load/test_z_image_fp8_wiring.py diff --git a/docs/src/content/docs/configuration/fp8-storage.mdx b/docs/src/content/docs/configuration/fp8-storage.mdx index 29b6ae95f3f..f9ebbcbab35 100644 --- a/docs/src/content/docs/configuration/fp8-storage.mdx +++ b/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 `. 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 `. 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 diff --git a/invokeai/app/invocations/z_image_denoise.py b/invokeai/app/invocations/z_image_denoise.py index 07658fe96eb..09af26bde92 100644 --- a/invokeai/app/invocations/z_image_denoise.py +++ b/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)) diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index 5d547193d6f..2357c2c3890 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -157,6 +157,34 @@ def _device_supports_fp8_storage(device: torch.device, logger: Optional[Logger] ) +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). @@ -425,18 +453,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 @@ -464,8 +485,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. Records the compute dtype on the model. After the cast, `model.dtype` reports the float8 diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index dbc4b109c6a..fbd681c8025 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/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 diff --git a/invokeai/backend/model_manager/load/model_loaders/z_image.py b/invokeai/backend/model_manager/load/model_loaders/z_image.py index 0ca5dfb010a..ffda6eaf9cd 100644 --- a/invokeai/backend/model_manager/load/model_loaders/z_image.py +++ b/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 diff --git a/tests/backend/model_manager/load/test_load_default_fp8.py b/tests/backend/model_manager/load/test_load_default_fp8.py index 17b96c3ad93..fe6ed8e9667 100644 --- a/tests/backend/model_manager/load/test_load_default_fp8.py +++ b/tests/backend/model_manager/load/test_load_default_fp8.py @@ -25,6 +25,7 @@ _FP8_STORAGE_SUPPORTED, 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, @@ -581,3 +582,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 diff --git a/tests/backend/model_manager/load/test_z_image_fp8_wiring.py b/tests/backend/model_manager/load/test_z_image_fp8_wiring.py new file mode 100644 index 00000000000..631dc4a48fe --- /dev/null +++ b/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