diff --git a/invokeai/app/invocations/denoise_latents.py b/invokeai/app/invocations/denoise_latents.py index 413c7bc5fa4..2d48dd87607 100644 --- a/invokeai/app/invocations/denoise_latents.py +++ b/invokeai/app/invocations/denoise_latents.py @@ -78,6 +78,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 @@ -497,7 +498,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, ) @@ -710,7 +711,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, ) @@ -1075,18 +1076,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, @@ -1104,7 +1108,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, @@ -1129,7 +1133,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/hidiffusion/hidiffusion.py b/invokeai/backend/hidiffusion/hidiffusion.py index a5a53197111..33437853214 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 ba617b2b55a..ca8ba6d5ab1 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 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 @@ -288,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 @@ -331,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 98cd4a2dec3..33558b4ced3 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/invokeai/backend/model_manager/load/model_loaders/krea2.py @@ -592,6 +592,8 @@ 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) self._logger.info( f"FP8 layerwise casting enabled for Qwen3-VL encoder '{config.name}' " 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..ce4d5b7a753 --- /dev/null +++ b/invokeai/backend/util/fp8.py @@ -0,0 +1,74 @@ +"""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 +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) + +# 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.""" + 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) + + +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 + + 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 new file mode 100644 index 00000000000..a77583c7189 --- /dev/null +++ b/tests/backend/util/test_fp8.py @@ -0,0 +1,143 @@ +"""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, 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) + # 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