From 9954ce3cecec39dfcde4269c773a4611b4e80def Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Fri, 31 Jul 2026 08:49:54 +0200 Subject: [PATCH 1/2] fix(fp8): resolve compute dtype instead of reading model.dtype MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SDXL with fp8_storage crashed before the UNet was ever called: NotImplementedError: "pow_cuda" not implemented for 'Float8_e4m3fn' After layerwise casting the UNet's weights are float8_e4m3fn, and diffusers derives `model.dtype` from the first parameter — so `unet.dtype` reports a storage-only dtype. The legacy SD/SDXL denoise path used it for every tensor it built, so the latents were created in float8 and the first bit of scheduler math (`sigma ** 2` in `add_noise`) blew up. torch has no arithmetic kernels for float8; it is only valid for weights that the forward hooks cast up per layer. Add `get_model_compute_dtype()`: returns `model.dtype` for normal models and the compute dtype for fp8 ones. The loader records the compute dtype on the model when it applies the cast; if the marker is missing (older cache entry, Krea2 encoder path) the resolver scans for the first non-fp8 float param, which works because the cast skips norm layers. Converted every site that derived a tensor dtype from a possibly-fp8 model: latents, noise, mask, masked_latents, conditioning, IP-Adapter and LoRA patch weights in denoise_latents and tiled_multi_diffusion_denoise_latents, plus the LoRA and T2I-Adapter extensions on the modular path. ControlNet and T2I-Adapter control images had the same latent bug — those configs expose an fp8_storage toggle too, so their control image would have been built in float8. Also point LayerPatcher at the shared FP8_STORAGE_DTYPES constant. Regression test covers the real loader path: `model.dtype` is float8 while the resolver returns fp16, and the resolved dtype survives the scheduler arithmetic that crashed. --- invokeai/app/invocations/denoise_latents.py | 22 +++-- .../tiled_multi_diffusion_denoise_latents.py | 15 ++- .../model_manager/load/load_default.py | 6 ++ .../model_manager/load/model_loaders/krea2.py | 4 + invokeai/backend/patches/layer_patcher.py | 9 +- .../stable_diffusion/extensions/lora.py | 3 +- .../extensions/t2i_adapter.py | 3 +- invokeai/backend/util/fp8.py | 54 +++++++++++ tests/backend/util/test_fp8.py | 94 +++++++++++++++++++ 9 files changed, 190 insertions(+), 20 deletions(-) create mode 100644 invokeai/backend/util/fp8.py create mode 100644 tests/backend/util/test_fp8.py diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 89508dfca68..901c7be40ad 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -76,6 +76,7 @@ from invokeai.backend.stable_diffusion.schedulers import SCHEDULER_MAP from invokeai.backend.stable_diffusion.schedulers.schedulers import SCHEDULER_NAME_VALUES from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.fp8 import get_model_compute_dtype from invokeai.backend.util.hotfixes import ControlNetModel from invokeai.backend.util.mask import to_standard_float_mask from invokeai.backend.util.silence_warnings import SilenceWarnings @@ -466,7 +467,7 @@ def prep_control_data( # batch_size=batch_size * num_images_per_prompt, # num_images_per_prompt=num_images_per_prompt, device=device, - dtype=control_model.dtype, + dtype=get_model_compute_dtype(control_model), control_mode=control_info.control_mode, resize_mode=control_info.resize_mode, ) @@ -671,7 +672,7 @@ def run_t2i_adapters( height=control_height_resize, num_channels=t2i_adapter_model.config["in_channels"], # mypy treats this as a FrozenDict device=device, - dtype=t2i_adapter_model.dtype, + dtype=get_model_compute_dtype(t2i_adapter_model), resize_mode=t2i_adapter_field.resize_mode, ) @@ -1018,18 +1019,21 @@ def _lora_loader() -> Iterator[PatchSpec]: model=unet, patches=_lora_loader(), prefix="lora_unet_", - dtype=unet.dtype, + # NOT unet.dtype: with fp8 storage that is float8_e4m3fn, which has no arithmetic + # kernels — every tensor in the denoise loop must use the compute dtype. + dtype=get_model_compute_dtype(unet), cached_weights=cached_weights, ), ): assert isinstance(unet, UNet2DConditionModel) - latents = latents.to(device=device, dtype=unet.dtype) + unet_dtype = get_model_compute_dtype(unet) + latents = latents.to(device=device, dtype=unet_dtype) if noise is not None: - noise = noise.to(device=device, dtype=unet.dtype) + noise = noise.to(device=device, dtype=unet_dtype) if mask is not None: - mask = mask.to(device=device, dtype=unet.dtype) + mask = mask.to(device=device, dtype=unet_dtype) if masked_latents is not None: - masked_latents = masked_latents.to(device=device, dtype=unet.dtype) + masked_latents = masked_latents.to(device=device, dtype=unet_dtype) scheduler = get_scheduler( context=context, @@ -1047,7 +1051,7 @@ def _lora_loader() -> Iterator[PatchSpec]: positive_conditioning_field=self.positive_conditioning, negative_conditioning_field=self.negative_conditioning, device=device, - dtype=unet.dtype, + dtype=unet_dtype, latent_height=latent_height, latent_width=latent_width, cfg_scale=self.cfg_scale, @@ -1072,7 +1076,7 @@ def _lora_loader() -> Iterator[PatchSpec]: exit_stack=exit_stack, latent_height=latent_height, latent_width=latent_width, - dtype=unet.dtype, + dtype=unet_dtype, ) timesteps, init_timestep, scheduler_step_kwargs = self.init_scheduler( diff --git a/invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py b/invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py index a946999070a..72e1308e733 100644 --- a/invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py +++ b/invokeai/app/invocations/tiled_multi_diffusion_denoise_latents.py @@ -35,6 +35,7 @@ ) from invokeai.backend.tiles.utils import TBLR from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.fp8 import get_model_compute_dtype def crop_controlnet_data(control_data: ControlNetData, latent_region: TBLR) -> ControlNetData: @@ -204,13 +205,19 @@ def _lora_loader() -> Iterator[PatchSpec]: ExitStack() as exit_stack, context.models.load(self.unet.unet) as unet, LayerPatcher.apply_smart_model_patches( - model=unet, patches=_lora_loader(), prefix="lora_unet_", dtype=unet.dtype + # NOT unet.dtype: with fp8 storage that is the float8 storage dtype, which has no + # arithmetic kernels (see get_model_compute_dtype). + model=unet, + patches=_lora_loader(), + prefix="lora_unet_", + dtype=get_model_compute_dtype(unet), ), ): assert isinstance(unet, UNet2DConditionModel) - latents = latents.to(device=device, dtype=unet.dtype) + unet_dtype = get_model_compute_dtype(unet) + latents = latents.to(device=device, dtype=unet_dtype) if noise is not None: - noise = noise.to(device=device, dtype=unet.dtype) + noise = noise.to(device=device, dtype=unet_dtype) scheduler = get_scheduler( context=context, scheduler_info=self.unet.scheduler, @@ -226,7 +233,7 @@ def _lora_loader() -> Iterator[PatchSpec]: positive_conditioning_field=self.positive_conditioning, negative_conditioning_field=self.negative_conditioning, device=device, - dtype=unet.dtype, + dtype=unet_dtype, latent_height=latent_tile_height, latent_width=latent_tile_width, cfg_scale=self.cfg_scale, diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index ba617b2b55a..5d527f3c740 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -28,6 +28,7 @@ SubModelType, ) from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.fp8 import set_fp8_compute_dtype # Layer classes that benefit from FP8 storage. Mirrors diffusers' # `_GO_LC_SUPPORTED_PYTORCH_LAYERS` so the plain-nn.Module fallback path makes the same @@ -314,6 +315,11 @@ def _apply_fp8_layerwise_casting( else: return model + # Record the compute dtype so callers can recover it. 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`). + set_fp8_compute_dtype(model, compute_dtype) + param_bytes = sum(p.nelement() * p.element_size() for p in model.parameters()) self._logger.info( f"FP8 layerwise casting enabled for {config.name} " diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index 98cd4a2dec3..670b13cebeb 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/invokeai/backend/model_manager/load/model_loaders/krea2.py @@ -26,6 +26,7 @@ ) from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.fp8 import set_fp8_compute_dtype def _normalize_qwen3vl_rope_config(config: Any) -> Any: @@ -593,6 +594,9 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod # when it shares the GPU with a large transformer. if source_is_fp8 and self._torch_device.type == "cuda": self._apply_fp8_to_nn_module(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=model_dtype) + # `model.dtype` now reports the float8 storage dtype; record the real compute dtype so + # callers can recover it via `get_model_compute_dtype`. + set_fp8_compute_dtype(model, model_dtype) self._logger.info( f"FP8 layerwise casting enabled for Qwen3-VL encoder '{config.name}' " f"(storage=float8_e4m3fn, compute={model_dtype})." diff --git a/invokeai/backend/patches/layer_patcher.py b/invokeai/backend/patches/layer_patcher.py index 232cc3cc142..0b836fccc9a 100644 --- a/invokeai/backend/patches/layer_patcher.py +++ b/invokeai/backend/patches/layer_patcher.py @@ -11,6 +11,7 @@ from invokeai.backend.patches.pad_with_zeros import pad_with_zeros from invokeai.backend.util import InvokeAILogger from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.fp8 import FP8_STORAGE_DTYPES from invokeai.backend.util.original_weights_storage import OriginalWeightsStorage # The optional third item pins the patch's model-cache record without moving the whole patch to VRAM. @@ -199,13 +200,11 @@ def apply_smart_model_patch( def _is_any_part_of_layer_on_cpu(layer: torch.nn.Module) -> bool: return any(p.device.type == "cpu" for p in layer.parameters()) - # FP8 storage dtypes. Direct patching does in-place arithmetic on the model weights, which has no - # CUDA kernel for these dtypes, so a layer with fp8 weights must be patched via the sidecar wrapper. - _FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) - + # Direct patching does in-place arithmetic on the model weights, which has no CUDA kernel for the + # FP8 storage dtypes, so a layer with fp8 weights must be patched via the sidecar wrapper. @staticmethod def _is_any_part_of_layer_fp8(layer: torch.nn.Module) -> bool: - return any(p.dtype in LayerPatcher._FP8_DTYPES for p in layer.parameters()) + return any(p.dtype in FP8_STORAGE_DTYPES for p in layer.parameters()) @staticmethod @torch.no_grad() diff --git a/invokeai/backend/stable_diffusion/extensions/lora.py b/invokeai/backend/stable_diffusion/extensions/lora.py index 229e4a42a34..8102b1dd8a5 100644 --- a/invokeai/backend/stable_diffusion/extensions/lora.py +++ b/invokeai/backend/stable_diffusion/extensions/lora.py @@ -8,6 +8,7 @@ from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.model_patch_raw import ModelPatchRaw from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase +from invokeai.backend.util.fp8 import get_model_compute_dtype if TYPE_CHECKING: from invokeai.app.invocations.model import ModelIdentifierField @@ -47,7 +48,7 @@ def patch_unet(self, unet: UNet2DConditionModel, original_weights: OriginalWeigh patch_weight=self._weight, original_weights=original_weights, original_modules={}, - dtype=unet.dtype, + dtype=get_model_compute_dtype(unet), force_direct_patching=True, force_sidecar_patching=False, ) diff --git a/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py b/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py index 67fede93664..63778953d5f 100644 --- a/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py +++ b/invokeai/backend/stable_diffusion/extensions/t2i_adapter.py @@ -13,6 +13,7 @@ from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.fp8 import get_model_compute_dtype if TYPE_CHECKING: from invokeai.app.invocations.model import ModelIdentifierField @@ -91,7 +92,7 @@ def _run_model( height=input_height, num_channels=model.config["in_channels"], device=TorchDevice.choose_torch_device(), - dtype=model.dtype, + dtype=get_model_compute_dtype(model), resize_mode=self._resize_mode, ) diff --git a/invokeai/backend/util/fp8.py b/invokeai/backend/util/fp8.py new file mode 100644 index 00000000000..c5670903791 --- /dev/null +++ b/invokeai/backend/util/fp8.py @@ -0,0 +1,54 @@ +"""Helpers for models loaded with FP8 layerwise-casting storage. + +See `ModelLoader._apply_fp8_layerwise_casting`: eligible layers keep their weights in +`float8_e4m3fn` and forward hooks cast them up to the model's compute dtype (fp16/bf16) for the +duration of each forward pass. + +The consequence for callers is that `model.dtype` — which diffusers derives from the first +parameter — reports `float8_e4m3fn`. That is a *storage* dtype: torch has no CUDA kernels for +arithmetic on it (`"pow_cuda" not implemented for 'Float8_e4m3fn'`). So any code that reads a +model's dtype in order to build or cast *tensors* (latents, noise, conditioning, control images, +LoRA patch weights) must use `get_model_compute_dtype()` instead of `model.dtype`. +""" + +import torch + +from invokeai.backend.util.devices import TorchDevice + +# Storage-only float8 dtypes. Weights may be held in these, but no math may be done in them. +FP8_STORAGE_DTYPES: tuple[torch.dtype, ...] = (torch.float8_e4m3fn, torch.float8_e5m2) + +# Attribute set on a model by the loader when FP8 layerwise casting is applied. It lives in the +# module's `__dict__` (torch.dtype is not a Parameter/Module), so it survives the deepcopy of the +# meta shell in the shared-CPU-weights adoption path. +FP8_COMPUTE_DTYPE_ATTR = "_invokeai_fp8_compute_dtype" + + +def set_fp8_compute_dtype(model: torch.nn.Module, compute_dtype: torch.dtype) -> None: + """Record the dtype that `model`'s fp8-cast layers compute in.""" + setattr(model, FP8_COMPUTE_DTYPE_ATTR, compute_dtype) + + +def get_model_compute_dtype(model: torch.nn.Module) -> torch.dtype: + """Return the dtype that `model` actually computes in. + + Equivalent to `model.dtype` for normally-loaded models. For models with FP8 storage it returns + the compute dtype (fp16/bf16) rather than the float8 storage dtype. + """ + marked = getattr(model, FP8_COMPUTE_DTYPE_ATTR, None) + if isinstance(marked, torch.dtype): + return marked + + dtype = getattr(model, "dtype", None) + if not isinstance(dtype, torch.dtype): + dtype = next((p.dtype for p in model.parameters()), None) + if isinstance(dtype, torch.dtype) and dtype not in FP8_STORAGE_DTYPES: + return dtype + + # Marker missing but the model is in fp8 storage. The cast skips precision-sensitive layers + # (norms, position/patch embeddings), so the first non-fp8 float param reveals the compute + # dtype. Fall back to the global torch dtype if every param happens to be fp8. + for param in model.parameters(): + if param.is_floating_point() and param.dtype not in FP8_STORAGE_DTYPES: + return param.dtype + return TorchDevice.choose_torch_dtype() diff --git a/tests/backend/util/test_fp8.py b/tests/backend/util/test_fp8.py new file mode 100644 index 00000000000..2338ed0197f --- /dev/null +++ b/tests/backend/util/test_fp8.py @@ -0,0 +1,94 @@ +"""Tests for `get_model_compute_dtype`. + +Regression coverage for the SDXL + fp8_storage crash: + + NotImplementedError: "pow_cuda" not implemented for 'Float8_e4m3fn' + +The legacy SD/SDXL denoise path derived every tensor dtype (latents, noise, conditioning, +control images, LoRA patch weights) from `unet.dtype`. With fp8 layerwise casting the UNet's +first parameter is `float8_e4m3fn`, so the latents were created in a storage-only dtype and the +first bit of scheduler math (`sigma ** 2` in `add_noise`) blew up before the UNet was ever +called. +""" + +from logging import getLogger +from types import SimpleNamespace + +import pytest +import torch + +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType +from invokeai.backend.util.fp8 import ( + FP8_COMPUTE_DTYPE_ATTR, + get_model_compute_dtype, + set_fp8_compute_dtype, +) + + +def _fp8_supported() -> bool: + return hasattr(torch, "float8_e4m3fn") + + +class _MiniUNet(torch.nn.Module): + """Stand-in for a diffusers model: castable layers plus a skipped norm, and a `dtype` + property that reports the first parameter's dtype (what `ModelMixin.dtype` does).""" + + def __init__(self): + super().__init__() + self.conv_in = torch.nn.Conv2d(4, 4, 3, padding=1) + self.norm1 = torch.nn.LayerNorm(4) + self.linear = torch.nn.Linear(4, 4) + + @property + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + +def test_returns_model_dtype_for_normal_model(): + model = _MiniUNet().to(torch.float16) + assert get_model_compute_dtype(model) == torch.float16 + + +def test_returns_marked_compute_dtype(): + model = _MiniUNet().to(torch.float16) + set_fp8_compute_dtype(model, torch.bfloat16) + assert get_model_compute_dtype(model) == torch.bfloat16 + + +@pytest.mark.skipif(not _fp8_supported(), reason="torch.float8_e4m3fn not available") +def test_returns_compute_dtype_after_real_fp8_cast(): + """End-to-end over the loader's real casting path: `model.dtype` reports the float8 storage + dtype, but the resolver must hand back the compute dtype so downstream tensors stay usable.""" + loader = ModelLoader.__new__(ModelLoader) + loader._torch_device = torch.device("cuda") + loader._torch_dtype = torch.float16 + loader._logger = getLogger("test") + config = SimpleNamespace( + type=ModelType.Main, + base=BaseModelType.StableDiffusionXL, + name="test", + default_settings=SimpleNamespace(fp8_storage=True), + ) + + model = loader._apply_fp8_layerwise_casting(_MiniUNet().to(torch.float16), config) + + # Precondition: the naive `model.dtype` read is the storage dtype — this is what crashed. + assert model.dtype == torch.float8_e4m3fn + assert get_model_compute_dtype(model) == torch.float16 + + # And the resolved dtype supports the arithmetic the scheduler does on it. + sigma = torch.tensor([1.5], dtype=get_model_compute_dtype(model)) + assert torch.isfinite(1 / ((sigma**2 + 1) ** 0.5)).all() + + +@pytest.mark.skipif(not _fp8_supported(), reason="torch.float8_e4m3fn not available") +def test_falls_back_to_scan_when_marker_missing(): + """A model that reached us without the marker (e.g. an older cache entry) must still resolve: + the fp8 cast skips norm layers, so their dtype reveals the compute dtype.""" + model = _MiniUNet().to(torch.bfloat16) + ModelLoader._apply_fp8_to_nn_module(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16) + assert not hasattr(model, FP8_COMPUTE_DTYPE_ATTR) + + assert model.dtype == torch.float8_e4m3fn + assert get_model_compute_dtype(model) == torch.bfloat16 From 79dc7f5202dc8654be8d8cbe946daab39ece2607 Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 3 Aug 2026 23:32:00 +0200 Subject: [PATCH 2/2] fix(fp8): harden compute-dtype marker against double-cast poisoning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review follow-up on the compute-dtype resolver. `_apply_fp8_layerwise_casting` derives the compute dtype from the first parameter and is not idempotent. Called on an already-cast model, the first param is float8, so it would record float8 as the *compute* dtype — and `get_model_compute_dtype` trusts the marker, silently reintroducing the "pow_cuda" not implemented for 'Float8_e4m3fn' crash. Two guards close the class: `set_fp8_compute_dtype` rejects any storage-only dtype, and the cast early-returns when the marker is already present. Move the marker-setting into `_apply_fp8_to_nn_module` itself. It was duplicated at both call sites (load_default and krea2's text encoder), so a third caller would have to remember it — the exact failure the fallback scan exists to paper over. Log a warning when the last-resort fallback fires (fp8 storage, no marker, no non-fp8 float param): it returns the global torch dtype, which is wrong for a bf16-compute model and would otherwise surface as an unexplained mismatch deep in the forward pass. Note the same bug in the vendored HiDiffusion pipeline, which builds control images from `controlnet.dtype` at four sites. Dead code today — only apply_hidiffusion/remove_hidiffusion are imported — but it would reproduce the crash if ever wired up with an fp8 ControlNet. Tests: the float8-marker guard for both fp8 dtypes, the marker is set by the cast itself, and a double cast is a no-op (skipped norm layer stays in compute dtype, hooks registered once). The marker-missing fallback test now simulates a legacy model with delattr instead of locking in the old split. --- invokeai/backend/hidiffusion/hidiffusion.py | 7 +++ .../model_manager/load/load_default.py | 20 +++++-- .../model_manager/load/model_loaders/krea2.py | 6 +- invokeai/backend/util/fp8.py | 22 +++++++- tests/backend/util/test_fp8.py | 55 ++++++++++++++++++- 5 files changed, 96 insertions(+), 14 deletions(-) diff --git a/invokeai/backend/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index 5d67f2554ec..18f5a499296 100644 --- a/invokeai/backend/hidiffusion/hidiffusion.py +++ b/invokeai/backend/hidiffusion/hidiffusion.py @@ -141,6 +141,13 @@ class sdxl_controlnet_ppl(block_class): # Save for unpatching later _parent = block_class + # NOTE: `__call__` below builds its control images with `dtype=controlnet.dtype`. That is + # unsafe for a ControlNet loaded with fp8_storage: `.dtype` then reports the float8 *storage* + # dtype, which has no arithmetic kernels (see `invokeai.backend.util.fp8`). It is inert today + # — InvokeAI imports only `apply_hidiffusion` / `remove_hidiffusion` from this vendored file + # and never runs this pipeline — but if it is ever wired up, those reads must go through + # `get_model_compute_dtype()`. + @torch.no_grad() def __call__( self, diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index 5d527f3c740..ca8ba6d5ab1 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -28,7 +28,7 @@ SubModelType, ) from invokeai.backend.util.devices import TorchDevice -from invokeai.backend.util.fp8 import set_fp8_compute_dtype +from invokeai.backend.util.fp8 import FP8_COMPUTE_DTYPE_ATTR, set_fp8_compute_dtype # Layer classes that benefit from FP8 storage. Mirrors diffusers' # `_GO_LC_SUPPORTED_PYTORCH_LAYERS` so the plain-nn.Module fallback path makes the same @@ -289,6 +289,12 @@ def _apply_fp8_layerwise_casting( if not self._should_use_fp8(config, submodel_type): return model + # The cast is not idempotent: on a second pass the first parameter is already fp8, so the + # compute dtype below would be derived as float8. The marker is set by + # `_apply_fp8_to_nn_module`, so its presence means this model has already been cast. + if isinstance(model, torch.nn.Module) and getattr(model, FP8_COMPUTE_DTYPE_ATTR, None) is not None: + return model + storage_dtype = torch.float8_e4m3fn compute_dtype = self._torch_dtype @@ -315,11 +321,6 @@ def _apply_fp8_layerwise_casting( else: return model - # Record the compute dtype so callers can recover it. 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`). - set_fp8_compute_dtype(model, compute_dtype) - param_bytes = sum(p.nelement() * p.element_size() for p in model.parameters()) self._logger.info( f"FP8 layerwise casting enabled for {config.name} " @@ -337,7 +338,14 @@ def _apply_fp8_to_nn_module(model: torch.nn.Module, storage_dtype: torch.dtype, `_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. + + 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 + call sites so a new caller cannot forget it. """ + set_fp8_compute_dtype(model, compute_dtype) + for module_name, module in model.named_modules(): if not isinstance(module, _FP8_SUPPORTED_PYTORCH_LAYERS): continue diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index 670b13cebeb..33558b4ced3 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/invokeai/backend/model_manager/load/model_loaders/krea2.py @@ -26,7 +26,6 @@ ) from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.util.devices import TorchDevice -from invokeai.backend.util.fp8 import set_fp8_compute_dtype def _normalize_qwen3vl_rope_config(config: Any) -> Any: @@ -593,10 +592,9 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod # halves the encoder's resident VRAM (~8.9GB bf16 -> ~4.4GB), which avoids partial-load thrashing # when it shares the GPU with a large transformer. if source_is_fp8 and self._torch_device.type == "cuda": + # `model.dtype` now reports the float8 storage dtype; `_apply_fp8_to_nn_module` records + # the real compute dtype so callers can recover it via `get_model_compute_dtype`. self._apply_fp8_to_nn_module(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=model_dtype) - # `model.dtype` now reports the float8 storage dtype; record the real compute dtype so - # callers can recover it via `get_model_compute_dtype`. - set_fp8_compute_dtype(model, model_dtype) self._logger.info( f"FP8 layerwise casting enabled for Qwen3-VL encoder '{config.name}' " f"(storage=float8_e4m3fn, compute={model_dtype})." diff --git a/invokeai/backend/util/fp8.py b/invokeai/backend/util/fp8.py index c5670903791..ce4d5b7a753 100644 --- a/invokeai/backend/util/fp8.py +++ b/invokeai/backend/util/fp8.py @@ -14,6 +14,9 @@ import torch from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.logging import InvokeAILogger + +logger = InvokeAILogger.get_logger(__name__) # Storage-only float8 dtypes. Weights may be held in these, but no math may be done in them. FP8_STORAGE_DTYPES: tuple[torch.dtype, ...] = (torch.float8_e4m3fn, torch.float8_e5m2) @@ -26,6 +29,16 @@ def set_fp8_compute_dtype(model: torch.nn.Module, compute_dtype: torch.dtype) -> None: """Record the dtype that `model`'s fp8-cast layers compute in.""" + if compute_dtype in FP8_STORAGE_DTYPES: + # A float8 compute dtype is never valid, and recording one would silently reintroduce the + # very crash this module exists to prevent: `get_model_compute_dtype` trusts the marker, so + # every downstream tensor would be built in a dtype torch has no arithmetic kernels for. + # The realistic way to get here is deriving the compute dtype from a model that is already + # cast (i.e. casting twice) — fail loudly at the source instead. + raise ValueError( + f"Refusing to record {compute_dtype} as an FP8 compute dtype; it is a storage-only dtype. " + "This usually means the compute dtype was derived from an already-fp8-cast model." + ) setattr(model, FP8_COMPUTE_DTYPE_ATTR, compute_dtype) @@ -51,4 +64,11 @@ def get_model_compute_dtype(model: torch.nn.Module) -> torch.dtype: for param in model.parameters(): if param.is_floating_point() and param.dtype not in FP8_STORAGE_DTYPES: return param.dtype - return TorchDevice.choose_torch_dtype() + + fallback = TorchDevice.choose_torch_dtype() + logger.warning( + f"{type(model).__name__} is in FP8 storage but carries no compute-dtype marker and has no non-fp8 float " + f"parameter to infer one from; falling back to {fallback}. If the model computes in a different dtype, " + "expect a dtype mismatch during the forward pass." + ) + return fallback diff --git a/tests/backend/util/test_fp8.py b/tests/backend/util/test_fp8.py index 2338ed0197f..a77583c7189 100644 --- a/tests/backend/util/test_fp8.py +++ b/tests/backend/util/test_fp8.py @@ -84,11 +84,60 @@ def test_returns_compute_dtype_after_real_fp8_cast(): @pytest.mark.skipif(not _fp8_supported(), reason="torch.float8_e4m3fn not available") def test_falls_back_to_scan_when_marker_missing(): - """A model that reached us without the marker (e.g. an older cache entry) must still resolve: - the fp8 cast skips norm layers, so their dtype reveals the compute dtype.""" + """A model that reached us without the marker (e.g. an older cache entry, cast before the + marker existed) must still resolve: the fp8 cast skips norm layers, so their dtype reveals the + compute dtype.""" model = _MiniUNet().to(torch.bfloat16) ModelLoader._apply_fp8_to_nn_module(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16) - assert not hasattr(model, FP8_COMPUTE_DTYPE_ATTR) + # Simulate the legacy, marker-less model. + delattr(model, FP8_COMPUTE_DTYPE_ATTR) assert model.dtype == torch.float8_e4m3fn assert get_model_compute_dtype(model) == torch.bfloat16 + + +@pytest.mark.skipif(not _fp8_supported(), reason="torch.float8_e4m3fn not available") +def test_cast_records_the_marker_itself(): + """The marker must be set by the cast itself, not by its callers — a caller that forgets it + would fall through to the scan, and a caller that derives the dtype from an already-cast model + would record float8.""" + model = _MiniUNet().to(torch.bfloat16) + ModelLoader._apply_fp8_to_nn_module(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16) + + assert getattr(model, FP8_COMPUTE_DTYPE_ATTR) == torch.bfloat16 + + +@pytest.mark.parametrize("storage_dtype", [torch.float8_e4m3fn, torch.float8_e5m2]) +def test_rejects_float8_as_compute_dtype(storage_dtype: torch.dtype): + """A float8 compute dtype would be handed straight back by `get_model_compute_dtype`, silently + reintroducing the crash. Recording one must fail loudly instead.""" + model = _MiniUNet().to(torch.float16) + with pytest.raises(ValueError, match="storage-only dtype"): + set_fp8_compute_dtype(model, storage_dtype) + + assert not hasattr(model, FP8_COMPUTE_DTYPE_ATTR) + + +@pytest.mark.skipif(not _fp8_supported(), reason="torch.float8_e4m3fn not available") +def test_double_cast_is_a_noop(): + """The cast is not idempotent: a second pass would derive the compute dtype from the already-fp8 + first parameter. An already-marked model must be left alone.""" + loader = ModelLoader.__new__(ModelLoader) + loader._torch_device = torch.device("cuda") + loader._torch_dtype = torch.float16 + loader._logger = getLogger("test") + config = SimpleNamespace( + type=ModelType.Main, + base=BaseModelType.StableDiffusionXL, + name="test", + default_settings=SimpleNamespace(fp8_storage=True), + ) + + model = loader._apply_fp8_layerwise_casting(_MiniUNet().to(torch.bfloat16), config) + model = loader._apply_fp8_layerwise_casting(model, config) + + assert get_model_compute_dtype(model) == torch.bfloat16 + # The skipped norm layer is still in the compute dtype — a second cast would have taken it and + # the hooks would have been registered twice. + assert model.norm1.weight.dtype == torch.bfloat16 + assert len(model.linear._forward_pre_hooks) == 1