diff --git a/docs/src/content/docs/configuration/fp8-storage.mdx b/docs/src/content/docs/configuration/fp8-storage.mdx index 56d73198ed0..89d4db63993 100644 --- a/docs/src/content/docs/configuration/fp8-storage.mdx +++ b/docs/src/content/docs/configuration/fp8-storage.mdx @@ -1,5 +1,5 @@ --- -title: FP8 Storage +title: FP8 Storage & Compute sidebar: order: 3 --- @@ -29,9 +29,9 @@ InvokeAI's FP8 path stores weights in FP8 and casts them back to BF16/FP16 on ea The toggle works as advertised: the UNet / transformer drops by roughly 50% on the GPU. Per-step latency is the same or marginally slower because every forward pass adds an FP8 → BF16 cast on entry and a BF16 → FP8 cast on exit. This is the **largest target group**: 3090 owners squeezing FLUX into 24 GB benefit the most. -### RTX 40-series, RTX 50-series, and Hopper — VRAM win today, compute win possible later +### RTX 40-series, RTX 50-series, and Hopper — VRAM win, plus a compute win via a separate setting -These GPUs have native FP8 tensor cores. The toggle still buys you the same ~50% VRAM reduction today, because the forward pass still runs in BF16 — the hook casts weights back up to compute precision before each layer. If InvokeAI later wires up a true FP8 matmul path (e.g. via `torchao`), the same toggle will *also* unlock compute speedups on this hardware. Until then, treat the benefit as "VRAM only, same as Ampere". +These GPUs have native FP8 tensor cores. FP8 *Storage* on its own still only buys the ~50% VRAM reduction, because the forward pass runs in BF16 — the hook casts weights back up to compute precision before each layer. To actually use the tensor cores you need the separate [FP8 Compute](#fp8-compute) setting, which applies to checkpoints that ship pre-quantized ("scaled fp8") rather than to full-precision ones. ### Older CUDA cards — still a VRAM win @@ -105,10 +105,40 @@ If you see unexpected quality regressions, disable FP8 Storage on the affected m ## Combining with Low-VRAM mode -**FP8 + partial loading**: fully supported. FP8 Storage shrinks the layers; partial loading streams them between RAM and VRAM as needed. Use both on tight VRAM budgets. +**FP8 Storage + partial loading**: fully supported. FP8 Storage shrinks the layers; partial loading streams them between RAM and VRAM as needed. Use both on tight VRAM budgets. + +**FP8 Compute + partial loading** is a different story — it still works, but it costs both speed and reproducibility. See [FP8 Compute](#fp8-compute) below. (For why FP8 Storage doesn't stack on top of GGUF / NF4 / int8 checkpoints, see the callout at the top of this page.) +## FP8 Compute + +Everything above describes FP8 *Storage*, which changes how weights are **stored** while the math still runs in BF16. `fp8_compute` is a separate, global setting in `invokeai.yaml` that also does the **math** in FP8, on GPUs that have hardware for it. That makes generation faster, not just smaller. + +It is for models that were **already saved in FP8** by whoever published them — you'll often see these labelled "fp8" or "fp8_scaled" in the filename. Normally InvokeAI unpacks them back to BF16 while loading; with `fp8_compute` on, they stay as they are and run directly on the GPU's FP8 hardware. FP8 Storage is the toggle for the other case: a full-precision model you want to shrink yourself. + +Not every part of a model stays in FP8 — the small, precision-sensitive pieces are always unpacked. Those are a tiny share of the weights, so you still get nearly the full VRAM saving. + +You need an FP8-capable GPU: RTX 40-series or newer, the datacenter cards of those generations, or AMD MI300 and newer. InvokeAI checks your card the first time it needs to — by actually trying a small FP8 operation rather than going by the model name — so a card that can't do it quietly falls back to the normal path instead of failing partway through a generation. + +:::danger[Reproducibility requires the model to be fully in VRAM] +With `fp8_compute` enabled, **the same seed only reproduces the same image if the model is 100% resident in VRAM.** + +FP8 math only happens for the parts of the model that are actually on the GPU. Anything still sitting in system RAM takes the normal path instead, which gives slightly different numbers — and *which* parts that affects depends on how much of the model happened to fit at that moment, which changes from run to run. + +Measured on a 24 GB card with a ~12 GB model at 88–95% loaded: two runs with an identical seed and identical settings differed in **98.7% of all pixels**. With the model fully loaded, repeated runs came out **identical**. + +To get repeatable output, check the model's load line in the log for `VRAM: … (100.0%)`. If it is below that, free up VRAM (lower resolution, fewer models loaded at once, a smaller text encoder) or set `enable_partial_loading: false`. +::: + +Keeping the whole model on the GPU is worth it for speed as well: on the same 24 GB card, having to stream the last 5–12% of the model over PCIe cost **+47% per step** (1.03 → 1.51 s/it at 1024², 8 steps). + +### When the model asks for full precision + +Some FP8 models come with a note from whoever made them, marking certain layers as ones that should not use FP8 math. InvokeAI follows those notes by default, so those layers run the slower way. On a model that marks a lot of layers, this can eat much of the FP8 Compute speedup. + +Setting `fp8_compute_full_precision_hints: false` ignores the notes and runs everything on the FP8 hardware. It is faster, but you are overriding the model author's judgement about which layers are sensitive — so compare a few images before sticking with it. + ## Troubleshooting ### "I toggled FP8 Storage but VRAM usage didn't change" diff --git a/docs/src/generated/settings.json b/docs/src/generated/settings.json index 5a73c0483af..3bd24caaa07 100644 --- a/docs/src/generated/settings.json +++ b/docs/src/generated/settings.json @@ -468,6 +468,28 @@ "type": "", "validation": {} }, + { + "category": "CACHE", + "default": false, + "description": "Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).", + "env_var": "INVOKEAI_FP8_COMPUTE", + "literal_values": [], + "name": "fp8_compute", + "required": false, + "type": "", + "validation": {} + }, + { + "category": "CACHE", + "default": true, + "description": "Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.", + "env_var": "INVOKEAI_FP8_COMPUTE_FULL_PRECISION_HINTS", + "literal_values": [], + "name": "fp8_compute_full_precision_hints", + "required": false, + "type": "", + "validation": {} + }, { "category": "CACHE", "default": null, diff --git a/invokeai/app/services/config/config_default.py b/invokeai/app/services/config/config_default.py index 7174742a66d..a0f32e15604 100644 --- a/invokeai/app/services/config/config_default.py +++ b/invokeai/app/services/config/config_default.py @@ -106,6 +106,8 @@ class InvokeAIAppConfig(BaseSettings): device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value. enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM. keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high. + fp8_compute: Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM). + fp8_compute_full_precision_hints: Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled. ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable. vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable. lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable. @@ -210,6 +212,8 @@ class InvokeAIAppConfig(BaseSettings): device_working_mem_gb: float = Field(default=3, description="The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.") enable_partial_loading: bool = Field(default=True, description="Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.") keep_ram_copy_of_weights: bool = Field(default=True, description="Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.") + fp8_compute: bool = Field(default=False, description="Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).") + fp8_compute_full_precision_hints: bool = Field(default=True, description="Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.") # Deprecated CACHE configs ram: Optional[float] = Field(default=None, gt=0, description="DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.") vram: Optional[float] = Field(default=None, ge=0, description="DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.") diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index e33d870263f..76007dabeb7 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -6,7 +6,7 @@ import re from logging import Logger from pathlib import Path -from typing import Optional +from typing import Callable, Optional import torch @@ -27,6 +27,7 @@ AnyModel, SubModelType, ) +from invokeai.backend.quantization.fp8_scaled import count_fp8_weights, should_keep_fp8_weights from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.fp8 import FP8_COMPUTE_DTYPE_ATTR, set_fp8_compute_dtype @@ -472,6 +473,22 @@ def _apply_fp8_layerwise_casting( if isinstance(model, torch.nn.Module) and getattr(model, FP8_COMPUTE_DTYPE_ATTR, None) is not None: return model + # A checkpoint that already ships fp8 weights is running (or is about to run) on the fp8 + # tensor cores. Layerwise casting would install hooks that restore the compute dtype before + # every forward, so `CustomLinear._can_use_fp8_matmul` would no longer see an fp8 weight and + # would silently fall back to the dequantized path — the VRAM toggle would make the model + # *slower* with no indication why. Storage has nothing to add here anyway: the weights are + # already 1 byte per parameter. + if isinstance(model, torch.nn.Module) and should_keep_fp8_weights(self._torch_device): + already_fp8 = count_fp8_weights(model) + if already_fp8: + self._logger.info( + f"FP8 storage skipped for {config.name}: {already_fp8} weight(s) are already fp8 and " + "are being run on the fp8 tensor cores (fp8_compute). Layerwise casting would " + "disable that matmul without saving any further VRAM." + ) + return model + storage_dtype = torch.float8_e4m3fn compute_dtype = self._torch_dtype @@ -517,6 +534,7 @@ def _apply_fp8_to_nn_module( storage_dtype: torch.dtype, compute_dtype: torch.dtype, extra_skip_patterns: tuple[str, ...] = (), + skip: Optional[Callable[[str, torch.nn.Module], bool]] = None, ) -> None: """Apply FP8 layerwise casting to a plain nn.Module. @@ -530,6 +548,12 @@ def _apply_fp8_to_nn_module( `_model_declared_skip_patterns`), which are model-specific and cannot be inferred from layer types or generic name patterns. + `skip` excludes further modules by (dotted name, module). Its one caller uses it to leave + scaled-fp8 layers alone: those already hold fp8 weights plus a `weight_scale`, and the cast + hooks installed here would upcast them *without* applying that scale — a silently wrong + weight. Casting only the remainder lets a partly-quantized checkpoint (fp8 language model, + bf16 visual tower) end up fully fp8-resident. + 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 @@ -548,6 +572,8 @@ def _apply_fp8_to_nn_module( continue if any(re.search(pattern, module_name) for pattern in skip_patterns): continue + if skip is not None and skip(module_name, module): + continue params = list(module.parameters(recurse=False)) if not params: continue diff --git a/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_linear.py b/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_linear.py index 451c5a9c86b..10c24f6c2b5 100644 --- a/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_linear.py +++ b/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_linear.py @@ -9,6 +9,13 @@ from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch from invokeai.backend.patches.layers.flux_control_lora_layer import FluxControlLoRALayer from invokeai.backend.patches.layers.lora_layer import LoRALayer +from invokeai.backend.quantization.fp8_scaled import ( + FP8_DTYPE, + dequantize_weight, + device_supports_fp8_matmul, + is_fp8_matmul_enabled, + scaled_mm_linear, +) from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor from invokeai.backend.quantization.sdnq.sdnq_tensor import SDNQTensor @@ -90,22 +97,77 @@ def _cast_tensor_for_input(self, tensor: torch.Tensor | None, input: torch.Tenso return tensor def _cast_weight_bias_for_input(self, input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: - weight = self._cast_tensor_for_input(self.weight, input) + weight = cast_to_device(self.weight, input.device) + assert weight is not None + if weight.dtype == FP8_DTYPE: + # A scaled fp8 weight must never be cast with a plain `.to(dtype)`: that drops the + # scale and yields a wrongly-scaled weight. `weight_scale` is None for fp8 weights that + # carry no scale (e.g. produced by fp8_storage layerwise casting), where a plain cast + # is correct. + weight = dequantize_weight(weight, self._fp8_weight_scale(input.device), input.dtype) + else: + weight = self._cast_tensor_for_input(weight, input) bias = self._cast_tensor_for_input(self.bias, input) assert weight is not None return weight, bias + def _fp8_weight_scale(self, device: torch.device) -> torch.Tensor | None: + scale = getattr(self, "weight_scale", None) + return cast_to_device(scale, device) if scale is not None else None + + def _can_use_fp8_matmul(self, input: torch.Tensor) -> bool: + """Whether this forward may run on the fp8 tensor cores. + + Every condition is a hard requirement of `torch._scaled_mm` or an explicit instruction from + the checkpoint producer; failing any of them falls back to the dequantized path rather than + raising, because a mid-generation failure is far worse than losing the speedup. + """ + if not is_fp8_matmul_enabled(): + return False + if self.weight.dtype != FP8_DTYPE or not input.is_floating_point(): + return False + if getattr(self, "_fp8_full_precision_matmul", False): + # The producer measured this layer as unsafe to multiply in fp8. + return False + if not device_supports_fp8_matmul(input.device): + return False + if self.in_features % 16 != 0 or self.out_features % 16 != 0: + return False + # Partial loading may leave the weight on the CPU; _scaled_mm needs both operands co-located. + return self.weight.device == input.device + + def _maybe_fp8_forward(self, input: torch.Tensor) -> torch.Tensor | None: + """Run the fp8 matmul if this module and input allow it, else None so callers fall back.""" + if not self._can_use_fp8_matmul(input): + return None + return scaled_mm_linear( + input, + self.weight, + self._fp8_weight_scale(input.device), + cast_to_device(self.bias, input.device), + input_scale=cast_to_device(getattr(self, "input_scale", None), input.device), + ) + def _autocast_forward_with_patches(self, input: torch.Tensor) -> torch.Tensor: return autocast_linear_forward_sidecar_patches(self, input, self._patches_and_weights) def _autocast_forward(self, input: torch.Tensor) -> torch.Tensor: + out = self._maybe_fp8_forward(input) + if out is not None: + return out weight, bias = self._cast_weight_bias_for_input(input) return torch.nn.functional.linear(input, weight, bias) def forward(self, input: torch.Tensor) -> torch.Tensor: if len(self._patches_and_weights) > 0: return self._autocast_forward_with_patches(input) - elif self._device_autocasting_enabled: + # Checked before the autocasting branch on purpose: `apply_custom_layers_to_model` leaves + # device autocasting *disabled* whenever the model is fully resident, so a check that only + # lived in `_autocast_forward` would silently never run in the common case. + fp8_out = self._maybe_fp8_forward(input) + if fp8_out is not None: + return fp8_out + if self._device_autocasting_enabled: return self._autocast_forward(input) elif input.is_floating_point() and ( (self.weight.is_floating_point() and self.weight.dtype != input.dtype) diff --git a/invokeai/backend/model_manager/load/model_loaders/anima.py b/invokeai/backend/model_manager/load/model_loaders/anima.py index 97782b07ffd..9af0d2c24c4 100644 --- a/invokeai/backend/model_manager/load/model_loaders/anima.py +++ b/invokeai/backend/model_manager/load/model_loaders/anima.py @@ -10,7 +10,7 @@ from invokeai.backend.model_manager.configs.controlnet import ControlNet_Checkpoint_Anima_Config from invokeai.backend.model_manager.configs.factory import AnyModelConfig from invokeai.backend.model_manager.configs.main import Main_Checkpoint_Anima_Config -from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.load_default import ModelLoader, _model_declared_skip_patterns from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.taxonomy import ( AnyModel, @@ -19,6 +19,21 @@ ModelType, SubModelType, ) +from invokeai.backend.quantization.fp8_scaled import ( + attach_fp8_scales, + cast_state_dict, + dequantize_fp8_scaled, + extract_comfy_quant_hints, + extract_fp8_scaled_layers, + full_precision_hints_respected, + parse_quantization_metadata, + predict_cast_state_dict_size, + read_safetensors_metadata, + should_keep_fp8_weights, + split_fp8_scaled_layers, + strip_layer_path_prefix, + warn_on_unattached_scales, +) from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger @@ -157,22 +172,58 @@ def _load_from_singlefile( # Drop runtime-derived buffers and exporter metadata that aren't model weights. sd = _filter_non_model_keys(sd) + target_device = TorchDevice.choose_torch_device() + model_dtype = TorchDevice.choose_anima_inference_dtype(target_device) + + # ComfyUI 'scaled fp8': an fp8 weight plus a `weight_scale`. `_filter_non_model_keys` above + # keeps those keys, and `load_state_dict` below rejects the checkpoint outright over them -- + # 500 unexpected keys on a plain scaled export, 749 on one that also ships `comfy_quant` + # markers. Such a checkpoint therefore does not load at all today. + # + # Anima keeps `q_proj`/`k_proj`/`v_proj` separate and the only key rewrite is a prefix strip, + # so a sibling scale travels with its weight and nothing has to be split. + keep_fp8 = should_keep_fp8_weights(target_device) + header_hints = parse_quantization_metadata(read_safetensors_metadata(model_path, logger)) + # The header names layers in the checkpoint's own scheme -- `net.`-prefixed on every Anima + # redistribution measured -- while the scales are read after `_strip_anima_bundle_prefix` + # has run. Without this the per-layer flags, `full_precision_matrix_mult` above all, match + # nothing and are silently ignored. + layer_hints = { + **extract_comfy_quant_hints(sd), + **strip_layer_path_prefix(header_hints), + } + fp8_layers = extract_fp8_scaled_layers(sd, layer_hints=layer_hints) + if fp8_layers and not keep_fp8: + # Without the matmul, keeping them quantized would halve VRAM but dequantize on every + # forward. Fold the scale into the weight instead. + dequantize_fp8_scaled(sd, fp8_layers, model_dtype) + fp8_layers = {} + # Create an empty AnimaTransformer with Anima's default architecture parameters with accelerate.init_empty_weights(): model = AnimaTransformer(**ANIMA_TRANSFORMER_CONFIG) - # Determine safe dtype - target_device = TorchDevice.choose_torch_device() - model_dtype = TorchDevice.choose_anima_inference_dtype(target_device) - - # Handle memory management - new_sd_size = sum(ten.nelement() * model_dtype.itemsize for ten in sd.values()) - self._ram_cache.make_room(new_sd_size) + skip_patterns = _model_declared_skip_patterns(model) + # Layers the cast would dequantize anyway are folded here, scale applied, so the cast never + # strips a scale that can no longer be put back. + # Reserve before the split, not after: `split_fp8_scaled_layers` dequantizes its unusable + # subset through fp32, so reserving afterwards lets that transient peak land on an + # unreserved cache. `scaled_layers` is what keeps that honest: the split also widens layers + # whose scale layout `scaled_mm` cannot apply, and without the mapping the prediction would + # charge those 1 byte/element and arrive at 2. + self._ram_cache.make_room( + predict_cast_state_dict_size( + sd, + model_dtype, + keep_fp8=keep_fp8, + model=model, + skip_patterns=skip_patterns, + scaled_layers=fp8_layers, + ) + ) - # Convert to target dtype (skip non-float tensors like embedding indices) - for k in sd.keys(): - if sd[k].is_floating_point(): - sd[k] = sd[k].to(model_dtype) + fp8_layers = split_fp8_scaled_layers(sd, fp8_layers, model_dtype, model=model, skip_patterns=skip_patterns) + kept = cast_state_dict(sd, model_dtype, keep_fp8=keep_fp8, model=model, skip_patterns=skip_patterns) load_result = model.load_state_dict(sd, assign=True, strict=False) if load_result.unexpected_keys: @@ -191,6 +242,19 @@ def _load_from_singlefile( # 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`. + if fp8_layers: + attached = attach_fp8_scales(model, fp8_layers) + logger.info(f"Anima: kept {attached} layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled)") + warn_on_unattached_scales(logger, "Anima", attached, fp8_layers) + marked = sum(1 for layer in fp8_layers.values() if layer.full_precision_matmul) + if marked and full_precision_hints_respected(): + logger.info( + f"Anima: {marked} of {len(fp8_layers)} layer(s) are marked full_precision_matrix_mult " + "and will dequantize per forward." + ) + elif kept: + logger.info(f"Anima: kept {kept} raw fp8 weight(s) quantized for the fp8 tensor cores.") + model = self._apply_fp8_layerwise_casting(model, config, SubModelType.Transformer) return model diff --git a/invokeai/backend/model_manager/load/model_loaders/flux.py b/invokeai/backend/model_manager/load/model_loaders/flux.py index 8ad9459735a..7e1dfc9a8a2 100644 --- a/invokeai/backend/model_manager/load/model_loaders/flux.py +++ b/invokeai/backend/model_manager/load/model_loaders/flux.py @@ -62,11 +62,16 @@ T5Encoder_T5Encoder_Config, ) from invokeai.backend.model_manager.configs.vae import VAE_Checkpoint_Config_Base, VAE_Checkpoint_Flux2_Config -from invokeai.backend.model_manager.load.load_default import ModelLoader, resolve_submodel_path +from invokeai.backend.model_manager.load.load_default import ( + ModelLoader, + _model_declared_skip_patterns, + resolve_submodel_path, +) from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.load.model_loaders.flux2_state_dict_utils import ( convert_flux2_bfl_to_diffusers, convert_flux2_vae_bfl_to_diffusers, + remap_flux2_layer_paths, ) from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader from invokeai.backend.model_manager.taxonomy import ( @@ -80,6 +85,24 @@ from invokeai.backend.model_manager.util.model_util import ( convert_bundle_to_flux_transformer_checkpoint, ) +from invokeai.backend.quantization.fp8_scaled import ( + attach_fp8_scales, + can_stay_quantized, + cast_state_dict, + dequantize_fp8_scaled, + extract_comfy_quant_hints, + extract_fp8_scaled_layers, + full_precision_hints_respected, + is_scale_metadata_key, + iter_weight_scale_pairs, + parse_quantization_metadata, + predict_cast_state_dict_size, + read_safetensors_metadata, + should_keep_fp8_weights, + split_fp8_scaled_layers, + strip_layer_path_prefix, + warn_on_unattached_scales, +) from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.quantization.gguf.utils import TORCH_COMPATIBLE_QTYPES from invokeai.backend.quantization.sdnq.detection import is_sdnq_folder @@ -715,12 +738,77 @@ def _load_from_singlefile( sd = load_file(model_path) if "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale" in sd: sd = convert_bundle_to_flux_transformer_checkpoint(sd) - new_sd_size = sum([ten.nelement() * torch.bfloat16.itemsize for ten in sd.values()]) - self._ram_cache.make_room(new_sd_size) - for k in sd.keys(): - # We need to cast to bfloat16 due to it being the only currently supported dtype for inference - sd[k] = sd[k].to(torch.bfloat16) + + # ComfyUI 'scaled fp8' (fp8 weight + .weight_scale/.scale_weight, optionally an input + # scale). Read *after* the bundle conversion: the scale keys carry the same + # `model.diffusion_model.` prefix as the weights and only line up with the module tree once + # it has been stripped. Until now the loader had no idea these keys existed and + # `load_state_dict` rejected the checkpoint outright as unexpected keys. + # + # Unlike FLUX.2 there is no key conversion here: the checkpoint is already in the BFL + # layout this model implements, and `qkv` stays one fused Linear, so a per-tensor scale + # attaches to exactly the module it was computed for. + # + # Hints ship either in the safetensors header or as per-layer `.comfy_quant` markers; the + # header wins on the rare checkpoint carrying both. + metadata = read_safetensors_metadata(model_path, self._logger) + # Header names still carry the checkpoint prefix that was stripped off `sd`; strip it from + # them too, or every per-layer flag matches nothing. + layer_hints = { + **extract_comfy_quant_hints(sd), + **strip_layer_path_prefix(parse_quantization_metadata(metadata)), + } + fp8_layers = extract_fp8_scaled_layers(sd, layer_hints=layer_hints) + + # A checkpoint that ships raw fp8 weights (fp8 tensors, no weight_scale) keeps them when the + # fp8 matmul is available; casting them to bf16 here would throw away both the VRAM saving + # and the tensor cores before the model is ever built. + keep_fp8 = should_keep_fp8_weights(self._torch_device) + if fp8_layers and not keep_fp8: + # Without the matmul, keeping them quantized would halve VRAM but dequantize on every + # forward. Fold the scale into the weight instead — the legacy result, except the scale + # is now actually applied rather than dropped. + dequantize_fp8_scaled(sd, fp8_layers, torch.bfloat16) + fp8_layers = {} + + skip_patterns = _model_declared_skip_patterns(model) + # Scaled layers that the cast would dequantize anyway are folded here, scale applied, so + # `cast_state_dict` never strips a scale that can no longer be put back. + # Reserve before the split, not after: `split_fp8_scaled_layers` dequantizes its unusable + # subset through fp32, so reserving afterwards lets that transient peak land on an + # unreserved cache. `scaled_layers` is what keeps that honest: the split also widens layers + # whose scale layout `scaled_mm` cannot apply, and without the mapping the prediction would + # charge those 1 byte/element and arrive at 2. + self._ram_cache.make_room( + predict_cast_state_dict_size( + sd, + torch.bfloat16, + keep_fp8=keep_fp8, + model=model, + skip_patterns=skip_patterns, + scaled_layers=fp8_layers, + ) + ) + + fp8_layers = split_fp8_scaled_layers(sd, fp8_layers, torch.bfloat16, model=model, skip_patterns=skip_patterns) + # Everything else is cast to bfloat16, the only dtype currently supported for inference. + kept = cast_state_dict(sd, torch.bfloat16, keep_fp8=keep_fp8, model=model, skip_patterns=skip_patterns) model.load_state_dict(sd, assign=True) + + if fp8_layers: + attached = attach_fp8_scales(model, fp8_layers) + self._logger.info(f"FLUX: kept {attached} layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled)") + warn_on_unattached_scales(self._logger, "FLUX", attached, fp8_layers) + marked = sum(1 for layer in fp8_layers.values() if layer.full_precision_matmul) + if marked and full_precision_hints_respected(): + self._logger.info( + f"FLUX: {marked} of {len(fp8_layers)} layer(s) are marked full_precision_matrix_mult " + "and will dequantize per forward. Set fp8_compute_full_precision_hints=false to run " + "them on the fp8 tensor cores instead." + ) + elif kept: + self._logger.info(f"FLUX: kept {kept} raw fp8 weight(s) quantized for the fp8 tensor cores.") + return model @@ -969,11 +1057,10 @@ def _load_from_singlefile( # Load state dict sd = load_file(model_path) - # Handle FP8 quantized weights (ComfyUI-style or scaled FP8) - # These store weights as: layer.weight (FP8) + layer.weight_scale (FP32 scalar) - sd = self._dequantize_fp8_weights(sd) + keep_fp8 = should_keep_fp8_weights(self._torch_device) - # Check if keys have ComfyUI-style prefix and strip if needed + # Check if keys have ComfyUI-style prefix and strip if needed. This runs before anything + # reads the quantization side-channel: the scales carry the same prefix as their weights. prefix_to_strip = None for prefix in ["model.diffusion_model.", "diffusion_model."]: if any(k.startswith(prefix) for k in sd.keys() if isinstance(k, str)): @@ -986,9 +1073,39 @@ def _load_from_singlefile( for k, v in sd.items() } - # Convert BFL format state dict to diffusers format + # ComfyUI 'scaled fp8' (fp8 weight + .weight_scale, optionally an input scale). Until now + # these were folded into bf16 at load, so a FLUX.2 checkpoint that ships them ran as fp8 + # *storage* at best and never reached the tensor cores. + # + # Hints ship in the safetensors header or as per-layer `.comfy_quant` markers, and both name + # layers in the checkpoint's BFL scheme. The scales are read after the rename, so the hints + # have to be renamed too -- a fused `qkv` becomes three diffusers layers. + header_hints = strip_layer_path_prefix( + parse_quantization_metadata(read_safetensors_metadata(model_path, self._logger)) + ) + layer_hints = {**extract_comfy_quant_hints(sd), **header_hints} + path_map = remap_flux2_layer_paths(layer_hints.keys()) + layer_hints = { + renamed: hints for name, hints in layer_hints.items() for renamed in (path_map.get(name) or [name]) + } + + # Convert BFL format state dict to diffusers format. Scales and markers are carried to + # wherever their weight landed, including across the fused-qkv split. converted_sd = convert_flux2_bfl_to_diffusers(sd) + fp8_layers = extract_fp8_scaled_layers(converted_sd, layer_hints=layer_hints) + if fp8_layers and not keep_fp8: + # Without the matmul, keeping them quantized would halve VRAM but dequantize on every + # forward. Fold the scale into the weight instead -- the legacy result, except reached + # through the shared helper. + dequantize_fp8_scaled(converted_sd, fp8_layers, torch.bfloat16) + fp8_layers = {} + + # Safety net for scale layouts the shared extractor does not model (block-wise scales whose + # shape has to be expanded to the weight's). It is a no-op on every checkpoint measured so + # far, because extraction has already taken the scales it understood. + converted_sd = self._dequantize_fp8_weights(converted_sd, keep_fp8=keep_fp8) + # Detect architecture from checkpoint keys double_block_indices = [ int(k.split(".")[1]) @@ -1063,16 +1180,45 @@ def _load_from_singlefile( out_features2, in_features2, dtype=torch.bfloat16 ) - # Convert to bfloat16 and load - for k in converted_sd.keys(): - converted_sd[k] = converted_sd[k].to(torch.bfloat16) + # Convert to bfloat16 and load, leaving raw fp8 weights quantized when the tensor cores can + # take them (the scaled-fp8 path above has already folded any weight_scale it found). The + # model's own precision-sensitive list is honored — see the Z-Image loader for why that is + # a correctness requirement and not just a quality nicety. + skip_patterns = _model_declared_skip_patterns(model) + # Scaled layers the cast would dequantize anyway are folded here, scale applied, so + # `cast_state_dict` never strips a scale that can no longer be put back. + fp8_layers = split_fp8_scaled_layers( + converted_sd, fp8_layers, torch.bfloat16, model=model, skip_patterns=skip_patterns + ) + + kept = cast_state_dict( + converted_sd, + torch.bfloat16, + keep_fp8=keep_fp8, + model=model, + skip_patterns=skip_patterns, + ) # Load the state dict - guidance weights were already initialized above if missing model.load_state_dict(converted_sd, assign=True) + if fp8_layers: + attached = attach_fp8_scales(model, fp8_layers) + self._logger.info(f"FLUX.2: kept {attached} layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled)") + warn_on_unattached_scales(self._logger, "FLUX.2", attached, fp8_layers) + marked = sum(1 for layer in fp8_layers.values() if layer.full_precision_matmul) + if marked and full_precision_hints_respected(): + self._logger.info( + f"FLUX.2: {marked} of {len(fp8_layers)} layer(s) are marked full_precision_matrix_mult " + "and will dequantize per forward. Set fp8_compute_full_precision_hints=false to run " + "them on the fp8 tensor cores instead." + ) + elif kept: + self._logger.info(f"FLUX.2: kept {kept} raw fp8 weight(s) quantized for the fp8 tensor cores.") + return model - def _dequantize_fp8_weights(self, sd: dict) -> dict: + def _dequantize_fp8_weights(self, sd: dict, keep_fp8: bool = False) -> dict: """Dequantize FP8 quantized weights in the state dict. ComfyUI and some FLUX.2 models store quantized weights as: @@ -1082,45 +1228,50 @@ def _dequantize_fp8_weights(self, sd: dict) -> dict: Dequantization formula: dequantized = weight.to(float) * weight_scale Also handles FP8 tensors stored with float8_e4m3fn dtype by converting to float. + + ``keep_fp8`` spares the *raw* fp8 weights (fp8 with no ``weight_scale``) that trailing + conversion, so they survive to `cast_state_dict` and reach the tensor cores. Without it this + method converts every float8 tensor unconditionally and nothing fp8 is left downstream — the + FLUX.2 half of the raw-fp8 path would never execute. Scaled weights are always folded here: + they have already had their scale applied a few lines up, so they are no longer float8 by + the time the loop below runs. + + The test is deliberately the model-less form of :func:`can_stay_quantized` — the module tree + does not exist yet at this point, and the keys are still in checkpoint naming. It is a + superset: `cast_state_dict` re-applies the same predicate later *with* the model and casts + whatever turns out not to be an ``nn.Linear`` weight. """ - # Check for ComfyUI-style scale factors - weight_scale_keys = [k for k in sd.keys() if isinstance(k, str) and k.endswith(".weight_scale")] - - for scale_key in weight_scale_keys: - # Get the corresponding weight key - weight_key = scale_key.replace(".weight_scale", ".weight") - if weight_key in sd: - weight = sd[weight_key] - scale = sd[scale_key] - - # Dequantize: convert FP8 to float and multiply by scale - # Note: Float8 types require .float() instead of .to(torch.float32) - weight_float = weight.float() - scale = scale.float() - - # Handle block-wise quantization where scale may have different shape - if scale.dim() > 0 and scale.shape != weight_float.shape and scale.numel() > 1: - for dim in range(len(weight_float.shape)): - if dim < len(scale.shape) and scale.shape[dim] != weight_float.shape[dim]: - block_size = weight_float.shape[dim] // scale.shape[dim] - if block_size > 1: - scale = scale.repeat_interleave(block_size, dim=dim) - - # Do the multiply in float32 for precision, but store bf16 (FLUX.2's compute dtype) - # immediately so the *whole* model is never materialized in float32. Holding every - # dequantized weight as float32 here doubled RAM transiently (~36GB vs ~17GB for a 9B - # model) and was the dominant cold-load spike, especially with two GPUs. The result is - # identical to the previous code, which cast the same values to bf16 a few steps later. - sd[weight_key] = (weight_float * scale).to(torch.bfloat16) - del weight_float + # Check for ComfyUI-style scale factors. Both spellings are folded here, because the + # metadata strip below removes both — reading only `.weight_scale` meant a `.scale_weight` + # checkpoint had its scales deleted without ever being applied, leaving every quantized + # weight off by 1/weight_scale with nothing logged. + for weight_key, scale_key in list(iter_weight_scale_pairs(sd)): + weight = sd[weight_key] + scale = sd[scale_key] + + # Dequantize: convert FP8 to float and multiply by scale + # Note: Float8 types require .float() instead of .to(torch.float32) + weight_float = weight.float() + scale = scale.float() + + # Handle block-wise quantization where scale may have different shape + if scale.dim() > 0 and scale.shape != weight_float.shape and scale.numel() > 1: + for dim in range(len(weight_float.shape)): + if dim < len(scale.shape) and scale.shape[dim] != weight_float.shape[dim]: + block_size = weight_float.shape[dim] // scale.shape[dim] + if block_size > 1: + scale = scale.repeat_interleave(block_size, dim=dim) + + # Do the multiply in float32 for precision, but store bf16 (FLUX.2's compute dtype) + # immediately so the *whole* model is never materialized in float32. Holding every + # dequantized weight as float32 here doubled RAM transiently (~36GB vs ~17GB for a 9B + # model) and was the dominant cold-load spike, especially with two GPUs. The result is + # identical to the previous code, which cast the same values to bf16 a few steps later. + sd[weight_key] = (weight_float * scale).to(torch.bfloat16) + del weight_float # Filter out scale metadata keys and other FP8 metadata - keys_to_remove = [ - k - for k in sd.keys() - if isinstance(k, str) - and (k.endswith(".weight_scale") or k.endswith(".scale_weight") or "comfy_quant" in k or k == "scaled_fp8") - ] + keys_to_remove = [k for k in sd.keys() if is_scale_metadata_key(k)] for k in keys_to_remove: del sd[k] @@ -1135,6 +1286,8 @@ def _dequantize_fp8_weights(self, sd: dict) -> dict: # 0-dimensional tensor (scalar) - likely metadata, remove it keys_to_remove_scalars.append(key) elif hasattr(tensor, "dtype") and "float8" in str(tensor.dtype): + if keep_fp8 and can_stay_quantized(key, tensor, None): + continue # Native FP8 tensor - mark for conversion keys_to_convert.append(key) diff --git a/invokeai/backend/model_manager/load/model_loaders/flux2_state_dict_utils.py b/invokeai/backend/model_manager/load/model_loaders/flux2_state_dict_utils.py index 86fd8478169..43db644ebdc 100644 --- a/invokeai/backend/model_manager/load/model_loaders/flux2_state_dict_utils.py +++ b/invokeai/backend/model_manager/load/model_loaders/flux2_state_dict_utils.py @@ -14,9 +14,15 @@ """ import re +from typing import Any import torch +from invokeai.backend.quantization.fp8_scaled import ( + QKV_SPLIT_SIDECHANNEL_SUFFIXES, + split_qkv_sidechannel, +) + def _flux2_chunk_tensor(tensor, chunks: int): """Chunk a tensor along dim 0, dequantizing GGUF tensors first. @@ -134,8 +140,14 @@ def _convert_flux2_single_block_key(key: str, tensor, converted: dict) -> str | return key -def convert_flux2_bfl_to_diffusers(sd: dict) -> dict: - """Convert a FLUX.2 transformer BFL-format state dict to diffusers format.""" +def _convert_flux2_weight_keys(sd: dict) -> dict: + """Convert the *weight* keys of a FLUX.2 BFL state dict to diffusers format. + + Quantization side-channel keys must not be routed through here. The block renames below are + substring tests, so `img_attn.proj.weight_scale` satisfies `"img_attn.proj.weight" in rest` + and would be written to the weight's destination key -- overwriting the weight with its own + scale. `convert_flux2_bfl_to_diffusers` keeps them out and places them separately. + """ converted: dict = {} # Basic key renames @@ -324,3 +336,92 @@ def convert_flux2_vae_bfl_to_diffusers(sd: dict) -> dict: converted[new_key] = tensor return converted + + +def _flux2_sidechannel_parts(key: Any) -> tuple[str, str] | None: + """Split ``.`` for a quantization side-channel key, else None.""" + if not isinstance(key, str): + return None + for suffix in QKV_SPLIT_SIDECHANNEL_SUFFIXES: + if key.endswith(f".{suffix}"): + return key[: -len(suffix) - 1], suffix + return None + + +# Divisible by both 3 (the fused QKV split) and 2 (the adaLN scale/shift swap), so a probe runs +# through every branch of the converter instead of tripping its "malformed, leave alone" guards. +_PROBE_ROWS = 6 + + +def _flux2_sidechannel_destinations(base: str) -> list[str]: + """Where a layer's weight ends up after conversion, as module paths. + + Rather than restating the rename rules -- which would drift from the converter the moment one + of them changes -- the layer is pushed through the real converter as a lone ``.weight`` + and the resulting keys are read back. A fused ``qkv`` maps to *three* destinations. + """ + probe = {f"{base}.weight": torch.zeros(_PROBE_ROWS, 1)} + return [k[: -len(".weight")] for k in _convert_flux2_weight_keys(probe) if k.endswith(".weight")] + + +# Converter transforms that reorder weight *rows*. A per-output-channel weight scale has one entry +# per row, so it has to be reordered identically or every row ends up scaled by another row's +# factor. The fused-QKV split is handled separately (`split_qkv_sidechannel`); this is the only +# other row-reordering transform the converter applies. +_ROW_PERMUTED_BY_CONVERSION = {"final_layer.adaLN_modulation.1"} + + +def _mirror_row_permutation(base: str, suffix: str, value: Any) -> Any: + """Apply the converter's row reordering to a per-output-channel weight scale. + + Only weight scales carry per-row structure: the activation scale is per-tensor and the + `comfy_quant` marker is a byte blob, both of which describe the whole layer and must be copied + verbatim. A per-tensor weight scale is likewise unaffected. + """ + if base not in _ROW_PERMUTED_BY_CONVERSION or suffix not in ("weight_scale", "scale_weight"): + return value + tensor = torch.as_tensor(value) if hasattr(value, "shape") else value + if not hasattr(tensor, "shape") or tensor.dim() < 1 or tensor.numel() <= 1 or tensor.shape[0] % 2 != 0: + return value + return _flux2_swap_scale_shift(tensor) + + +def convert_flux2_bfl_to_diffusers(sd: dict) -> dict: + """Convert a FLUX.2 transformer BFL-format state dict to diffusers format. + + Quantization scales and markers are carried to wherever their weight landed. Doing that is not + optional for a scaled-fp8 checkpoint: a scale left on the fused ``qkv`` path is keyed on a + module the diffusers model does not have, so `attach_fp8_scales` finds nothing and the three + split weights stay quantized but *unscaled* -- off by 1/weight_scale, with nothing logged. + """ + weights = {k: v for k, v in sd.items() if _flux2_sidechannel_parts(k) is None} + converted = _convert_flux2_weight_keys(weights) + + for key, value in sd.items(): + parts = _flux2_sidechannel_parts(key) + if parts is None: + continue + base, suffix = parts + destinations = _flux2_sidechannel_destinations(base) + if not destinations: + # Unknown layer: keep the key as it is. `extract_fp8_scaled_layers` drops a scale with + # no matching fp8 weight, which is the safe outcome -- better than guessing a target. + converted[key] = value + elif len(destinations) == 1: + converted[f"{destinations[0]}.{suffix}"] = _mirror_row_permutation(base, suffix, value) + else: + for destination, part in zip(destinations, split_qkv_sidechannel(key, value), strict=True): + converted[f"{destination}.{suffix}"] = part + + return converted + + +def remap_flux2_layer_paths(layer_names: Any) -> dict[str, list[str]]: + """Map BFL layer paths to their diffusers equivalents, one-to-many for a fused ``qkv``. + + ``_quantization_metadata`` names its layers in the checkpoint's own scheme, but the scales are + extracted *after* the state dict has been renamed. Without this the per-layer flags -- notably + ``full_precision_matrix_mult`` -- match nothing and are silently ignored, which is the exact + mistake that already cost a debugging round on Krea-2. + """ + return {name: _flux2_sidechannel_destinations(name) for name in layer_names} diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index fbd681c8025..b82e514a439 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/invokeai/backend/model_manager/load/model_loaders/krea2.py @@ -2,7 +2,7 @@ """Class for Krea-2 model loading in InvokeAI.""" from pathlib import Path -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional import accelerate from transformers import AutoConfig, AutoTokenizer @@ -14,7 +14,11 @@ Qwen3VLEncoder_Checkpoint_Config, Qwen3VLEncoder_Qwen3VLEncoder_Config, ) -from invokeai.backend.model_manager.load.load_default import ModelLoader, _device_supports_fp8_storage +from invokeai.backend.model_manager.load.load_default import ( + ModelLoader, + _device_supports_fp8_storage, + _model_declared_skip_patterns, +) from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader from invokeai.backend.model_manager.taxonomy import ( @@ -24,13 +28,26 @@ ModelType, SubModelType, ) +from invokeai.backend.quantization.fp8_scaled import ( + attach_fp8_scales, + cast_state_dict, + dequantize_fp8_scaled, + detach_layer_sidechannel, + extract_comfy_quant_hints, + extract_fp8_scaled_layers, + full_precision_hints_respected, + parse_quantization_metadata, + predict_cast_state_dict_size, + read_safetensors_metadata, + reattach_layer_sidechannel, + should_keep_fp8_weights, + split_fp8_scaled_layers, + strip_layer_path_prefix, + warn_on_unattached_scales, +) from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.util.devices import TorchDevice -if TYPE_CHECKING: - # torch is imported lazily inside the helpers below; this is annotations-only. - import torch - def _normalize_qwen3vl_rope_config(config: Any) -> Any: """Mirror Qwen3-VL rope_parameters into rope_scaling for Transformers compatibility.""" @@ -93,39 +110,29 @@ def _is_native_krea2_format(sd: dict[str, Any]) -> bool: ) -def _dequantize_scaled_fp8(sd: dict[str, Any], dtype: "torch.dtype") -> dict[str, Any]: - """Dequantize ComfyUI 'scaled fp8' weights: ``dequant = weight.float() * weight_scale``. - - Each quantized layer stores an fp8 ``.weight`` plus a (usually scalar) ``.weight_scale``. - Returns a new dict with the weights dequantized and the ``.weight_scale`` keys removed. No-op if - there are no scale keys. +def _remap_native_layer_paths(layer_names: Any) -> dict[str, str]: + """Map native/ComfyUI layer paths to their diffusers equivalents. - The multiply runs in float32 for precision, but each result is stored as ``dtype`` immediately so - the *whole model* is never materialized in float32. Krea-2's ~12 GB fp8 checkpoint would otherwise - peak at ~50 GB of RAM (4 bytes/param) before the caller's later bf16 cast brings it down to ~25 GB, - which puts a 32 GB machine into swap during a cold load. This mirrors the same fix already applied - to the FLUX.2 loader. - - ``dtype`` is required on purpose. It used to default to bfloat16, which is wrong on a device where - ``choose_bfloat16_safe_dtype`` picks float16: the weights would land in bf16 and then take a second - rounding step on the caller's later float16 cast. Callers already know the compute dtype, so there - is no reason to guess one here. + ``_quantization_metadata`` names its layers in the checkpoint's own (native) scheme, but the + scales are extracted after the state dict has been renamed. Rather than restating the rename + rules - which would drift - each name is pushed through the real converter as a lone + ``.weight`` entry and the resulting key is read back. """ import torch - scale_keys = [k for k in sd if isinstance(k, str) and k.endswith(".weight_scale")] - if not scale_keys: - return sd - out = dict(sd) - for scale_key in scale_keys: - weight_key = scale_key.replace(".weight_scale", ".weight") - if weight_key in out: - weight = torch.as_tensor(_to_plain_tensor(out[weight_key])).float() - scale = torch.as_tensor(_to_plain_tensor(out[scale_key])).float() - out[weight_key] = (weight * scale).to(dtype) - del weight - del out[scale_key] - return out + mapping: dict[str, str] = {} + for name in layer_names: + if not isinstance(name, str): + continue + try: + converted = _convert_krea2_native_to_diffusers({f"{name}.weight": torch.empty(0)}) + except Exception: + continue + for key in converted: + if isinstance(key, str) and key.endswith(".weight"): + mapping[name] = key[: -len(".weight")] + break + return mapping def _convert_krea2_native_to_diffusers(sd: dict[str, Any]) -> dict[str, Any]: @@ -351,22 +358,84 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) sd = load_file(model_path) + metadata = read_safetensors_metadata(model_path, self._logger) sd = _strip_comfyui_prefix(sd) - # ComfyUI 'scaled fp8' checkpoints: fold the per-tensor weight_scale into the weights. The - # compute dtype is resolved first so the dequantized weights land there directly instead of - # transiently materializing the whole model in float32. - sd = _dequantize_scaled_fp8(sd, model_dtype) - # Native/ComfyUI key naming → diffusers Krea2Transformer2DModel keys. + # Per-layer `.comfy_quant` markers are read first (they are popped out of `sd` here, before + # the key conversion). Checkpoints ship the flags in either the header or these markers; + # without both, a checkpoint using only the per-tensor form has its + # full_precision_matrix_mult layers silently multiplied in fp8. The header wins on the rare + # checkpoint carrying both. Header names carry the prefix that was just stripped off the + # state dict, so strip it from them too or they match nothing. + layer_hints = { + **extract_comfy_quant_hints(sd), + **strip_layer_path_prefix(parse_quantization_metadata(metadata)), + } if _is_native_krea2_format(sd): + # Take the quantization side channel out before renaming. The converter renames + # ".weight"-suffixed keys by substring and five more by whole-key equality, so a sibling + # ".scale_weight", ".input_scale", or any scale on one of the equality-renamed keys + # (e.g. `last.linear.weight_scale`) would be left behind at its old path while the + # weight moves — and then silently dropped, leaving the weight unscaled. + detached = detach_layer_sidechannel(sd) sd = _convert_krea2_native_to_diffusers(sd) + # The metadata and the detached scales still name layers natively; rename both the same + # way or the per-layer flags (notably full_precision_matrix_mult) match nothing. + path_map = _remap_native_layer_paths({*detached, *layer_hints}) + orphaned = reattach_layer_sidechannel(sd, detached, path_map) + if orphaned: + # INFO, not DEBUG: a dropped scale leaves its weight off by 1/weight_scale with no + # other symptom, so the one line that mentions it must be visible by default. + self._logger.info( + f"Krea-2: dropped quantization side-channel for {len(orphaned)} module(s) with no " + f"diffusers counterpart (e.g. {orphaned[0]})." + ) + layer_hints = {path_map.get(name, name): hints for name, hints in layer_hints.items()} + + # ComfyUI 'scaled fp8' checkpoints (fp8 weight + .weight_scale, optionally .input_scale). + fp8_layers = extract_fp8_scaled_layers(sd, layer_hints=layer_hints) + keep_fp8 = should_keep_fp8_weights(target_device) + if fp8_layers and not keep_fp8: + # Legacy behavior: fold the scales into the weights. Keeping them quantized without the + # fp8 matmul would halve VRAM but run slower, so both are tied to the same setting. + dequantize_fp8_scaled(sd, fp8_layers, model_dtype) + fp8_layers = {} with accelerate.init_empty_weights(): model = Krea2Transformer2DModel(**KREA2_TRANSFORMER_CONFIG) - new_sd_size = sum(ten.nelement() * model_dtype.itemsize for ten in sd.values()) - self._ram_cache.make_room(new_sd_size) - for k in sd.keys(): - sd[k] = sd[k].to(model_dtype) + skip_patterns = _model_declared_skip_patterns(model) + # Scaled layers the cast would dequantize anyway (skip patterns, non-Linear weights) are + # folded here, with their scale applied. Left to `cast_state_dict` they would be cast + # *without* it and `attach_fp8_scales` would then skip them for no longer being fp8 — + # a weight silently off by 1/weight_scale. Krea-2's `time_embed.linear_1/linear_2` are + # ordinary quantized Linears in a ComfyUI export and match the model's `time_embed` pattern. + # Reserve before the split, not after: `split_fp8_scaled_layers` dequantizes its unusable + # subset through fp32, so reserving afterwards lets that transient peak land on an + # unreserved cache. `scaled_layers` is what keeps that honest: the split also widens layers + # whose scale layout `scaled_mm` cannot apply, and without the mapping the prediction would + # charge those 1 byte/element and arrive at 2. + self._ram_cache.make_room( + predict_cast_state_dict_size( + sd, + model_dtype, + keep_fp8=keep_fp8, + model=model, + skip_patterns=skip_patterns, + scaled_layers=fp8_layers, + ) + ) + + fp8_layers = split_fp8_scaled_layers(sd, fp8_layers, model_dtype, model=model, skip_patterns=skip_patterns) + # A checkpoint with raw fp8 weights (fp8 tensors and no weight_scale) yields no fp8_layers at + # all, but its weights are still usable on the tensor cores, so the same `keep_fp8` covers + # both kinds. + kept = cast_state_dict( + sd, + model_dtype, + keep_fp8=keep_fp8, + model=model, + skip_patterns=skip_patterns, + ) model.load_state_dict(sd, assign=True, strict=False) _reject_incomplete_load(model, what="Krea-2 single-file checkpoint") @@ -374,6 +443,50 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: # 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() + + if kept and not fp8_layers: + # Raw fp8: no scales to attach, but say so — otherwise the tensor-core path is invisible, + # and the only other fp8 log line (the scaled one below) never fires for this checkpoint. + self._logger.info( + f"Krea-2: kept {kept} raw fp8 weight(s) quantized (no weight_scale in the checkpoint); " + "they will run on the fp8 tensor cores with unit scaling." + ) + + if fp8_layers: + attached = attach_fp8_scales(model, fp8_layers) + self._logger.info(f"Krea-2: kept {attached} layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled)") + warn_on_unattached_scales(self._logger, "Krea-2", attached, fp8_layers) + # Marked layers dequantize on every forward instead of using the tensor cores, which is + # the single biggest lever on how much fp8_compute actually buys for a given checkpoint. + # Surface it: otherwise "fp8_compute is on but barely faster" has no visible cause. + marked = sum(1 for layer in fp8_layers.values() if layer.full_precision_matmul) + if marked: + if full_precision_hints_respected(): + self._logger.info( + f"Krea-2: {marked} of {len(fp8_layers)} layer(s) are marked " + "full_precision_matrix_mult and will dequantize per forward. Set " + "fp8_compute_full_precision_hints=false to run them on the fp8 tensor cores " + "instead (faster, but overrides the checkpoint producer's instruction)." + ) + else: + self._logger.info( + f"Krea-2: ignoring the full_precision_matrix_mult marker on {marked} layer(s) " + "(fp8_compute_full_precision_hints=false)." + ) + # fp8_storage exists to *create* fp8 weights from full-precision ones; here they already + # are fp8, so it is bypassed entirely. Say so, otherwise a user who enabled it is left + # wondering whether it took effect. + default_settings = getattr(config, "default_settings", None) + if default_settings is not None and getattr(default_settings, "fp8_storage", None): + self._logger.info( + "Krea-2: the model's fp8_storage setting is redundant here and was skipped - the " + "checkpoint already ships fp8 weights." + ) + # The layerwise-casting path exists to *produce* fp8 weights from full-precision ones. The + # checkpoint already is fp8, and its hooks would cast back to the compute dtype without + # applying weight_scale, so it must not run here. + return model + # 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 @@ -479,12 +592,26 @@ def _load_model( ) -def _remap_qwen3vl_singlefile_keys(sd: dict[str, Any]) -> dict[str, Any]: - """Remap ComfyUI single-file Qwen3-VL keys to the transformers ``Qwen3VLModel`` layout. +def _qwen3vl_target_key(key: str) -> str: + """Map one ComfyUI single-file Qwen3-VL key (or module path) to the transformers layout. ComfyUI/native layout uses a single ``model.`` prefix for both towers; transformers splits them: ``model.visual.*`` -> ``visual.*`` and ``model.`` (layers/embed_tokens/norm) -> ``language_model.``. + + Shared by the state-dict remap and the fp8 layer-hint remap so the two cannot drift apart: a hint + keyed by a path the model does not have matches nothing and is silently ignored. """ + # Strip a leading "model." (some checkpoints prefix everything with it), then route by tower. + key = key[len("model.") :] if key.startswith("model.") else key + if key.startswith("visual.") or key.startswith("language_model."): + # Already the transformers layout (e.g. "model.language_model.*" / "model.visual.*"). + return key + # Bare language-model keys (layers.* / embed_tokens / norm) belong under language_model. + return "language_model." + key + + +def _remap_qwen3vl_singlefile_keys(sd: dict[str, Any]) -> dict[str, Any]: + """Remap ComfyUI single-file Qwen3-VL keys to the transformers ``Qwen3VLModel`` layout.""" out: dict[str, Any] = {} source_of: dict[Any, Any] = {} what = "Qwen3-VL encoder checkpoint" @@ -492,14 +619,7 @@ def _remap_qwen3vl_singlefile_keys(sd: dict[str, Any]) -> dict[str, Any]: if not isinstance(k, str): _put_unique_key(out, k, v, source=k, source_of=source_of, what=what) continue - # Strip a leading "model." (some checkpoints prefix everything with it), then route by tower. - key = k[len("model.") :] if k.startswith("model.") else k - if key.startswith("visual.") or key.startswith("language_model."): - # Already the transformers layout (e.g. "model.language_model.*" / "model.visual.*"). - _put_unique_key(out, key, v, source=k, source_of=source_of, what=what) - else: - # Bare language-model keys (layers.* / embed_tokens / norm) belong under language_model. - _put_unique_key(out, "language_model." + key, v, source=k, source_of=source_of, what=what) + _put_unique_key(out, _qwen3vl_target_key(k), v, source=k, source_of=source_of, what=what) return out @@ -583,37 +703,109 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) sd = load_file(str(model_path)) - # Detect an fp8 source (ComfyUI 'scaled fp8' weight_scale keys, or raw float8 weights) BEFORE - # dequantizing. An fp8-on-disk encoder is kept fp8-resident with layerwise upcasting below, so - # it occupies ~half the VRAM of the dequantized bf16 model (the whole point of shipping fp8). - source_is_fp8 = any(isinstance(k, str) and k.endswith(".weight_scale") for k in sd) or any( + metadata = read_safetensors_metadata(model_path, self._logger) + # Per-layer markers must be read before extract_fp8_scaled_layers() drops them, and before + # the key remap, which would not carry a ".comfy_quant" suffix to a sensible destination. + layer_hints = {**extract_comfy_quant_hints(sd), **parse_quantization_metadata(metadata)} + + # Remap BEFORE pulling the scales out. The remap rewrites whole keys, so each + # ".weight_scale" travels with its ".weight" and the recovered layer paths already match the + # model's module paths — which is what attach_fp8_scales() resolves them against. + sd = _remap_qwen3vl_singlefile_keys(sd) + layer_hints = {_qwen3vl_target_key(path): hints for path, hints in layer_hints.items()} + + # ComfyUI 'scaled fp8' (fp8 weight + .weight_scale). Only the language-model linears are + # quantized in the checkpoints seen so far; the visual tower stays bf16 either way. + fp8_layers = extract_fp8_scaled_layers(sd, layer_hints=layer_hints) + source_is_fp8 = bool(fp8_layers) or any( getattr(t, "dtype", None) in (torch.float8_e4m3fn, torch.float8_e5m2) for t in sd.values() ) - # ComfyUI 'scaled fp8': fold weight_scale into the weights, then drop quantization metadata. - sd = _dequantize_scaled_fp8(sd, model_dtype) - for k in list(sd.keys()): - if isinstance(k, str) and (k.endswith(".comfy_quant") or "scale_input" in k): - del sd[k] - sd = _remap_qwen3vl_singlefile_keys(sd) + # Resolved once. `device_supports_fp8_matmul` deliberately does not cache an inconclusive + # probe, so two calls can disagree: a transient failure here followed by a success below + # would leave the scales already folded while the raw Linears stay quantized, i.e. the + # encoder silently on the storage path with the matmul log line never printed. + keep_matmul_fp8 = should_keep_fp8_weights(target_device) + if fp8_layers and not keep_matmul_fp8: + # Legacy behavior: fold the scales into the weights. Without the fp8 matmul, staying + # quantized would save VRAM but cost speed, so the two are tied to the same setting. + dequantize_fp8_scaled(sd, fp8_layers, model_dtype) + fp8_layers = {} te_config = self._load_hf_config() with accelerate.init_empty_weights(): model = Qwen3VLModel._from_config(te_config) - new_sd_size = sum(ten.nelement() * model_dtype.itemsize for ten in sd.values()) - self._ram_cache.make_room(new_sd_size) - for k in sd.keys(): - sd[k] = sd[k].to(model_dtype) + # Same ordering contract as every other loader in this series: split *before* the cast. + # A per-key `if dtype is not FP8_DTYPE` cast looks equivalent and is not — it keeps every + # fp8 tensor quantized, including the ones that must not stay: + # + # - a 1-D fp8 norm (checkpoints that "quantize everything" ship these) would keep its + # scale as an unused buffer on a non-Linear and the forward would compute on the raw + # fp8 codes, i.e. off by 1/weight_scale with nothing logged; + # - an e5m2 scaled weight would be cast *without* its scale, since only e4m3fn is spared; + # - a block-wise scale would reach `scaled_mm_linear` unchecked and raise mid-generation. + # + # `split_fp8_scaled_layers` applies exactly those filters and dequantizes the affected + # layers *with* their scale, so what remains is what the matmul can actually consume. + # The storage path below re-quantizes to fp8 anyway, so casting the raw fp8 Linears to + # `model_dtype` here would double both this reservation and the host-RAM peak (~4.4 -> ~8.9 + # GiB on the 4B encoder) for a round trip that ends where it started -- e4m3fn is a subset + # of bf16, so it is value-exact. Keep them for either consumer. + use_fp8_storage = source_is_fp8 and _device_supports_fp8_storage(self._torch_device, self._logger) + keep_fp8 = keep_matmul_fp8 or use_fp8_storage + # Reserve before the split: it dequantizes its unusable subset through fp32, so reserving + # afterwards lets that transient peak land on an unreserved cache. `scaled_layers` keeps the + # prediction in step with the split's own scale-layout filter. + self._ram_cache.make_room( + predict_cast_state_dict_size(sd, model_dtype, keep_fp8=keep_fp8, model=model, scaled_layers=fp8_layers) + ) + fp8_layers = split_fp8_scaled_layers(sd, fp8_layers, model_dtype, model=model) + cast_state_dict(sd, model_dtype, keep_fp8=keep_fp8, model=model) model.load_state_dict(sd, assign=True, strict=False) _reject_incomplete_load(model, what="Qwen3-VL encoder checkpoint") + if fp8_layers: + # Keep the weights quantized and let CustomLinear run their matmuls on the fp8 tensor + # cores. Same resident VRAM as the layerwise-casting path below, but without paying a + # fp8->bf16 round trip on every forward. + attached = attach_fp8_scales(model, fp8_layers) + warn_on_unattached_scales(self._logger, f"Qwen3-VL encoder '{config.name}'", attached, fp8_layers) + # The checkpoint quantizes only the language-model linears; the visual tower ships bf16 + # and would make the encoder ~0.4GB larger resident than the plain fp8_storage path. On a + # GPU already holding a ~12GB transformer that is enough to push the transformer into + # partial loading, so cast the rest to fp8 storage. Two exclusions: + # + # - anything carrying a `weight_scale`: those keep their scale and go through + # _scaled_mm; the cast hooks would upcast them without it, i.e. a wrong weight. + # - `nn.Embedding`: the token embedding table is this model's *input* representation + # and it is large (389M params). Measured on qwen3vl_4b_fp8_scaled, quantizing it + # doubles the encoder's error against bf16 (relative L2 0.0079 -> 0.0163) to save + # 371MiB. That is a bad trade for a model whose entire job is text fidelity. (The + # old fp8_storage path did cast it — this is strictly more accurate than before.) + self._apply_fp8_to_nn_module( + model, + storage_dtype=torch.float8_e4m3fn, + compute_dtype=model_dtype, + skip=lambda _name, module: getattr(module, "weight_scale", None) is not None + or isinstance(module, torch.nn.Embedding), + ) + self._logger.info( + f"FP8 compute enabled for Qwen3-VL encoder '{config.name}': kept {attached} layer(s) " + f"quantized (storage=float8_e4m3fn, matmul=torch._scaled_mm, compute={model_dtype}); " + "remaining layers cast to fp8 storage." + ) + # The layerwise-casting path below exists to *produce* fp8 weights from full-precision + # ones. These already are fp8, and its hooks would cast them to the compute dtype without + # applying weight_scale — a silently wrong weight. It must not run here. + return model + # Keep an fp8 encoder running in fp8 (storage=float8_e4m3fn, per-layer upcast to the compute # dtype during forward) on devices that support fp8 storage. `_should_use_fp8` deliberately # excludes text encoders (and the config has no fp8_storage toggle), so apply the hook-based # casting directly here. This roughly 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 _device_supports_fp8_storage(self._torch_device, self._logger): + if use_fp8_storage: # `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) diff --git a/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py b/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py index 6aad71a71c3..d751a209e60 100644 --- a/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py +++ b/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py @@ -37,6 +37,23 @@ ModelType, SubModelType, ) +from invokeai.backend.quantization.fp8_scaled import ( + TRANSFORMER_KEY_PREFIXES, + attach_fp8_scales, + cast_state_dict, + expand_weight_scale, + extract_comfy_quant_hints, + extract_fp8_scaled_layers, + full_precision_hints_respected, + is_scale_metadata_key, + iter_weight_scale_pairs, + parse_quantization_metadata, + read_safetensors_metadata, + should_keep_fp8_weights, + split_fp8_scaled_layers, + strip_layer_path_prefix, + warn_on_unattached_scales, +) from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.util.devices import TorchDevice @@ -189,6 +206,13 @@ def _read_gguf_metadata_int(path: Path, key: str) -> int | None: return int(value) if isinstance(value, (int, float)) else None +# Wrapper prefixes this loader strips from the state dict. Kept as a named constant because the +# `_quantization_metadata` header names its layers *before* the strip, so the hints have to be +# re-keyed with exactly the same list -- see `_load_text_encoder`. Restating it there would be a +# second copy free to drift. +MISTRAL_KEY_PREFIXES = ("text_encoder.", "language_model.") + + def _strip_known_prefixes(sd: dict[str, Any]) -> dict[str, Any]: """Strip wrapper prefixes used by some FLUX.2 single-file redistributions. @@ -202,7 +226,7 @@ def _strip_known_prefixes(sd: dict[str, Any]) -> dict[str, Any]: out[key] = value continue new_key = key - for prefix in ("text_encoder.", "language_model."): + for prefix in MISTRAL_KEY_PREFIXES: if new_key.startswith(prefix): new_key = new_key[len(prefix) :] break @@ -357,31 +381,25 @@ def _drop_quantization_metadata(sd: dict[str, Any], logger, target_dtype: torch. 24B fp8 encoder that difference is tens of GB — enough to OOM machines that can otherwise load the model. """ - weight_scale_keys = [k for k in sd.keys() if isinstance(k, str) and k.endswith(".weight_scale")] dequantized = 0 - for scale_key in weight_scale_keys: - weight_key = scale_key[: -len(".weight_scale")] + ".weight" - if weight_key not in sd: - continue + for weight_key, scale_key in list(iter_weight_scale_pairs(sd)): weight = sd[weight_key].float() - scale = sd[scale_key].float() - if scale.shape != weight.shape and scale.numel() > 1: - for dim in range(len(weight.shape)): - if dim < len(scale.shape) and scale.shape[dim] != weight.shape[dim]: - block = weight.shape[dim] // scale.shape[dim] - if block > 1: - scale = scale.repeat_interleave(block, dim=dim) + # `expand_weight_scale` rather than a local broadcast: a per-output-channel scale is 1-D of + # length `out`, and `(out, in) * (out,)` aligns on the *last* axis, so it scales input + # channels instead of output channels -- wrong on a square weight, a shape error otherwise. + scale = expand_weight_scale(weight, sd[scale_key].float()) result = weight * scale sd[weight_key] = result.to(target_dtype) if target_dtype is not None else result dequantized += 1 if dequantized: logger.info(f"Dequantized {dequantized} Comfy-Org-style quantized weights") - drop_suffixes = (".weight_scale", ".input_scale", ".scale") + # `is_scale_metadata_key` covers both spellings of the weight and input scales; `.scale` stays + # here because it is this producer's own spelling and is not a scaled-fp8 key. drop_keys = [ k for k in sd.keys() - if isinstance(k, str) and (k.endswith(drop_suffixes) or "comfy_quant" in k or k.startswith("scaled_fp8")) + if isinstance(k, str) and (is_scale_metadata_key(k) or k.endswith(".scale") or k.startswith("scaled_fp8")) ] for k in drop_keys: del sd[k] @@ -922,10 +940,41 @@ def _load_text_encoder(self, config: MistralEncoder_Checkpoint_Config) -> AnyMod target_device = TorchDevice.choose_torch_device() model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) - sd = load_file(Path(config.path)) + model_path = Path(config.path) + sd = load_file(model_path) sd = _strip_known_prefixes(sd) - # Dequantize straight to the compute dtype (per-tensor peak, not whole-dict fp32). - sd = _drop_quantization_metadata(sd, logger, target_dtype=model_dtype) + + # These redistributions are ComfyUI 'scaled fp8': an fp8 weight plus a `weight_scale`. + # Folding the scale doubles the encoder -- 16.8 GiB on disk becomes 32.3 GiB in bf16, which + # does not fit on a 24 GB card even on its own. Keeping the weights quantized is therefore + # not just a speed question here, it decides whether the model loads at all. + # + # Both key rewrites in this loader (`_strip_known_prefixes` above and + # `_convert_for_bare_mistral_model` below) are plain prefix operations, so a sibling + # `.weight_scale` travels with its weight automatically -- no fused projections to split. + keep_fp8 = should_keep_fp8_weights(target_device) + fp8_layers: dict[str, Any] = {} + if keep_fp8: + # This loader strips its own wrapper prefixes on top of the generic ones, so the hints + # need both lists. With only the generic tuple, a `language_model.`-prefixed + # redistribution keeps its hint names while the sd keys lose the prefix, and every + # `full_precision_matrix_mult` is silently ignored. + header_hints = strip_layer_path_prefix( + parse_quantization_metadata(read_safetensors_metadata(model_path, logger)), + prefixes=(*MISTRAL_KEY_PREFIXES, *TRANSFORMER_KEY_PREFIXES), + ) + layer_hints = {**extract_comfy_quant_hints(sd), **header_hints} + fp8_layers = extract_fp8_scaled_layers(sd, layer_hints=layer_hints) + if not fp8_layers: + # Dequantize straight to the compute dtype (per-tensor peak, not whole-dict fp32). + sd = _drop_quantization_metadata(sd, logger, target_dtype=model_dtype) + else: + # `extract_fp8_scaled_layers` removes the keys it interprets, but this producer also + # emits `.scale` and `scaled_fp8*`, which it does not recognize. On the dequantizing + # path `_drop_quantization_metadata` takes those; here nothing did, so they reached + # `load_state_dict` and padded the "ignored N unexpected keys" line. + for k in [k for k in sd if isinstance(k, str) and (k.endswith(".scale") or k.startswith("scaled_fp8"))]: + del sd[k] mistral_config = _build_mistral_config(sd, torch_dtype=model_dtype) logger.info( @@ -940,18 +989,22 @@ def _load_text_encoder(self, config: MistralEncoder_Checkpoint_Config) -> AnyMod for k in [k for k in sd.keys() if isinstance(k, str) and k.startswith("lm_head.")]: del sd[k] - # Cast remaining tensors to compute dtype before loading. Dequantized weights are - # already at model_dtype; this covers the un-quantized ones (norms, embeddings). - for k in list(sd.keys()): - if sd[k].dtype != model_dtype: - sd[k] = sd[k].to(model_dtype) - - # Adapt CausalLM-prefixed keys for bare MistralModel. + # Adapt CausalLM-prefixed keys for bare MistralModel. The recovered scales are keyed on the + # checkpoint's paths, so they need the same `model.` strip or `attach_fp8_scales` resolves + # nothing and every weight stays quantized but unscaled. sd = _convert_for_bare_mistral_model(sd) + fp8_layers = { + (path[len("model.") :] if path.startswith("model.") else path): layer for path, layer in fp8_layers.items() + } with accelerate.init_empty_weights(): model = MistralModel(mistral_config) + # Layers the cast would dequantize anyway are folded here, scale applied, so the cast never + # strips a scale that can no longer be put back. + fp8_layers = split_fp8_scaled_layers(sd, fp8_layers, model_dtype, model=model) + cast_state_dict(sd, model_dtype, keep_fp8=bool(fp8_layers), model=model) + missing, unexpected = model.load_state_dict(sd, strict=False, assign=True) if unexpected: logger.debug(f"Mistral encoder: ignored {len(unexpected)} unexpected keys") @@ -979,6 +1032,19 @@ def _load_text_encoder(self, config: MistralEncoder_Checkpoint_Config) -> AnyMod _strip_final_norm_for_cow(model, config.variant, logger) _warn_if_40_layer_mistral(config.variant, logger) + if fp8_layers: + attached = attach_fp8_scales(model, fp8_layers) + logger.info( + f"Mistral encoder: kept {attached} layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled)" + ) + warn_on_unattached_scales(logger, "Mistral encoder", attached, fp8_layers) + marked = sum(1 for layer in fp8_layers.values() if layer.full_precision_matmul) + if marked and full_precision_hints_respected(): + logger.info( + f"Mistral encoder: {marked} of {len(fp8_layers)} layer(s) are marked " + "full_precision_matrix_mult and will dequantize per forward." + ) + 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 ffda6eaf9cd..62f175f0e55 100644 --- a/invokeai/backend/model_manager/load/model_loaders/z_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/z_image.py @@ -24,7 +24,11 @@ Qwen3Encoder_SDNQ_Config, Qwen3Encoder_SDNQ_Folder_Config, ) -from invokeai.backend.model_manager.load.load_default import ModelLoader, resolve_submodel_path +from invokeai.backend.model_manager.load.load_default import ( + ModelLoader, + _model_declared_skip_patterns, + resolve_submodel_path, +) from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader from invokeai.backend.model_manager.taxonomy import ( @@ -34,6 +38,26 @@ ModelType, SubModelType, ) +from invokeai.backend.quantization.fp8_scaled import ( + QKV_SPLIT_SIDECHANNEL_SUFFIXES, + attach_fp8_scales, + cast_state_dict, + dequantize_fp8_scaled, + expand_weight_scale, + extract_comfy_quant_hints, + extract_fp8_scaled_layers, + full_precision_hints_respected, + is_scale_metadata_key, + iter_weight_scale_pairs, + parse_quantization_metadata, + predict_cast_state_dict_size, + read_safetensors_metadata, + should_keep_fp8_weights, + split_fp8_scaled_layers, + split_qkv_sidechannel, + strip_layer_path_prefix, + warn_on_unattached_scales, +) from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.quantization.sdnq.detection import is_sdnq_folder from invokeai.backend.quantization.sdnq.loaders import raise_on_incomplete_sdnq_load, sdnq_sd_loader @@ -41,6 +65,30 @@ from invokeai.backend.util.devices import TorchDevice +def _remap_z_image_layer_paths(layer_names: Any) -> dict[str, list[str]]: + """Map native Z-Image layer paths to their diffusers equivalents. + + ``_quantization_metadata`` names its layers in the checkpoint's own scheme, but the scales are + extracted after the state dict has been renamed. Rather than restating the rename rules — which + would drift — each name is pushed through the real converter as a lone ``.weight`` entry + and the resulting keys are read back. A fused ``qkv`` maps to *three* diffusers layers, so the + mapping is one-to-many. + """ + mapping: dict[str, list[str]] = {} + for name in layer_names: + if not isinstance(name, str): + continue + try: + # 3 rows so the qkv split is well-defined; the values themselves are never read. + converted = _convert_z_image_gguf_to_diffusers({f"{name}.weight": torch.empty(3, 1)}) + except Exception: + continue + targets = [k[: -len(".weight")] for k in converted if isinstance(k, str) and k.endswith(".weight")] + if targets: + mapping[name] = targets + return mapping + + def _convert_z_image_gguf_to_diffusers(sd: dict[str, Any]) -> dict[str, Any]: """Convert Z-Image GGUF state dict keys to diffusers format. @@ -97,9 +145,15 @@ def _convert_z_image_gguf_to_diffusers(sd: dict[str, Any]) -> dict[str, Any]: prefix = key.rsplit(".attention.qkv.", 1)[0] suffix = key.rsplit(".attention.qkv.", 1)[1] # "weight" or "bias" - # Skip non-weight/bias tensors (e.g., FP8 scale_weight tensors) - # These are quantization metadata and should not be split if suffix not in ("weight", "bias"): + # Quantization side-channel for the fused weight. It has to travel with the split, + # or the recovered scale is keyed on `...attention.qkv`, a module path that no + # longer exists — `attach_fp8_scales` then finds nothing and the three split + # weights stay quantized but *unscaled*, i.e. off by 1/weight_scale. + if suffix in QKV_SPLIT_SIDECHANNEL_SUFFIXES: + for name, part in zip(("to_q", "to_k", "to_v"), split_qkv_sidechannel(key, value), strict=True): + new_sd[f"{prefix}.attention.{name}.{suffix}"] = part + continue new_sd[key] = value continue @@ -415,6 +469,14 @@ def _load_from_singlefile( stripped_sd[key] = value sd = stripped_sd + # Per-layer `full_precision_matrix_mult` hints, from the safetensors header and/or the + # per-tensor `.comfy_quant` markers. The header names layers in the checkpoint's own scheme, + # so it is remapped below; the markers ride along through the key conversion instead. + # The names in the header still carry the checkpoint prefix stripped off `sd` above. + header_hints = strip_layer_path_prefix( + parse_quantization_metadata(read_safetensors_metadata(model_path, self._logger)) + ) + # Check if the state dict is in original format (not diffusers format) # Original format has keys like "x_embedder.weight" instead of "all_x_embedder.2-1.weight" needs_conversion = any(k.startswith("x_embedder.") for k in sd.keys() if isinstance(k, str)) @@ -422,6 +484,10 @@ def _load_from_singlefile( if needs_conversion: # Convert from original format to diffusers format sd = _convert_z_image_gguf_to_diffusers(sd) + path_map = _remap_z_image_layer_paths(header_hints.keys()) + header_hints = { + target: hints for name, hints in header_hints.items() for target in path_map.get(name, [name]) + } # Create an empty model with the default Z-Image config # Z-Image-Turbo uses these default parameters from diffusers @@ -451,7 +517,8 @@ def _load_from_singlefile( # Filter out keys that don't belong to the ZImageTransformer2DModel. # Merged checkpoints (e.g. LoRA-baked models) may bundle text encoder weights # (text_encoders.*) or other non-transformer keys alongside the transformer weights. - # Also filter FP8 quantization metadata (scale_weight, scaled_fp8). + # This runs *before* the scales are extracted so a bundled encoder's own scale keys are + # dropped here rather than being recovered as transformer layers that resolve to nothing. valid_prefixes = ( "all_x_embedder.", "all_final_layer.", @@ -463,23 +530,56 @@ def _load_from_singlefile( "rope_embedder.", ) valid_exact = {"x_pad_token", "cap_pad_token"} - keys_to_remove = [ - k - for k in sd.keys() - if not (k.startswith(valid_prefixes) or k in valid_exact) - or k.endswith(".scale_weight") - or k == "scaled_fp8" - ] + keys_to_remove = [k for k in sd.keys() if not (k.startswith(valid_prefixes) or k in valid_exact)] for k in keys_to_remove: del sd[k] - # Handle memory management and dtype conversion - new_sd_size = sum([ten.nelement() * model_dtype.itemsize for ten in sd.values()]) - self._ram_cache.make_room(new_sd_size) + # ComfyUI 'scaled fp8' (fp8 weight + .weight_scale/.scale_weight). Until now the loader + # deleted those scales and cast the weight — silently producing a weight off by + # 1/weight_scale — and had no way to tell such a checkpoint from a raw fp8 one. + layer_hints = {**extract_comfy_quant_hints(sd), **header_hints} + fp8_layers = extract_fp8_scaled_layers(sd, layer_hints=layer_hints) + + # Handle memory management and dtype conversion. A checkpoint that ships raw fp8 weights + # (fp8 tensors, no weight_scale) keeps them when the fp8 matmul is available — casting them + # here would discard both the VRAM saving and the tensor cores before the model is built. + keep_fp8 = should_keep_fp8_weights(self._torch_device) + if fp8_layers and not keep_fp8: + # Legacy behavior, but now with the scale actually applied: fold it into the weight. + dequantize_fp8_scaled(sd, fp8_layers, model_dtype) + fp8_layers = {} + + # Honor the model's own precision-sensitive list. Z-Image declares + # ["t_embedder", "cap_embedder"], and `TimestepEmbedder.forward` casts its activations to + # `self.mlp[0].weight.dtype` — an fp8 weight there turns the activations fp8 and the forward + # dies in `x.abs()`. Those layers must be dequantized even though the rest stays quantized. + skip_patterns = _model_declared_skip_patterns(model) + # Scaled layers that the cast would dequantize anyway are folded here, scale applied, so + # `cast_state_dict` never strips a scale it cannot put back. + # Reserve before the split, not after: `split_fp8_scaled_layers` dequantizes its unusable + # subset through fp32, so reserving afterwards lets that transient peak land on an + # unreserved cache. `scaled_layers` is what keeps that honest: the split also widens layers + # whose scale layout `scaled_mm` cannot apply, and without the mapping the prediction would + # charge those 1 byte/element and arrive at 2. + self._ram_cache.make_room( + predict_cast_state_dict_size( + sd, + model_dtype, + keep_fp8=keep_fp8, + model=model, + skip_patterns=skip_patterns, + scaled_layers=fp8_layers, + ) + ) - # Convert to target dtype - for k in sd.keys(): - sd[k] = sd[k].to(model_dtype) + fp8_layers = split_fp8_scaled_layers(sd, fp8_layers, model_dtype, model=model, skip_patterns=skip_patterns) + kept = cast_state_dict( + sd, + model_dtype, + keep_fp8=keep_fp8, + model=model, + skip_patterns=skip_patterns, + ) model.load_state_dict(sd, assign=True) # `assign=True` aliases every param to its `sd` tensor, so the dict keeps the whole model @@ -489,13 +589,28 @@ def _load_from_singlefile( # 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. + if fp8_layers: + attached = attach_fp8_scales(model, fp8_layers) + self._logger.info(f"Z-Image: kept {attached} layer(s) in fp8 (scaled fp8 checkpoint, fp8_compute enabled)") + warn_on_unattached_scales(self._logger, "Z-Image", attached, fp8_layers) + marked = sum(1 for layer in fp8_layers.values() if layer.full_precision_matmul) + if marked and full_precision_hints_respected(): + self._logger.info( + f"Z-Image: {marked} of {len(fp8_layers)} layer(s) are marked full_precision_matrix_mult " + "and will dequantize per forward. Set fp8_compute_full_precision_hints=false to run " + "them on the fp8 tensor cores instead." + ) + elif kept: + self._logger.info( + f"Z-Image: kept {kept} raw fp8 weight(s) quantized (no weight_scale in the checkpoint); " + "they will run on the fp8 tensor cores with unit scaling." + ) + + # FP8 *storage* on top. When nothing was kept quantized above, every param is uniform + # `model_dtype` here, so the layerwise cast has one unambiguous compute dtype to restore to. + # When weights *were* kept fp8, `_apply_fp8_layerwise_casting` bails out on its own (and + # says so in the log): its hooks would restore the compute dtype before every forward and + # silently disable the fp8 matmul, for no VRAM saving. model = self._apply_fp8_layerwise_casting(model, config, SubModelType.Transformer) return model @@ -911,6 +1026,37 @@ def _load_control_adapter( @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.Qwen3Encoder, format=ModelFormat.Checkpoint) +def _fold_comfy_scaled_weights(sd: dict[str, Any], dtype: torch.dtype) -> int: + """Fold every ComfyUI-style ``weight_scale`` into its weight, in place. Returns how many. + + ComfyUI stores quantized weights with accompanying scale factors (``layer.weight`` quantized, + ``layer.weight_scale`` the factor, both spellings), so ``dequantized = weight * weight_scale``. + See https://github.com/Comfy-Org/ComfyUI/blob/master/QUANTIZATION.md. + + A named function rather than a loop inside the loader so the scale-axis contract below is + reachable from a test. `expand_weight_scale` handles all three layouts (per-tensor, + per-output-channel, block-wise); the local loop this replaced left a 1-D per-channel scale + untouched, and ``(out, in) * (out,)`` then broadcasts on the *last* axis — scaling input + channels instead of output channels, which is a shape error on a non-square weight and a + silently wrong weight on a square one. + + The multiply runs in float32 for precision but each result is stored as ``dtype`` immediately, + so the whole model is never materialized in float32: holding every dequantized weight at fp32 + until the caller's later cast quadruples the per-parameter cost (4 bytes vs 1 on disk) and + dominates the cold-load RAM peak — enough to swap a 32 GB machine. Same fix as in the FLUX.2 + and Krea-2 loaders. + """ + folded = 0 + for weight_key, scale_key in list(iter_weight_scale_pairs(sd)): + # Float8 needs `.float()`; torch has no direct type promotion for it. + weight_float = sd[weight_key].float() + scale = expand_weight_scale(weight_float, sd[scale_key].float()) + sd[weight_key] = (weight_float * scale).to(dtype) + del weight_float + folded += 1 + return folded + + class Qwen3EncoderCheckpointLoader(ModelLoader): """Class to load single-file Qwen3 Encoder models for Z-Image (safetensors format).""" @@ -975,46 +1121,14 @@ def _load_from_singlefile( # Dequantization formula: dequantized = weight.to(dtype) * weight_scale # Reference: https://github.com/Comfy-Org/ComfyUI/blob/master/QUANTIZATION.md original_key_count = len(sd) - weight_scale_keys = [k for k in sd.keys() if k.endswith(".weight_scale")] - dequantized_count = 0 - - for scale_key in weight_scale_keys: - # Get the corresponding weight key (remove "_scale" suffix) - weight_key = scale_key.replace(".weight_scale", ".weight") - if weight_key in sd: - weight = sd[weight_key] - scale = sd[scale_key] - # Dequantize: convert to float and multiply by scale - # Handle block-wise quantization (e.g., FP4 with block_size=8) - # where scale has shape [weight_dim / block_size, ...] - # Note: Float8 types (e.g., float8_e4m3fn) require .float() instead of .to(torch.float32) - # as PyTorch doesn't support direct type promotion for Float8 types - weight_float = weight.float() - scale = scale.float() - if scale.shape != weight_float.shape and scale.numel() > 1: - # Block-wise quantization: need to expand scale to match weight shape - # Find which dimension differs and repeat scale along that dimension - for dim in range(len(weight_float.shape)): - if dim < len(scale.shape) and scale.shape[dim] != weight_float.shape[dim]: - block_size = weight_float.shape[dim] // scale.shape[dim] - if block_size > 1: - # Repeat scale along this dimension to match weight shape - scale = scale.repeat_interleave(block_size, dim=dim) - # Multiply in float32 for precision, but store the compute dtype immediately so the - # *whole model* is never materialized in float32. Keeping every dequantized weight as - # float32 until the caller's later cast quadruples the per-parameter cost (4 bytes vs - # 1 on disk) and dominates the cold-load RAM peak — enough to swap a 32 GB machine. - # Same fix as in the FLUX.2 and Krea-2 loaders. - sd[weight_key] = (weight_float * scale).to(model_dtype) - del weight_float - dequantized_count += 1 + dequantized_count = _fold_comfy_scaled_weights(sd, model_dtype) if dequantized_count > 0: logger.info(f"Dequantized {dequantized_count} ComfyUI quantized weights") # Filter out ComfyUI quantization metadata keys (comfy_quant, weight_scale) # These are no longer needed after dequantization - comfy_metadata_keys = [k for k in sd.keys() if "comfy_quant" in k or "weight_scale" in k] + comfy_metadata_keys = [k for k in sd.keys() if is_scale_metadata_key(k)] for k in comfy_metadata_keys: del sd[k] if comfy_metadata_keys: diff --git a/invokeai/backend/model_manager/util/model_util.py b/invokeai/backend/model_manager/util/model_util.py index c153129353b..f85c1be490a 100644 --- a/invokeai/backend/model_manager/util/model_util.py +++ b/invokeai/backend/model_manager/util/model_util.py @@ -10,6 +10,7 @@ from invokeai.app.services.config.config_default import get_config from invokeai.backend.model_manager.taxonomy import ClipVariantType +from invokeai.backend.quantization.fp8_scaled import is_scale_metadata_key from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.util.logging import InvokeAILogger @@ -172,9 +173,14 @@ def convert_bundle_to_flux_transformer_checkpoint( if not k.startswith("model.diffusion_model"): keys_to_remove.append(k) # This can be removed in the future if we only want to delete transformer keys continue - if k.endswith("scale"): + if k.endswith("scale") and not is_scale_metadata_key(k): # Scale math must be done at bfloat16 due to our current flux model - # support limitations at inference time + # support limitations at inference time. + # + # fp8 quantization scales are excluded. `.weight_scale` also ends in "scale", so the + # unguarded test folded a scaled-fp8 checkpoint's f32 scales down to bf16 before the + # loader ever saw them — 8 mantissa bits for a value every quantized weight is + # multiplied by. Only the model's own RMSNorm `.scale` parameters belong here. v = v.to(dtype=torch.bfloat16) new_key = k.replace("model.diffusion_model.", "") original_state_dict[new_key] = v diff --git a/invokeai/backend/quantization/fp8_scaled.py b/invokeai/backend/quantization/fp8_scaled.py new file mode 100644 index 00000000000..c14340fc88c --- /dev/null +++ b/invokeai/backend/quantization/fp8_scaled.py @@ -0,0 +1,1033 @@ +"""Shared handling for ComfyUI-style "scaled fp8" checkpoints. + +These checkpoints store each quantized Linear as: + + .weight float8_e4m3fn + .weight_scale float32 (usually a scalar; per-output-channel also occurs) + .input_scale float32 (optional; a calibrated *static* activation scale) + +so that ``w_real ≈ weight.to(float) * weight_scale``. Some producers use ``.scale_weight`` instead +of ``.weight_scale``. A ``_quantization_metadata`` entry in the safetensors header may additionally +mark individual layers with ``full_precision_matrix_mult``, meaning the producer determined that +this layer must not be multiplied in fp8. + +Historically InvokeAI dequantized these to bf16 at load time (three near-identical implementations +in the FLUX.2, Z-Image and Qwen-Image loaders). That throws away both the VRAM saving and the +ability to run the matmul on the fp8 tensor cores. This module keeps the quantization intact so +that :class:`CustomLinear` can decide per forward what to do with it. +""" + +from dataclasses import dataclass +from logging import Logger +from pathlib import Path +from typing import Any, Iterable, Mapping + +import torch + +FP8_DTYPE = torch.float8_e4m3fn + +# Every float8 dtype a checkpoint may store weights in. Scale *recovery* must cover all of them: +# only `float8_e4m3fn` can stay quantized (see `can_stay_quantized`), but an `e5m2` weight still +# needs its `weight_scale` folded in on the way to bf16. Gating extraction on e4m3fn alone dropped +# the scale key and then cast the weight unscaled — off by `1/weight_scale`, silently. +FP8_WEIGHT_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2) + +WEIGHT_SCALE_SUFFIXES = (".weight_scale", ".scale_weight") +# Both spellings occur, exactly as for the weight scale. ComfyUI normalizes `.scale_input` to +# `.input_scale` on load (comfy/utils.py, convert_old_quants); reading only one of them means a +# calibrated activation scale is silently discarded and every forward pays the amax reduction. +INPUT_SCALE_SUFFIXES = (".input_scale", ".scale_input") + +QUANT_METADATA_KEY = "_quantization_metadata" + +# Per-layer marker tensor written by ComfyUI exports that do not use the header entry above. +COMFY_QUANT_SUFFIX = ".comfy_quant" + +# Standalone marker keys some producers emit alongside the tensors. They carry no per-layer data. +STRAY_METADATA_KEYS = ("scaled_fp8",) + + +@dataclass +class Fp8ScaledLayer: + """The quantization parameters recovered for a single Linear.""" + + weight_scale: torch.Tensor + """Scalar, or shape ``(out_features,)`` for per-output-channel quantization.""" + + input_scale: torch.Tensor | None = None + """Calibrated static activation scale, if the checkpoint ships one.""" + + full_precision_matmul: bool = False + """The producer marked this layer as unsafe for an fp8 matmul; only dequantized use is allowed.""" + + def is_per_tensor(self) -> bool: + return self.weight_scale.numel() == 1 + + +def read_safetensors_metadata(path: Path, logger: Logger | None = None) -> dict[str, str] | None: + """Read the safetensors header metadata, or None if it cannot be read. + + Only used to enrich fp8 handling (per-layer ``full_precision_matrix_mult`` hints), so an + unreadable header must not fail the model load. It is warned about rather than swallowed: without + the hints, layers the quantizer marked as unsafe would silently be multiplied in fp8. + """ + try: + from safetensors import safe_open + + with safe_open(path, framework="pt") as f: + return f.metadata() + except Exception as e: + if logger is not None: + logger.warning(f"Could not read safetensors metadata from {path.name} ({e}); fp8 layer hints unavailable.") + return None + + +def parse_quantization_metadata(metadata: Mapping[str, Any] | None) -> dict[str, dict[str, Any]]: + """Parse the safetensors ``_quantization_metadata`` header entry into a per-layer mapping. + + Returns an empty dict when the entry is absent or unparseable - the scales themselves are the + source of truth, the metadata only adds per-layer hints. + """ + if not metadata: + return {} + raw = metadata.get(QUANT_METADATA_KEY) + if not raw: + return {} + if isinstance(raw, (bytes, bytearray)): + raw = raw.decode("utf-8", errors="replace") + if isinstance(raw, str): + import json + + try: + raw = json.loads(raw) + except ValueError: + return {} + if not isinstance(raw, dict): + return {} + layers = raw.get("layers") + return layers if isinstance(layers, dict) else {} + + +def extract_comfy_quant_hints(sd: dict[str, Any]) -> dict[str, dict[str, Any]]: + """Decode per-layer ``.comfy_quant`` markers into the same mapping + :func:`parse_quantization_metadata` produces, popping them out of ``sd``. + + ComfyUI writes the per-layer quantization flags in one of two places. Newer exports use the + safetensors *header* (``_quantization_metadata``); others store one uint8 JSON blob per + quantized layer as an ordinary tensor, e.g.:: + + model.layers.0.mlp.down_proj.comfy_quant + -> {"format": "float8_e4m3fn", "full_precision_matrix_mult": false} + + Both forms carry ``full_precision_matrix_mult`` - the producer's instruction that a layer must + not be multiplied in fp8. A loader that reads only the header therefore *silently ignores that + instruction* on a checkpoint using the per-tensor form, and runs an fp8 matmul the producer + measured as unsafe. Checkpoints using the per-tensor form and marking layers this way exist in + the wild, so both forms must be read. + + Malformed blobs are skipped rather than raised on: the marker is an optimization hint, and the + scales themselves remain the source of truth. + """ + import json + + hints: dict[str, dict[str, Any]] = {} + for key in list(sd.keys()): + if not isinstance(key, str) or not key.endswith(COMFY_QUANT_SUFFIX): + continue + raw = sd.pop(key) + path = key[: -len(COMFY_QUANT_SUFFIX)] + try: + # A uint8 tensor holding UTF-8 JSON, sometimes NUL-padded to a fixed width. + blob = bytes(raw.flatten().tolist()).decode("utf-8", errors="replace").rstrip("\x00") + parsed = json.loads(blob) + except Exception: + continue + if isinstance(parsed, dict): + hints[path] = parsed + return hints + + +# Key prefixes redistributors wrap a transformer in. Loaders strip these off the state dict before +# anything else, so `_quantization_metadata` — which is read from the file and still carries them — +# has to be stripped the same way. +TRANSFORMER_KEY_PREFIXES = ("model.diffusion_model.", "diffusion_model.", "net.") + + +def strip_layer_path_prefix( + layer_hints: Mapping[str, Any], + prefixes: Iterable[str] = TRANSFORMER_KEY_PREFIXES, +) -> dict[str, Any]: + """Re-key ``layer_hints`` as if the checkpoint prefix had been stripped from their names. + + ``_quantization_metadata`` lives in the safetensors header, so its layer names are in the + file's own scheme — ``model.diffusion_model.blocks.0.attn.wq`` — while the state dict has had + that prefix removed before the scales are extracted. A hint whose name still carries the prefix + matches no layer, so ``full_precision_matrix_mult`` is silently ignored and the producer's + "do not multiply this one in fp8" instruction is disregarded: exactly the failure the hint + plumbing exists to prevent. + + Names that carry none of ``prefixes`` are passed through unchanged. Dropping them instead — as + running the names through a strip function that filters by prefix would — turns a + partially-prefixed header into a silently truncated one, and can abort the load. + """ + out: dict[str, Any] = {} + for name, hints in layer_hints.items(): + if isinstance(name, str): + for prefix in prefixes: + if name.startswith(prefix): + name = name[len(prefix) :] + break + out[name] = hints + return out + + +# Every per-layer side-channel suffix that belongs to a module rather than being a tensor of its +# own. Used by the detach/reattach pair below. +LAYER_SIDECHANNEL_SUFFIXES = WEIGHT_SCALE_SUFFIXES + INPUT_SCALE_SUFFIXES + (COMFY_QUANT_SUFFIX,) + + +def detach_layer_sidechannel(sd: dict[str, Any]) -> dict[str, list[tuple[str, Any]]]: + """Pop every per-layer quantization side-channel entry, keyed by the module path it belongs to. + + For loaders that rename checkpoint keys. Key converters are written against ``.weight`` — they + match it as a substring, or test whole keys for equality — so a sibling ``.scale_weight`` or + ``.input_scale`` is *not* carried along, and neither is any scale on a key the converter renames + by equality. The scale is then orphaned under its old path while its weight moves, and + :func:`extract_fp8_scaled_layers` drops it because no fp8 weight sits at the old path any more. + The weight stays quantized with no scale attached and is off by ``1/weight_scale``, in silence: + the layer never enters ``fp8_layers``, so :func:`warn_on_unattached_scales` cannot see it either. + + Take the side channel out of the way, convert, then :func:`reattach_layer_sidechannel`. + """ + detached: dict[str, list[tuple[str, Any]]] = {} + for key in list(sd.keys()): + if not isinstance(key, str): + continue + for suffix in LAYER_SIDECHANNEL_SUFFIXES: + if key.endswith(suffix): + detached.setdefault(key[: -len(suffix)], []).append((suffix, sd.pop(key))) + break + return detached + + +def reattach_layer_sidechannel( + sd: dict[str, Any], + detached: Mapping[str, list[tuple[str, Any]]], + path_map: Mapping[str, str], +) -> list[str]: + """Put detached side-channel entries back under their renamed module paths. + + Returns the module paths that could not be placed. A path with no entry in ``path_map`` had no + destination in the converted state dict — usually because the converter drops that module + outright — so its scale is dropped with it. Returning them rather than swallowing them lets the + caller say so: a *silently* dropped scale is exactly the failure this pair exists to prevent. + + A destination counts as present when *any* tensor sits under it, not specifically a + ``.weight``. Requiring that name assumes every quantizable module stores its parameter as + ``weight``, which is not true of the norms: a producer that quantizes them (the + "quantizes everything" class documented in :func:`can_stay_quantized`) writes ``.scale``, + and the guard would reject the destination for the wrong reason and drop a scale it could have + placed. + """ + present_modules = {key.rsplit(".", 1)[0] for key in sd if isinstance(key, str) and "." in key} + orphaned: list[str] = [] + for path, entries in detached.items(): + destination = path_map.get(path, path) + if destination not in present_modules: + orphaned.append(path) + continue + for suffix, value in entries: + sd[f"{destination}{suffix}"] = value + return orphaned + + +def iter_weight_scale_pairs(sd: Mapping[str, Any]) -> Iterable[tuple[str, str]]: + """Yield ``(weight_key, scale_key)`` for every weight scale in ``sd``, in either spelling. + + For loaders that fold the scale into the weight themselves instead of going through + :func:`extract_fp8_scaled_layers`. Matching only ``.weight_scale`` is the failure that keeps + recurring: a ``.scale_weight`` checkpoint then either loses its scales silently (if the loader + strips both spellings afterwards, leaving the weight off by ``1/weight_scale``) or trips + ``load_state_dict(..., strict=True)`` on the leftover key. Pairs whose ``.weight`` is absent are + skipped, so a stray scale cannot invent one. + """ + for key in list(sd.keys()): + if not isinstance(key, str): + continue + for suffix in WEIGHT_SCALE_SUFFIXES: + if key.endswith(suffix): + weight_key = f"{key[: -len(suffix)]}.weight" + if weight_key in sd: + yield weight_key, key + break + + +# Per-layer quantization side-channel entries that sit next to a fused `qkv.weight` and therefore +# have to be carried through a split of it. The scale spellings are the ones this module accepts; +# the marker is a JSON blob describing the layer, identical for all three parts of the split. +QKV_SPLIT_SIDECHANNEL_SUFFIXES = ("weight_scale", "scale_weight", "input_scale", "scale_input", "comfy_quant") + + +def split_qkv_sidechannel(key: str, value: Any) -> tuple[Any, Any, Any]: + """Split a fused-QKV scale/marker into the parts belonging to Q, K and V. + + A per-tensor scale (and any marker blob) describes the whole fused tensor, so each third + inherits it unchanged. A per-output-channel scale has one entry per row and is split exactly + like the weight. + + Getting this wrong is silent: a scale left on the fused path is keyed on a module the split + model does not have, so `attach_fp8_scales` finds nothing and the three weights stay quantized + but *unscaled* -- off by 1/weight_scale, with no error anywhere. + """ + tensor = torch.as_tensor(value) if hasattr(value, "shape") else value + if not hasattr(tensor, "shape") or tensor.dim() == 0 or tensor.shape[0] == 1: + return (tensor, tensor, tensor) + if tensor.numel() == 1 or key.endswith(("comfy_quant", "input_scale", "scale_input")): + # A marker blob is a 1-D byte string, not a per-channel vector -- never split it. + return (tensor, tensor, tensor) + if tensor.shape[0] % 3 != 0: + raise ValueError( + f"Cannot split fused QKV quantization data '{key}': first dimension ({tensor.shape[0]}) is " + "neither 1 nor divisible by 3, so it matches neither a per-tensor nor a per-channel scale." + ) + third = tensor.shape[0] // 3 + return (tensor[:third], tensor[third : 2 * third], tensor[2 * third :]) + + +def is_scale_metadata_key(key: Any) -> bool: + """Whether ``key`` is fp8 scale/quantization metadata rather than a model tensor. + + Covers both spellings of the weight and input scales plus the marker keys producers emit, so + callers strip exactly what they were able to interpret. + """ + if not isinstance(key, str): + return False + return ( + key.endswith(WEIGHT_SCALE_SUFFIXES) + or key.endswith(INPUT_SCALE_SUFFIXES) + or COMFY_QUANT_SUFFIX.strip(".") in key + or key in STRAY_METADATA_KEYS + ) + + +def _strip_scale_suffix(key: str) -> tuple[str, bool] | None: + """Return ``(module path, is_input_scale)``, or None if ``key`` is not a scale key.""" + for suffix in WEIGHT_SCALE_SUFFIXES: + if key.endswith(suffix): + return key[: -len(suffix)], False + for suffix in INPUT_SCALE_SUFFIXES: + if key.endswith(suffix): + return key[: -len(suffix)], True + return None + + +def _usable_input_scale(scale: torch.Tensor | None) -> torch.Tensor | None: + """A calibrated static activation scale, or None to fall back to per-forward ``amax`` scaling. + + An ``input_scale`` of exactly 1.0 is a placeholder: the producer wrote the field without + calibrating it. Taking it at face value replaces the dynamic scale with *no scaling at all*, so + every activation above the fp8 maximum saturates — far worse than the dynamic path it + suppresses. ComfyUI drops the key in exactly this case (comfy/utils.py, convert_old_quants). + + Non-finite and non-positive scales are rejected for the same reason: they cannot be a valid + divisor, and using one would produce inf/NaN activations instead of a slightly worse image. + + A multi-element ``input_scale`` is dropped rather than raising. ``scaled_mm_linear`` scales the + activations by a single value, so there is nothing to do with a per-channel or per-block + activation scale, and the dynamic per-forward ``amax`` path is a correct fallback. Reshaping it + to a scalar unconditionally turned such a checkpoint into a `RuntimeError: shape '[]' is invalid + for input of size N` at load time -- the one malformed side-channel in this module that neither + skipped nor explained itself. + """ + if scale is None: + return None + scale = scale.float() + if scale.numel() != 1: + return None + scale = scale.reshape(()) + if not torch.isfinite(scale) or scale <= 0 or scale.item() == 1.0: + return None + return scale + + +# OCP Microscaling (MXFP8) stores one E8M0 exponent per 32-element block. safetensors has no E8M0 +# dtype, so producers write the exponents as `uint8`. +# +# We refuse such checkpoints rather than guessing at them. Decoding the byte as `2**(v-127)` and +# expanding it 32-wide is *not* sufficient, established against a real file: the MXFP8 and the +# scaled-fp8 build of `krea2TurboOfficialComfy` share all 174 bf16 tensors bit-for-bit, so the +# scaled build is an exact reference for the same weights -- and against it the decoded weights +# reach a correlation of only 0.60, producing pure noise end to end. The measured per-block scale +# has no monotonic relation to the byte (112 and 116 yield the same true scale), which points at a +# swizzled scale layout rather than a wrong exponent bias. Supporting it means implementing that +# de-swizzle, not adding a constant. +# +# Refusing is the point: with the block-wise expansion in place, such a file otherwise *loads* and +# generates a garbage image with nothing in the log. +_MX_SCALE_DTYPES = (torch.uint8,) + + +def _reject_mx_scale(path: str) -> None: + raise NotImplementedError( + f"'{path}' carries an MXFP8 (OCP Microscaling) block scale, which InvokeAI cannot decode " + "yet: the exponents are stored in a swizzled layout. Use the scaled-fp8 or bf16 build of " + "this checkpoint instead. Loading it anyway would produce a noise image, not a warning." + ) + + +def _normalize_weight_scale(scale: torch.Tensor) -> torch.Tensor: + """Canonical float32 form of a weight scale, preserving its layout. + + Per-tensor scales become 0-d and per-output-channel scales 1-D, so downstream code can branch on + ``numel()``. A scale with more than one dimension is *block-wise* — one entry per block of + weight elements — and is returned with its shape intact: flattening it destroys the block + geometry that :func:`expand_weight_scale` needs to line it back up with the weight. + """ + scale = scale.float() + if scale.numel() == 1: + return scale.reshape(()) + if scale.dim() > 1: + return scale + return scale.flatten() + + +def expand_weight_scale(weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Broadcast ``scale`` to line up with ``weight`` for an elementwise multiply. + + Handles the three layouts producers emit: + + - per-tensor (0-d / single element) — returned unchanged, broadcasting handles it; + - per-output-channel (one entry per row) — reshaped to ``(rows, 1, ...)``; + - block-wise (one entry per block along one or more dims) — each axis is + ``repeat_interleave``d by that axis' block size. + + Without the block-wise case a 2-D scale reaches the multiply as-is and raises a shape error, so + a checkpoint using that layout fails to load outright. That is the layout ComfyUI's own + dequantizer expands, and the FLUX.2 loader used to expand before this module centralized the + logic. + """ + if scale.numel() == 1: + return scale + if scale.dim() <= 1: + if scale.numel() != weight.shape[0]: + # A 1-D scale is per-output-channel by definition, so any other length means the file + # does not describe this weight — most likely a block-wise scale flattened by the + # producer, whose block structure is not recoverable from the tensor alone. Say so: + # left to broadcast, torch raises "size of tensor a (32) must match tensor b (7)" from + # inside the multiply, which names neither the layer nor the file. + raise ValueError( + f"fp8 weight_scale has {scale.numel()} entries but the weight has {weight.shape[0]} output " + "channels; the scale is neither per-tensor nor per-output-channel and cannot be applied. " + "The checkpoint's quantization metadata appears to be malformed." + ) + return scale.reshape(-1, *([1] * (weight.dim() - 1))) + for dim in range(weight.dim()): + if dim < scale.dim() and scale.shape[dim] != weight.shape[dim]: + block = weight.shape[dim] // scale.shape[dim] + if block > 1: + scale = scale.repeat_interleave(block, dim=dim) + return scale + + +def is_matmul_usable_scale(weight: Any, scale: torch.Tensor) -> bool: + """Whether ``scaled_mm_linear`` can apply ``scale`` without materializing the weight. + + It handles exactly two layouts: a per-tensor scalar, which goes into the ``_scaled_mm`` call, + and a per-output-channel vector, which is applied to the *result* (scaling weight row ``j`` + scales output column ``j``). A block-wise scale is separable in neither sense, so such a layer + has to be dequantized up front instead of failing mid-generation inside the kernel. + """ + if scale.numel() == 1: + return True + rows = getattr(weight, "shape", (None,))[0] + return scale.dim() == 1 and scale.numel() == rows + + +def extract_fp8_scaled_layers( + sd: dict[str, Any], + metadata: Mapping[str, Any] | None = None, + layer_hints: Mapping[str, Mapping[str, Any]] | None = None, +) -> dict[str, Fp8ScaledLayer]: + """Pop the quantization side-channel out of ``sd`` and return it keyed by module path. + + ``sd`` is modified in place: scale and marker keys are removed so the remaining state dict + loads cleanly into a model that knows nothing about fp8. The ``.weight`` tensors are left as + float8 - the caller decides whether to keep or dequantize them. + + Only layers whose ``.weight`` is actually float8 are reported; a stray scale without a matching + fp8 weight is dropped rather than silently mis-scaling a bf16 weight. + + ``layer_hints`` overrides the parsed metadata. Loaders that rename checkpoint keys (native → + diffusers) must pass hints keyed by the *renamed* paths, otherwise the per-layer flags - + including ``full_precision_matrix_mult`` - silently match nothing. + """ + layer_meta = dict(layer_hints) if layer_hints is not None else parse_quantization_metadata(metadata) + + weight_scales: dict[str, torch.Tensor] = {} + input_scales: dict[str, torch.Tensor] = {} + for key in list(sd.keys()): + if not isinstance(key, str): + continue + parsed = _strip_scale_suffix(key) + if parsed is None: + continue + path, is_input_scale = parsed + if is_input_scale: + input_scales[path] = sd.pop(key) + else: + weight_scales[path] = sd.pop(key) + + for key in list(sd.keys()): + if isinstance(key, str) and (key in STRAY_METADATA_KEYS or "comfy_quant" in key): + del sd[key] + + layers: dict[str, Fp8ScaledLayer] = {} + for path, scale in weight_scales.items(): + weight = sd.get(f"{path}.weight") + if weight is None or getattr(weight, "dtype", None) not in FP8_WEIGHT_DTYPES: + # A scale without an fp8 weight means the weight was already dequantized (or the key + # naming does not line up). Applying the scale later would corrupt it, so drop it. + continue + if getattr(scale, "dtype", None) in _MX_SCALE_DTYPES: + _reject_mx_scale(path) + hints = layer_meta.get(path, {}) + layers[path] = Fp8ScaledLayer( + weight_scale=_normalize_weight_scale(scale), + input_scale=_usable_input_scale(input_scales.get(path)), + full_precision_matmul=bool(hints.get("full_precision_matrix_mult", False)), + ) + return layers + + +def dequantize_fp8_scaled( + sd: dict[str, Any], + layers: Mapping[str, Fp8ScaledLayer], + dtype: torch.dtype = torch.bfloat16, +) -> dict[str, Any]: + """Fold the scales back into the weights, producing a plain ``dtype`` state dict. + + This is the legacy behavior, kept as the fallback for models/devices that cannot use the fp8 + path. The multiply runs in float32 for precision but the result is stored as ``dtype`` + immediately, so a cold load never holds the whole model in float32. + """ + for path, layer in layers.items(): + key = f"{path}.weight" + weight = sd.get(key) + if weight is None: + continue + weight = weight.float() + sd[key] = (weight * expand_weight_scale(weight, layer.weight_scale)).to(dtype) + return sd + + +def attach_fp8_scales( + model: torch.nn.Module, + layers: Mapping[str, Fp8ScaledLayer], + module_paths: Iterable[str] | None = None, +) -> int: + """Register the recovered scales as buffers on the matching modules. + + ``CustomLinear`` looks for ``weight_scale`` / ``input_scale`` buffers and a + ``_fp8_full_precision_matmul`` flag. The buffers are non-persistent: they must never end up in + ``state_dict()`` output, or a re-save would produce a checkpoint that gets scaled twice. + + The producer's ``full_precision_matrix_mult`` markers are applied here rather than dropped at + parse time, so the recovered metadata stays inspectable, and are suppressed when the user has + turned the hints off (see :func:`full_precision_hints_respected`). + + Returns the number of modules that were annotated. A caller that expects *every* recovered + layer to be annotated should check that count with :func:`warn_on_unattached_scales`. + """ + wanted = set(module_paths) if module_paths is not None else None + respect_hints = full_precision_hints_respected() + count = 0 + for path, layer in layers.items(): + if wanted is not None and path not in wanted: + continue + try: + module = model.get_submodule(path) if path else model + except AttributeError: + continue + weight = getattr(module, "weight", None) + if weight is None or weight.dtype != FP8_DTYPE: + continue + module.register_buffer("weight_scale", layer.weight_scale.to(weight.device), persistent=False) + if layer.input_scale is not None: + module.register_buffer("input_scale", layer.input_scale.to(weight.device), persistent=False) + module._fp8_full_precision_matmul = layer.full_precision_matmul and respect_hints + count += 1 + return count + + +def warn_on_unattached_scales(logger: Logger, what: str, attached: int, layers: Mapping[str, Any]) -> None: + """Complain when :func:`attach_fp8_scales` annotated fewer modules than there were layers. + + Every recovered layer should reach a module: :func:`split_fp8_scaled_layers` has already folded + the ones that cannot stay quantized. A shortfall therefore means a scale went nowhere, and a + weight is now off by ``1/weight_scale`` — visually a broken or washed-out generation, with + nothing in the log to point at it. Loaders otherwise report ``attached`` as if it were the whole + story, which reads as success. + """ + missing = len(layers) - attached + if missing > 0: + logger.warning( + f"{what}: {missing} of {len(layers)} scaled fp8 layer(s) did not receive their weight_scale. " + "Those weights are quantized but unscaled, which will degrade output. This is a bug — " + "please report the checkpoint." + ) + + +# ----------------------------------------------------------------------------------- runtime path + +FP8_MAX = torch.finfo(FP8_DTYPE).max + +# torch._scaled_mm requires every GEMM dimension to be a multiple of 16. +_MM_ALIGNMENT = 16 + +_fp8_mm_supported: dict[int, bool] = {} + +# fp8 compute quantizes the *activations* as well, so it changes numerics: an existing install would +# start producing different images at the same seed. It is therefore opt-in (`fp8_compute` in +# invokeai.yaml) for one release before becoming the default. +# +# The same flag also decides whether scaled fp8 checkpoints stay quantized at load. Keeping them +# quantized without the fp8 matmul would halve VRAM but make generation *slower* (the dequantize +# round trip costs more than it saves), so the two must be switched together. +_fp8_matmul_override: bool | None = None + + +def set_fp8_matmul_enabled(enabled: bool | None) -> None: + """Override the configured setting process-wide. Pass ``None`` to revert to the config value.""" + global _fp8_matmul_override + _fp8_matmul_override = enabled + + +def is_fp8_matmul_enabled() -> bool: + if _fp8_matmul_override is not None: + return _fp8_matmul_override + try: + from invokeai.app.services.config.config_default import get_config + + return bool(get_config().fp8_compute) + except Exception: + # Backend code may run outside a configured app (scripts, tests). Default to the safe path. + return False + + +_full_precision_hints_override: bool | None = None + + +def set_full_precision_hints_respected(respected: bool | None) -> None: + """Override the configured setting process-wide. Pass ``None`` to revert to the config value.""" + global _full_precision_hints_override + _full_precision_hints_override = respected + + +def full_precision_hints_respected() -> bool: + """Whether ``full_precision_matrix_mult`` markers are obeyed. + + Honoring a marker means that layer dequantizes on every forward instead of using the fp8 tensor + cores, and that is not cheap: on a checkpoint that marks the attention output, gate and FFN-down + projections (a common choice) the marked layers can be ~40% of the quantized weights, and + measurably erase most of the fp8_compute speedup. Whether that trade is worth it depends on how + much the producer's flags actually buy in a given checkpoint, so it is a user-facing setting + rather than a hard-coded policy. + """ + if _full_precision_hints_override is not None: + return _full_precision_hints_override + try: + from invokeai.app.services.config.config_default import get_config + + return bool(get_config().fp8_compute_full_precision_hints) + except Exception: + # Backend code may run outside a configured app (scripts, tests). Default to obeying them. + return True + + +# Wordings torch uses when `_scaled_mm` is genuinely unavailable on the hardware, as opposed to +# having failed under load. Kept narrow on purpose: anything not matched here is treated as +# transient and re-probed, which is the safe direction (see `_probe_fp8_matmul`). +_UNSUPPORTED_MM_MARKERS = ("compute capability", "not supported", "no kernel image", "not implemented") + + +def _is_definitive_unsupported_error(exc: BaseException) -> bool: + """Whether ``exc`` says the device cannot do fp8 matmul at all, rather than not right now.""" + message = str(exc).lower() + return any(marker in message for marker in _UNSUPPORTED_MM_MARKERS) + + +def _probe_fp8_matmul(index: int) -> bool | None: + """Run one minimal ``_scaled_mm`` on device ``index`` and report whether it worked. + + Asking the device is the only reliable test. ``get_device_capability`` reports the *gfx arch* + on ROCm, not an SM version, so an RDNA3 card (gfx1100) answers ``(11, 0)`` and sails past a + ``>= (8, 9)`` check — then every forward raises ``torch._scaled_mm is only supported on CUDA + devices with compute capability >= 9.0 or 8.9, or ROCm MI300+``. The whole point of the + capability gate is to *fall back* rather than raise mid-generation, so it must not itself be a + guess. The probe is one 16x16 matmul, run once per device and cached. + + Returns ``None`` when the probe could not be *carried out* rather than having established that + the operation is unsupported. The probe runs during a model load, i.e. under real VRAM pressure + and alongside whatever else the driver is doing, and caching a momentary failure as "this GPU + cannot do fp8" would disable the fp8 matmul for the rest of the process. Same reasoning as + `_device_supports_fp8_storage`, which also refuses to cache a transient failure. + + A ``RuntimeError`` is therefore treated as inconclusive *unless* it names the support + constraint. Listing the transient wordings instead would be the wrong way round: an OOM and a + cuBLAS workspace failure are only two of the ways a loaded machine can fail this call, and every + wording not on such a list would be cached as permanent. The definitive answer has one stable + shape — torch says the op needs a given compute capability — so match that and treat the open + set as inconclusive. + + An inconclusive answer costs one 16x16 matmul on the next load; a wrongly cached one costs the + fp8 path for the lifetime of the process. + """ + try: + device = torch.device("cuda", index) + # Column-major right operand, i.e. exactly the layout `scaled_mm_linear` feeds it. + lhs = torch.zeros((_MM_ALIGNMENT, _MM_ALIGNMENT), device=device, dtype=FP8_DTYPE) + rhs = torch.zeros((_MM_ALIGNMENT, _MM_ALIGNMENT), device=device, dtype=FP8_DTYPE).t() + scale = torch.ones((1, 1), device=device, dtype=torch.float32) + torch._scaled_mm(lhs, rhs, scale, scale, out_dtype=torch.bfloat16) + except torch.OutOfMemoryError: + return None + except RuntimeError as e: + return False if _is_definitive_unsupported_error(e) else None + except Exception: + # Not a RuntimeError: a missing op or a rejected dtype, i.e. structural rather than + # situational. Those do not become true on a less busy machine. + return False + return True + + +def device_supports_fp8_matmul(device: torch.device) -> bool: + """Whether ``torch._scaled_mm`` can run on this device (Ada/SM 8.9 and newer, or MI300+).""" + if device.type != "cuda" or not torch.cuda.is_available(): + return False + index = device.index if device.index is not None else torch.cuda.current_device() + cached = _fp8_mm_supported.get(index) + if cached is not None: + return cached + # The capability check is only a cheap pre-filter that spares older CUDA cards the probe; + # it is deliberately not trusted on its own (see `_probe_fp8_matmul`). + if not (torch.version.hip is not None or torch.cuda.get_device_capability(index) >= (8, 9)): + _fp8_mm_supported[index] = False + return False + probed = _probe_fp8_matmul(index) + if probed is None: + # Inconclusive: answer this call conservatively but leave the cache empty so the next load + # re-probes instead of the process being stuck without fp8 after one transient OOM. + return False + _fp8_mm_supported[index] = probed + return probed + + +def reset_fp8_matmul_support_cache() -> None: + """Forget the probed per-device support. For tests; the answer cannot change at runtime.""" + _fp8_mm_supported.clear() + + +def should_keep_fp8_weights(device: torch.device) -> bool: + """Whether fp8 weights in a checkpoint should survive the load instead of being dequantized. + + Only true when the fp8 matmul is both enabled and usable, because keeping weights quantized + without it is the worst of both worlds: the same VRAM as fp8 but a dequantize round trip on + every forward (measured slower than plain bf16). + """ + return is_fp8_matmul_enabled() and device_supports_fp8_matmul(device) + + +def _is_fp8_matmul_weight(key: str, tensor: Any, model: torch.nn.Module | None) -> bool: + """Whether this state-dict entry is a weight `scaled_mm_linear` can actually consume. + + Only the ``.weight`` of an ``nn.Linear`` qualifies. This matters because checkpoints exist that + quantize *everything* — biases, norm weights, even learned pad tokens. Keeping those in fp8 does + not save anything worth having and actively breaks inference: an fp8 norm or pad token flows + into the activations, and the next Linear then receives an fp8 *input*, which dies in + ``x.abs()`` with ``"abs_cuda" not implemented for 'Float8_e4m3fn'``. Observed on a Z-Image + checkpoint where 243 of 453 fp8 tensors were 1-D. + """ + if not key.endswith(".weight") or getattr(tensor, "dim", None) is None or tensor.dim() < 2: + return False + if model is None: + # No model to resolve against: the 2-D + `.weight` shape test above is the safe subset. + return True + try: + module = model.get_submodule(key[: -len(".weight")]) + except AttributeError: + return False + return isinstance(module, torch.nn.Linear) + + +def can_stay_quantized( + key: str, + tensor: Any, + model: torch.nn.Module | None, + skip_patterns: Iterable[str] = (), +) -> bool: + """Whether this state-dict entry may be left in fp8 by :func:`cast_state_dict`. + + Single source of truth for that decision: the loaders reserve RAM against it + (:func:`predict_cast_state_dict_size`) and decide which scaled layers survive with it + (:func:`split_fp8_scaled_layers`), so the three must never drift apart. + + Only ``float8_e4m3fn`` qualifies. ``float8_e5m2`` is always cast because + :func:`scaled_mm_linear` cannot use it as the weight operand on Ada, so keeping it quantized + would buy VRAM at the cost of a per-forward dequantize. + + ``skip_patterns`` are substrings of the state-dict key whose weights must be dequantized even + when the rest stays fp8. Pass the model's ``_skip_layerwise_casting_patterns``: diffusers uses + it to mark precision-sensitive modules, and some of them *read their own weight's dtype and cast + their activations to it*. Z-Image's ``TimestepEmbedder.forward`` does exactly that + (``t_freq.to(self.mlp[0].weight.dtype)``), so leaving its weight in fp8 hands the next Linear an + fp8 activation and the forward dies in ``x.abs()`` with + ``"abs_cuda" not implemented for 'Float8_e4m3fn'``. + """ + return ( + getattr(tensor, "dtype", None) is FP8_DTYPE + and _is_fp8_matmul_weight(key, tensor, model) + and not any(pattern in key for pattern in skip_patterns) + ) + + +def scaled_layer_for(key: Any, scaled_layers: Mapping[str, "Fp8ScaledLayer"] | None) -> "Fp8ScaledLayer | None": + """The recovered scale for the module whose ``.weight`` this key is, if there is one.""" + if not scaled_layers or not isinstance(key, str) or not key.endswith(".weight"): + return None + return scaled_layers.get(key[: -len(".weight")]) + + +def survives_split_and_cast( + key: str, + tensor: Any, + model: torch.nn.Module | None, + skip_patterns: Iterable[str] = (), + layer: "Fp8ScaledLayer | None" = None, +) -> bool: + """Whether this entry is *still* fp8 after both :func:`split_fp8_scaled_layers` and + :func:`cast_state_dict` have run over it. + + :func:`can_stay_quantized` answers only for the cast. The split applies one filter more — a + scale layout :func:`scaled_mm_linear` cannot use — so a block-wise-scaled 2-D Linear weight + passes ``can_stay_quantized`` and is nonetheless widened to ``dtype``. Predicting RAM with the + weaker predicate therefore charges 1 byte/element for a tensor that ends up at 2, and the loaders + reserve half of what they need. This is the predicate both the split and the prediction use, so + the two cannot drift. + + ``layer`` is None for a raw fp8 weight (no ``weight_scale`` at all); there is no layout to + reject, and the tensor-core path takes it with unit scaling. + """ + if not can_stay_quantized(key, tensor, model, skip_patterns): + return False + return layer is None or is_matmul_usable_scale(tensor, layer.weight_scale) + + +def _is_castable_float(tensor: Any) -> bool: + """Whether ``tensor`` is a floating-point payload that may be cast to the compute dtype.""" + is_floating_point = getattr(tensor, "is_floating_point", None) + if not callable(is_floating_point): + return False + try: + return bool(is_floating_point()) + except Exception: + return False + + +def cast_state_dict( + sd: dict[str, Any], + dtype: torch.dtype, + *, + keep_fp8: bool, + model: torch.nn.Module | None = None, + skip_patterns: Iterable[str] = (), +) -> int: + """Cast every tensor in ``sd`` to ``dtype`` in place, optionally leaving fp8 weights quantized. + + Loaders historically cast the whole state dict unconditionally, which silently dequantizes a + checkpoint that ships raw fp8 weights (fp8 tensors with no ``weight_scale`` alongside them) — + the VRAM saving and the tensor cores are both thrown away before the model is ever built. + + A *scaled* fp8 weight must never reach this function still carrying an unapplied scale: the + plain ``tensor.to(dtype)`` below drops the scale silently, leaving the weight off by + ``1/weight_scale``. Run :func:`split_fp8_scaled_layers` first — it folds the scale into exactly + those layers this function would cast. + + Returns the number of tensors left in fp8. + """ + patterns = tuple(skip_patterns) + kept = 0 + for key in sd: + tensor = sd[key] + if keep_fp8 and can_stay_quantized(key, tensor, model, patterns): + kept += 1 + continue + if not _is_castable_float(tensor): + # Integer payloads (embedding indices, packed buffers) are not weights and must keep + # their dtype. Loaders used to guard this themselves; centralizing it here means a + # loader that switches to `cast_state_dict` does not silently lose the guard. + continue + sd[key] = tensor.to(dtype) + return kept + + +def predict_cast_state_dict_size( + sd: Mapping[str, Any], + dtype: torch.dtype, + *, + keep_fp8: bool, + model: torch.nn.Module | None = None, + skip_patterns: Iterable[str] = (), + scaled_layers: Mapping[str, Fp8ScaledLayer] | None = None, +) -> int: + """Bytes the state dict will occupy once the split and :func:`cast_state_dict` have run over it. + + Loaders call this to size their ``make_room()`` reservation. Charging 1 byte/element for every + fp8 tensor is wrong in the direction that hurts: only 2-D ``nn.Linear`` weights outside the skip + patterns stay quantized, and everything else — biases, norms, learned pad tokens, the + deliberately-dequantized precision-sensitive Linears — lands at ``dtype.itemsize``. On a + checkpoint that quantized all 453 of its tensors that under-count is most of the difference. + + ``scaled_layers`` is the mapping the caller is about to hand :func:`split_fp8_scaled_layers`. + Pass it: the reservation is made *before* the split (so the split's own fp32 transient lands on + a reserved cache), and the split widens every layer whose scale layout ``scaled_mm`` cannot + apply. Without the mapping those layers are predicted at 1 byte/element and arrive at 2 — on a + block-wise-scaled checkpoint that is the entire quantized-Linear set. + """ + patterns = tuple(skip_patterns) + total = 0 + for key, tensor in sd.items(): + stays = keep_fp8 and survives_split_and_cast(key, tensor, model, patterns, scaled_layer_for(key, scaled_layers)) + if stays or not _is_castable_float(tensor): + total += tensor.nelement() * tensor.element_size() + else: + total += tensor.nelement() * dtype.itemsize + return total + + +def split_fp8_scaled_layers( + sd: dict[str, Any], + layers: Mapping[str, Fp8ScaledLayer], + dtype: torch.dtype, + *, + model: torch.nn.Module | None = None, + skip_patterns: Iterable[str] = (), +) -> dict[str, Fp8ScaledLayer]: + """Dequantize the scaled layers that cannot stay quantized; return the ones that can. + + Every filter that keeps a weight *out* of fp8 — a skip pattern, a weight that is not a 2-D + ``nn.Linear.weight`` — is a filter that would otherwise let :func:`cast_state_dict` do a plain + ``.to(dtype)`` on a scaled weight, i.e. drop its ``weight_scale`` and leave the weight off by + ``1/weight_scale``. :func:`attach_fp8_scales` cannot repair that afterwards: it skips any module + whose weight is no longer fp8, so the scale is lost for good. Krea-2 hits this on ordinary + ComfyUI exports, where ``time_embed.linear_1/linear_2`` are quantized like any other Linear and + match the model's ``time_embed`` skip pattern. + + So the filters are applied here instead, *before* the cast, and the affected layers go through + :func:`dequantize_fp8_scaled`, which applies the scale properly. They are then dropped from the + returned mapping — they are no longer fp8, so there is nothing left to attach. + + A scale layout :func:`scaled_mm_linear` cannot apply — block-wise, or a vector that does not + match the weight's row count — is dequantized here too. Left quantized it would fail inside the + kernel mid-generation instead. That filter lives here rather than in :func:`can_stay_quantized` + because it needs the recovered scale, which the cast never sees. The RAM prediction runs + *before* this function, so it has to account for the same widening — it does, via the shared + :func:`survives_split_and_cast`, provided the caller passes it ``scaled_layers``. + + ``float8_e5m2`` layers land here by way of :func:`can_stay_quantized`, which admits only + ``float8_e4m3fn`` — they are dequantized *with* their scale applied rather than losing it. + """ + patterns = tuple(skip_patterns) + usable: dict[str, Fp8ScaledLayer] = {} + unusable: dict[str, Fp8ScaledLayer] = {} + for path, layer in layers.items(): + key = f"{path}.weight" + tensor = sd.get(key) + if tensor is not None and survives_split_and_cast(key, tensor, model, patterns, layer): + usable[path] = layer + else: + unusable[path] = layer + if unusable: + dequantize_fp8_scaled(sd, unusable, dtype) + return usable + + +def count_fp8_weights(model: torch.nn.Module) -> int: + """Number of parameters already stored as ``float8_e4m3fn``. + + Used to tell a checkpoint that arrived quantized apart from one this loader is about to + quantize itself — the two must not both happen (see `ModelLoader._apply_fp8_layerwise_casting`). + """ + return sum(1 for p in model.parameters() if p.dtype is FP8_DTYPE) + + +def dequantize_weight(weight: torch.Tensor, weight_scale: torch.Tensor | None, dtype: torch.dtype) -> torch.Tensor: + """Cast an fp8 weight up to ``dtype``, applying its scale if it has one. + + Used by the fallback path. Casting *without* the scale - which is what a plain ``.to(dtype)`` + does - silently produces a wrongly-scaled weight, so every dequantization of a scaled fp8 + weight must go through here. + """ + out = weight.to(dtype) + if weight_scale is None: + return out + scale = weight_scale.to(device=out.device, dtype=dtype) + return out * expand_weight_scale(out, scale) + + +def scaled_mm_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor | None, + bias: torch.Tensor | None = None, + input_scale: torch.Tensor | None = None, +) -> torch.Tensor: + """``F.linear`` executed on the fp8 tensor cores via ``torch._scaled_mm``. + + ``weight`` stays float8 and is never materialized in a wider dtype - that is where both the + VRAM and the speed come from. + + Scaling notes: + - The activation scale is per-tensor. Ada rejects per-row activation scaling outright, and + a static ``input_scale`` from the checkpoint is used when available (calibrated, and it saves + the per-forward ``amax`` reduction). + - A per-output-channel weight scale cannot be handed to ``_scaled_mm`` on Ada either, but it is + separable: scaling row ``j`` of the weight scales output column ``j``, so it is applied to the + result instead. Per-tensor scales go straight into the kernel. + """ + orig_shape = input.shape + x = input.reshape(-1, orig_shape[-1]) + + # The transpose must be produced here rather than cached: a stored transposed view stops being a + # view the moment the tensor is moved between devices (which partial loading does constantly), + # silently doubling the weight memory. + weight_t = weight.t() + + pad = (-x.shape[0]) % _MM_ALIGNMENT + if pad: + x = torch.nn.functional.pad(x, (0, 0, 0, pad)) + + if input_scale is not None: + x_scale = input_scale.to(device=x.device, dtype=torch.float32).reshape(1, 1) + x_fp8 = (x / x_scale.to(x.dtype)).clamp(-FP8_MAX, FP8_MAX).to(FP8_DTYPE) + else: + amax = x.abs().amax().clamp(min=1e-12) + x_scale = (amax / FP8_MAX).float().reshape(1, 1) + x_fp8 = (x / x_scale.to(x.dtype)).to(FP8_DTYPE) + + per_tensor = weight_scale is not None and weight_scale.numel() == 1 + if per_tensor: + w_scale = weight_scale.to(device=x.device, dtype=torch.float32).reshape(1, 1) + else: + w_scale = torch.ones(1, 1, device=x.device, dtype=torch.float32) + + out = torch._scaled_mm(x_fp8.contiguous(), weight_t, x_scale, w_scale, out_dtype=input.dtype) + + if pad: + out = out[: x.shape[0] - pad] + if weight_scale is not None and not per_tensor: + out = out * weight_scale.to(device=out.device, dtype=out.dtype).reshape(1, -1) + + out = out.reshape(*orig_shape[:-1], weight.shape[0]) + if bias is not None: + out = out + bias.to(dtype=out.dtype) + return out diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index cc470ea5c3d..218105f8215 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -49500,6 +49500,18 @@ "description": "Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.", "default": true }, + "fp8_compute": { + "type": "boolean", + "title": "Fp8 Compute", + "description": "Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).", + "default": false + }, + "fp8_compute_full_precision_hints": { + "type": "boolean", + "title": "Fp8 Compute Full Precision Hints", + "description": "Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.", + "default": true + }, "ram": { "anyOf": [ { @@ -49881,7 +49893,7 @@ "additionalProperties": false, "type": "object", "title": "InvokeAIAppConfig", - "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set.\n http_compression_level: Compression level for gzipped HTTP API responses. 0 disables response compression entirely, 1 is fastest, 9 (the default) is smallest. Compression runs on the event loop and blocks the whole server while it works, and level 9 costs about 5.5x the time of level 1 for 0.4 percentage points of extra compression, so lowering this makes the app noticeably more responsive on large libraries. Set to 0 when a reverse proxy already compresses responses." + "description": "Invoke's global app configuration.\n\nTypically, you won't need to interact with this class directly. Instead, use the `get_config` function from `invokeai.app.services.config` to get a singleton config object.\n\nAttributes:\n host: IP address to bind to. Use `0.0.0.0` to serve to your local network.\n port: Port to bind to.\n allow_origins: Allowed CORS origins.\n allow_credentials: Allow CORS credentials.\n allow_methods: Methods allowed for CORS.\n allow_headers: Headers allowed for CORS.\n ssl_certfile: SSL certificate file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n ssl_keyfile: SSL key file for HTTPS. See https://www.uvicorn.dev/settings/#https.\n log_tokenization: Enable logging of parsed prompt tokens.\n patchmatch: Enable patchmatch inpaint code.\n models_dir: Path to the models directory.\n convert_cache_dir: Path to the converted models cache directory (DEPRECATED, but do not delete because it is needed for migration from previous versions).\n download_cache_dir: Path to the directory that contains dynamically downloaded models.\n legacy_conf_dir: Path to directory of legacy checkpoint config files.\n db_dir: Path to InvokeAI databases directory.\n outputs_dir: Path to directory for outputs.\n image_subfolder_strategy: Strategy for organizing images into subfolders. 'flat' stores all images in a single folder. 'date' organizes by YYYY/MM/DD. 'type' organizes by image category. 'hash' uses first 2 characters of UUID for filesystem performance.
Valid values: `flat`, `date`, `type`, `hash`\n custom_nodes_dir: Path to directory for custom nodes.\n style_presets_dir: Path to directory for style presets.\n workflow_thumbnails_dir: Path to directory for workflow thumbnails.\n log_handlers: Log handler. Valid options are \"console\", \"file=\", \"syslog=path|address:host:port\", \"http=\".\n log_format: Log format. Use \"plain\" for text-only, \"color\" for colorized output, \"legacy\" for 2.3-style logging and \"syslog\" for syslog-style.
Valid values: `plain`, `color`, `syslog`, `legacy`\n log_level: Emit logging messages at this level or higher.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n log_sql: Log SQL queries. `log_level` must be `debug` for this to do anything. Extremely verbose.\n log_level_network: Log level for network-related messages. 'info' and 'debug' are very verbose.
Valid values: `debug`, `info`, `warning`, `error`, `critical`\n use_memory_db: Use in-memory database. Useful for development.\n dev_reload: Automatically reload when Python sources are changed. Does not reload node definitions.\n profile_graphs: Enable graph profiling using `cProfile`.\n profile_prefix: An optional prefix for profile output files.\n profiles_dir: Path to profiles output directory.\n max_cache_ram_gb: The maximum amount of CPU RAM to use for model caching in GB. If unset, the limit will be configured based on the available RAM. In most cases, it is recommended to leave this unset.\n max_cache_vram_gb: The amount of VRAM to use for model caching in GB. If unset, the limit will be configured based on the available VRAM and the device_working_mem_gb. In most cases, it is recommended to leave this unset.\n log_memory_usage: If True, a memory snapshot will be captured before and after every model cache operation, and the result will be logged (at debug level). There is a time cost to capturing the memory snapshots, so it is recommended to only enable this feature if you are actively inspecting the model cache's behaviour.\n model_cache_keep_alive_min: How long to keep models in cache after last use, in minutes. A value of 0 (the default) means models are kept in cache indefinitely. If no model generations occur within the timeout period, the model cache is cleared using the same logic as the 'Clear Model Cache' button.\n device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value.\n enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM.\n keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high.\n fp8_compute: Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM).\n fp8_compute_full_precision_hints: Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled.\n ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.\n lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.\n pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.\n device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.
Valid values: `auto`, `cpu`, `cuda`, `mps`, `xpu`, `cuda:N`, `xpu:N` (where N is a device number)\n precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.
Valid values: `auto`, `float16`, `bfloat16`, `float32`\n sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.\n wan_memory_optimization: Enable experimental Wan memory optimizations at the cost of slower generation.\n pid_memory_optimization: Enable experimental PiD decode memory optimizations. Roughly halves the peak activation memory of a PiD decode; in exchange the decoded image changes slightly, because neither the chunked pixel pathway nor the float32 sampler intermediates are bit-exact with the default path.\n attention_type: Attention type.
Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`\n attention_slice_size: Slice size, valid when attention_type==\"sliced\".
Valid values: `auto`, `balanced`, `max`, `1`, `2`, `3`, `4`, `5`, `6`, `7`, `8`\n force_tiled_decode: Whether to enable tiled VAE decode (reduces memory consumption with some performance penalty).\n pil_compress_level: The compress_level setting of PIL.Image.save(), used for PNG encoding. All settings are lossless. 0 = no compression, 1 = fastest with slightly larger filesize, 9 = slowest with smallest filesize. 1 is typically the best setting.\n max_queue_size: Maximum number of items in the session queue.\n session_queue_mode: Session queue mode. Use 'FIFO' for traditional first-in-first-out, or 'round_robin' to serve each user's jobs in turn. In single-user mode, FIFO is always used regardless of this setting.
Valid values: `FIFO`, `round_robin`\n clear_queue_on_startup: Empties session queue on startup. If true, disables `max_queue_history`.\n max_queue_history: Keep the last N completed, failed, and canceled queue items. Older items are deleted on startup. Set to 0 to prune all terminal items. Ignored if `clear_queue_on_startup` is true.\n allow_nodes: List of nodes to allow. Omit to allow all.\n deny_nodes: List of nodes to deny. Omit to deny none.\n node_cache_size: How many cached nodes to keep in memory.\n hashing_algorithm: Model hashing algorthim for model installs. 'blake3_multi' is best for SSDs. 'blake3_single' is best for spinning disk HDDs. 'random' disables hashing, instead assigning a UUID to models. Useful when using a memory db to reduce model installation time, or if you don't care about storing stable hashes for models. Alternatively, any other hashlib algorithm is accepted, though these are not nearly as performant as blake3.
Valid values: `blake3_multi`, `blake3_single`, `random`, `md5`, `sha1`, `sha224`, `sha256`, `sha384`, `sha512`, `blake2b`, `blake2s`, `sha3_224`, `sha3_256`, `sha3_384`, `sha3_512`, `shake_128`, `shake_256`\n remote_api_tokens: List of regular expression and token pairs used when downloading models from URLs. The download URL is tested against the regex, and if it matches, the token is provided in as a Bearer token.\n scan_models_on_startup: Scan the models directory on startup, registering orphaned models. This is typically only used in conjunction with `use_memory_db` for testing purposes.\n allow_private_download_urls: Allow the download queue to fetch from loopback, link-local and private-network addresses. Disabled by default so that a download URL cannot be used to reach services that are only reachable from the server. Enable this only if you install models from a mirror on your own network.\n download_proxy: Optional HTTP proxy for model downloads. The proxy must enforce the public-address policy because proxy-side DNS cannot be checked by InvokeAI.\n unsafe_disable_picklescan: UNSAFE. Disable the picklescan security check during model installation. Recommended only for development and testing purposes. This will allow arbitrary code execution during model installation, so should never be used in production.\n allow_unknown_models: Allow installation of models that we are unable to identify. If enabled, models will be marked as `unknown` in the database, and will not have any metadata associated with them. If disabled, unknown models will be rejected during installation.\n multiuser: Enable multiuser support. When disabled, the application runs in single-user mode using a default system account with administrator privileges. When enabled, requires user authentication and authorization.\n strict_password_checking: Enforce strict password requirements. When True, passwords must contain uppercase, lowercase, and numbers. When False (default), any password is accepted but its strength (weak/moderate/strong) is reported to the user.\n external_alibabacloud_api_key: API key for Alibaba Cloud DashScope image generation.\n external_alibabacloud_base_url: Base URL override for Alibaba Cloud DashScope image generation.\n external_gemini_api_key: API key for Gemini image generation.\n external_openai_api_key: API key for OpenAI image generation.\n external_gemini_base_url: Base URL override for Gemini image generation.\n external_openai_base_url: Base URL override for OpenAI image generation.\n external_seedream_api_key: API key for Seedream image generation.\n external_seedream_base_url: Base URL override for Seedream image generation.\n base_url: Public base path when running behind a reverse proxy under a sub-path, e.g. `/invoke`. Set only when the proxy PRESERVES the sub-path (the backend receives `/invoke/api/...`). Leave unset when the proxy strips the sub-path or when serving at the domain root.\n forwarded_allow_ips: Comma-separated list of IPs (or `*`) allowed to set X-Forwarded-* headers. Set to the reverse proxy's IP. Only used when `base_url` is set.\n http_compression_level: Compression level for gzipped HTTP API responses. 0 disables response compression entirely, 1 is fastest, 9 (the default) is smallest. Compression runs on the event loop and blocks the whole server while it works, and level 9 costs about 5.5x the time of level 1 for 0.4 percentage points of extra compression, so lowering this makes the app noticeably more responsive on large libraries. Set to 0 when a reverse proxy already compresses responses." }, "InvokeAIAppConfigWithSetFields": { "properties": { diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 2365dd00afa..d6b729b6948 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -19048,6 +19048,8 @@ export type components = { * device_working_mem_gb: The amount of working memory to keep available on the compute device (in GB). Has no effect if running on CPU. If you are experiencing OOM errors, try increasing this value. * enable_partial_loading: Enable partial loading of models. This enables models to run with reduced VRAM requirements (at the cost of slower speed) by streaming the model from RAM to VRAM as its used. In some edge cases, partial loading can cause models to run more slowly if they were previously being fully loaded into VRAM. * keep_ram_copy_of_weights: Whether to keep a full RAM copy of a model's weights when the model is loaded in VRAM. Keeping a RAM copy increases average RAM usage, but speeds up model switching and LoRA patching (assuming there is sufficient RAM). Set this to False if RAM pressure is consistently high. + * fp8_compute: Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM). + * fp8_compute_full_precision_hints: Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled. * ram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable. * vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable. * lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable. @@ -19355,6 +19357,18 @@ export type components = { * @default true */ keep_ram_copy_of_weights?: boolean; + /** + * Fp8 Compute + * @description Keep ComfyUI 'scaled fp8' checkpoints quantized instead of dequantizing them at load, and run their matmuls on the fp8 tensor cores (requires an Ada/SM 8.9 or newer NVIDIA GPU; falls back automatically otherwise). Roughly halves the transformer's VRAM and speeds up denoising, but quantizes activations as well, so images will differ from previous versions at the same seed. Reproducibility also requires the model to be FULLY resident in VRAM: a layer whose weights are still in RAM falls back to the dequantized path, and since which layers are resident shifts from run to run, the same seed then yields visibly different images. For repeatable output, ensure the model loads at 100% (e.g. enable_partial_loading=false with enough free VRAM). + * @default false + */ + fp8_compute?: boolean; + /** + * Fp8 Compute Full Precision Hints + * @description Honor the per-layer 'full_precision_matrix_mult' flags that some scaled-fp8 checkpoints ship. Those layers then dequantize on every forward instead of using the fp8 tensor cores, which can cost a large part of the fp8_compute speedup - on checkpoints that mark many layers, most of it. Set to false to run every quantized layer on the fp8 tensor cores, ignoring the producer's instruction; faster, but the marked layers were flagged as numerically sensitive, so quality may suffer. Only has an effect when fp8_compute is enabled. + * @default true + */ + fp8_compute_full_precision_hints?: boolean; /** * Ram * @description DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_ram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable. diff --git a/tests/backend/model_manager/load/state_dicts/anima_transformer_scaled_fp8_keys.py b/tests/backend/model_manager/load/state_dicts/anima_transformer_scaled_fp8_keys.py new file mode 100644 index 00000000000..115e8287b0e --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/anima_transformer_scaled_fp8_keys.py @@ -0,0 +1,92 @@ +"""Representative key layout of a scaled-fp8 Anima checkpoint. + +Captured from `pachiiahri/anima-fp8-comfyui` (`anima-preview_tcfp8_mixed`). It is the only +checkpoint in this series that ships **both** hint transports at once: a safetensors-header +`_quantization_metadata` block *and* per-layer `.comfy_quant` marker tensors. It also carries +exactly one `full_precision_matrix_mult` layer -- every other captured checkpoint marks either +none or a large fraction, so the single-marked case was previously untested. + +The header names its layers `net.`-prefixed, i.e. in the checkpoint's own scheme, while the scales +are read after `_strip_anima_bundle_prefix` has run. Reading the header without renaming matches +nothing and drops every flag silently. + +Subsetting rule: block 24 (the one carrying the marked layer) plus every non-block key. Values are +`(shape, dtype)`; `layer_hints` is the header block, subset the same way. +""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "net.blocks.24.adaln_modulation_cross_attn.1.weight": ([256, 2048], "BF16"), + "net.blocks.24.adaln_modulation_cross_attn.2.weight": ([6144, 256], "BF16"), + "net.blocks.24.adaln_modulation_mlp.1.weight": ([256, 2048], "BF16"), + "net.blocks.24.adaln_modulation_mlp.2.weight": ([6144, 256], "BF16"), + "net.blocks.24.adaln_modulation_self_attn.1.weight": ([256, 2048], "BF16"), + "net.blocks.24.adaln_modulation_self_attn.2.weight": ([6144, 256], "BF16"), + "net.blocks.24.cross_attn.k_norm.weight": ([128], "BF16"), + "net.blocks.24.cross_attn.k_proj.comfy_quant": ([64], "U8"), + "net.blocks.24.cross_attn.k_proj.input_scale": ([], "F32"), + "net.blocks.24.cross_attn.k_proj.weight": ([2048, 1024], "F8_E4M3"), + "net.blocks.24.cross_attn.k_proj.weight_scale": ([], "F32"), + "net.blocks.24.cross_attn.output_proj.comfy_quant": ([64], "U8"), + "net.blocks.24.cross_attn.output_proj.input_scale": ([], "F32"), + "net.blocks.24.cross_attn.output_proj.weight": ([2048, 2048], "F8_E4M3"), + "net.blocks.24.cross_attn.output_proj.weight_scale": ([], "F32"), + "net.blocks.24.cross_attn.q_norm.weight": ([128], "BF16"), + "net.blocks.24.cross_attn.q_proj.comfy_quant": ([63], "U8"), + "net.blocks.24.cross_attn.q_proj.weight": ([2048, 2048], "F8_E4M3"), + "net.blocks.24.cross_attn.q_proj.weight_scale": ([], "F32"), + "net.blocks.24.cross_attn.v_proj.comfy_quant": ([64], "U8"), + "net.blocks.24.cross_attn.v_proj.input_scale": ([], "F32"), + "net.blocks.24.cross_attn.v_proj.weight": ([2048, 1024], "F8_E4M3"), + "net.blocks.24.cross_attn.v_proj.weight_scale": ([], "F32"), + "net.blocks.24.mlp.layer1.comfy_quant": ([64], "U8"), + "net.blocks.24.mlp.layer1.input_scale": ([], "F32"), + "net.blocks.24.mlp.layer1.weight": ([8192, 2048], "F8_E4M3"), + "net.blocks.24.mlp.layer1.weight_scale": ([], "F32"), + "net.blocks.24.mlp.layer2.comfy_quant": ([64], "U8"), + "net.blocks.24.mlp.layer2.input_scale": ([], "F32"), + "net.blocks.24.mlp.layer2.weight": ([2048, 8192], "F8_E4M3"), + "net.blocks.24.mlp.layer2.weight_scale": ([], "F32"), + "net.blocks.24.self_attn.k_norm.weight": ([128], "BF16"), + "net.blocks.24.self_attn.k_proj.comfy_quant": ([64], "U8"), + "net.blocks.24.self_attn.k_proj.input_scale": ([], "F32"), + "net.blocks.24.self_attn.k_proj.weight": ([2048, 2048], "F8_E4M3"), + "net.blocks.24.self_attn.k_proj.weight_scale": ([], "F32"), + "net.blocks.24.self_attn.output_proj.comfy_quant": ([64], "U8"), + "net.blocks.24.self_attn.output_proj.input_scale": ([], "F32"), + "net.blocks.24.self_attn.output_proj.weight": ([2048, 2048], "F8_E4M3"), + "net.blocks.24.self_attn.output_proj.weight_scale": ([], "F32"), + "net.blocks.24.self_attn.q_norm.weight": ([128], "BF16"), + "net.blocks.24.self_attn.q_proj.comfy_quant": ([64], "U8"), + "net.blocks.24.self_attn.q_proj.input_scale": ([], "F32"), + "net.blocks.24.self_attn.q_proj.weight": ([2048, 2048], "F8_E4M3"), + "net.blocks.24.self_attn.q_proj.weight_scale": ([], "F32"), + "net.blocks.24.self_attn.v_proj.comfy_quant": ([64], "U8"), + "net.blocks.24.self_attn.v_proj.input_scale": ([], "F32"), + "net.blocks.24.self_attn.v_proj.weight": ([2048, 2048], "F8_E4M3"), + "net.blocks.24.self_attn.v_proj.weight_scale": ([], "F32"), + "net.final_layer.adaln_modulation.1.weight": ([256, 2048], "BF16"), + "net.final_layer.adaln_modulation.2.weight": ([4096, 256], "BF16"), + "net.final_layer.linear.weight": ([64, 2048], "BF16"), + "net.llm_adapter.embed.weight": ([32128, 1024], "BF16"), + "net.llm_adapter.norm.weight": ([1024], "BF16"), + "net.llm_adapter.out_proj.bias": ([1024], "BF16"), + "net.llm_adapter.out_proj.weight": ([1024, 1024], "BF16"), + "net.t_embedder.1.linear_1.weight": ([2048, 2048], "BF16"), + "net.t_embedder.1.linear_2.weight": ([6144, 2048], "BF16"), + "net.t_embedding_norm.weight": ([2048], "BF16"), + "net.x_embedder.proj.1.weight": ([2048, 68], "BF16"), +} + +# The `_quantization_metadata` header block, layer names exactly as the producer wrote them. +layer_hints: dict[str, dict[str, object]] = { + "net.blocks.24.cross_attn.k_proj": {"format": "float8_e4m3fn", "full_precision_matrix_mult": False}, + "net.blocks.24.cross_attn.output_proj": {"format": "float8_e4m3fn", "full_precision_matrix_mult": False}, + "net.blocks.24.cross_attn.q_proj": {"format": "float8_e4m3fn", "full_precision_matrix_mult": True}, + "net.blocks.24.cross_attn.v_proj": {"format": "float8_e4m3fn", "full_precision_matrix_mult": False}, + "net.blocks.24.mlp.layer1": {"format": "float8_e4m3fn", "full_precision_matrix_mult": False}, + "net.blocks.24.mlp.layer2": {"format": "float8_e4m3fn", "full_precision_matrix_mult": False}, + "net.blocks.24.self_attn.k_proj": {"format": "float8_e4m3fn", "full_precision_matrix_mult": False}, + "net.blocks.24.self_attn.output_proj": {"format": "float8_e4m3fn", "full_precision_matrix_mult": False}, + "net.blocks.24.self_attn.q_proj": {"format": "float8_e4m3fn", "full_precision_matrix_mult": False}, + "net.blocks.24.self_attn.v_proj": {"format": "float8_e4m3fn", "full_precision_matrix_mult": False}, +} diff --git a/tests/backend/model_manager/load/state_dicts/flux1_transformer_scaled_fp8_keys.py b/tests/backend/model_manager/load/state_dicts/flux1_transformer_scaled_fp8_keys.py new file mode 100644 index 00000000000..99cef0f12df --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/flux1_transformer_scaled_fp8_keys.py @@ -0,0 +1,122 @@ +"""Representative key layout of a *scaled* fp8 FLUX.1 dev transformer checkpoint. + +Captured from `comfyanonymous/flux_dev_scaled_fp8_test`, the reference ComfyUI scaled-fp8 export +for FLUX.1. Every key in that file carries the `model.diffusion_model.` bundle prefix; this fixture +records them as they look *after* `convert_bundle_to_flux_transformer_checkpoint` has stripped it, +which is where the loader reads them. + +Worth capturing for three reasons. It spells the scales `.scale_weight`/`.scale_input` rather than +`.weight_scale`/`.input_scale`, so it exercises the second spelling end to end. Every one of its +314 quantized Linears carries a *calibrated* `scale_input` -- not the uncalibrated 1.0 placeholder +that `_usable_input_scale` rejects -- which no other captured fixture provides. And its `qkv` stays +one fused `[9216, 3072]` Linear with a single scalar scale, matching the BFL layout InvokeAI's own +`Flux` implements, so no scale has to be split or copied. + +Before scaled-fp8 support reached this loader the file did not load at all: 629 scale and marker +keys reached `load_state_dict` as unexpected keys. + +Subsetting rule as for the sibling fixtures: block 0 of each stack plus every non-block key. +Values are `(shape, dtype)`. +""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "double_blocks.0.img_attn.norm.key_norm.scale": ([128], "BF16"), + "double_blocks.0.img_attn.norm.query_norm.scale": ([128], "BF16"), + "double_blocks.0.img_attn.proj.bias": ([3072], "F8_E4M3"), + "double_blocks.0.img_attn.proj.scale_input": ([], "F32"), + "double_blocks.0.img_attn.proj.scale_weight": ([], "F32"), + "double_blocks.0.img_attn.proj.weight": ([3072, 3072], "F8_E4M3"), + "double_blocks.0.img_attn.qkv.bias": ([9216], "F8_E4M3"), + "double_blocks.0.img_attn.qkv.scale_input": ([], "F32"), + "double_blocks.0.img_attn.qkv.scale_weight": ([], "F32"), + "double_blocks.0.img_attn.qkv.weight": ([9216, 3072], "F8_E4M3"), + "double_blocks.0.img_mlp.0.bias": ([12288], "F8_E4M3"), + "double_blocks.0.img_mlp.0.scale_input": ([], "F32"), + "double_blocks.0.img_mlp.0.scale_weight": ([], "F32"), + "double_blocks.0.img_mlp.0.weight": ([12288, 3072], "F8_E4M3"), + "double_blocks.0.img_mlp.2.bias": ([3072], "F8_E4M3"), + "double_blocks.0.img_mlp.2.scale_input": ([], "F32"), + "double_blocks.0.img_mlp.2.scale_weight": ([], "F32"), + "double_blocks.0.img_mlp.2.weight": ([3072, 12288], "F8_E4M3"), + "double_blocks.0.img_mod.lin.bias": ([18432], "F8_E4M3"), + "double_blocks.0.img_mod.lin.scale_input": ([], "F32"), + "double_blocks.0.img_mod.lin.scale_weight": ([], "F32"), + "double_blocks.0.img_mod.lin.weight": ([18432, 3072], "F8_E4M3"), + "double_blocks.0.txt_attn.norm.key_norm.scale": ([128], "BF16"), + "double_blocks.0.txt_attn.norm.query_norm.scale": ([128], "BF16"), + "double_blocks.0.txt_attn.proj.bias": ([3072], "F8_E4M3"), + "double_blocks.0.txt_attn.proj.scale_input": ([], "F32"), + "double_blocks.0.txt_attn.proj.scale_weight": ([], "F32"), + "double_blocks.0.txt_attn.proj.weight": ([3072, 3072], "F8_E4M3"), + "double_blocks.0.txt_attn.qkv.bias": ([9216], "F8_E4M3"), + "double_blocks.0.txt_attn.qkv.scale_input": ([], "F32"), + "double_blocks.0.txt_attn.qkv.scale_weight": ([], "F32"), + "double_blocks.0.txt_attn.qkv.weight": ([9216, 3072], "F8_E4M3"), + "double_blocks.0.txt_mlp.0.bias": ([12288], "F8_E4M3"), + "double_blocks.0.txt_mlp.0.scale_input": ([], "F32"), + "double_blocks.0.txt_mlp.0.scale_weight": ([], "F32"), + "double_blocks.0.txt_mlp.0.weight": ([12288, 3072], "F8_E4M3"), + "double_blocks.0.txt_mlp.2.bias": ([3072], "F8_E4M3"), + "double_blocks.0.txt_mlp.2.scale_input": ([], "F32"), + "double_blocks.0.txt_mlp.2.scale_weight": ([], "F32"), + "double_blocks.0.txt_mlp.2.weight": ([3072, 12288], "F8_E4M3"), + "double_blocks.0.txt_mod.lin.bias": ([18432], "F8_E4M3"), + "double_blocks.0.txt_mod.lin.scale_input": ([], "F32"), + "double_blocks.0.txt_mod.lin.scale_weight": ([], "F32"), + "double_blocks.0.txt_mod.lin.weight": ([18432, 3072], "F8_E4M3"), + "final_layer.adaLN_modulation.1.bias": ([6144], "F8_E4M3"), + "final_layer.adaLN_modulation.1.scale_input": ([], "F32"), + "final_layer.adaLN_modulation.1.scale_weight": ([], "F32"), + "final_layer.adaLN_modulation.1.weight": ([6144, 3072], "F8_E4M3"), + "final_layer.linear.bias": ([64], "F8_E4M3"), + "final_layer.linear.scale_input": ([], "F32"), + "final_layer.linear.scale_weight": ([], "F32"), + "final_layer.linear.weight": ([64, 3072], "F8_E4M3"), + "guidance_in.in_layer.bias": ([3072], "F8_E4M3"), + "guidance_in.in_layer.scale_input": ([], "F32"), + "guidance_in.in_layer.scale_weight": ([], "F32"), + "guidance_in.in_layer.weight": ([3072, 256], "F8_E4M3"), + "guidance_in.out_layer.bias": ([3072], "F8_E4M3"), + "guidance_in.out_layer.scale_input": ([], "F32"), + "guidance_in.out_layer.scale_weight": ([], "F32"), + "guidance_in.out_layer.weight": ([3072, 3072], "F8_E4M3"), + "img_in.bias": ([3072], "F8_E4M3"), + "img_in.scale_input": ([], "F32"), + "img_in.scale_weight": ([], "F32"), + "img_in.weight": ([3072, 64], "F8_E4M3"), + "scaled_fp8": ([0], "F8_E4M3"), + "single_blocks.0.linear1.bias": ([21504], "F8_E4M3"), + "single_blocks.0.linear1.scale_input": ([], "F32"), + "single_blocks.0.linear1.scale_weight": ([], "F32"), + "single_blocks.0.linear1.weight": ([21504, 3072], "F8_E4M3"), + "single_blocks.0.linear2.bias": ([3072], "F8_E4M3"), + "single_blocks.0.linear2.scale_input": ([], "F32"), + "single_blocks.0.linear2.scale_weight": ([], "F32"), + "single_blocks.0.linear2.weight": ([3072, 15360], "F8_E4M3"), + "single_blocks.0.modulation.lin.bias": ([9216], "F8_E4M3"), + "single_blocks.0.modulation.lin.scale_input": ([], "F32"), + "single_blocks.0.modulation.lin.scale_weight": ([], "F32"), + "single_blocks.0.modulation.lin.weight": ([9216, 3072], "F8_E4M3"), + "single_blocks.0.norm.key_norm.scale": ([128], "BF16"), + "single_blocks.0.norm.query_norm.scale": ([128], "BF16"), + "time_in.in_layer.bias": ([3072], "F8_E4M3"), + "time_in.in_layer.scale_input": ([], "F32"), + "time_in.in_layer.scale_weight": ([], "F32"), + "time_in.in_layer.weight": ([3072, 256], "F8_E4M3"), + "time_in.out_layer.bias": ([3072], "F8_E4M3"), + "time_in.out_layer.scale_input": ([], "F32"), + "time_in.out_layer.scale_weight": ([], "F32"), + "time_in.out_layer.weight": ([3072, 3072], "F8_E4M3"), + "txt_in.bias": ([3072], "F8_E4M3"), + "txt_in.scale_input": ([], "F32"), + "txt_in.scale_weight": ([], "F32"), + "txt_in.weight": ([3072, 4096], "F8_E4M3"), + "vector_in.in_layer.bias": ([3072], "F8_E4M3"), + "vector_in.in_layer.scale_input": ([], "F32"), + "vector_in.in_layer.scale_weight": ([], "F32"), + "vector_in.in_layer.weight": ([3072, 768], "F8_E4M3"), + "vector_in.out_layer.bias": ([3072], "F8_E4M3"), + "vector_in.out_layer.scale_input": ([], "F32"), + "vector_in.out_layer.scale_weight": ([], "F32"), + "vector_in.out_layer.weight": ([3072, 3072], "F8_E4M3"), +} diff --git a/tests/backend/model_manager/load/state_dicts/flux2_klein_4b_scaled_fp8_keys.py b/tests/backend/model_manager/load/state_dicts/flux2_klein_4b_scaled_fp8_keys.py new file mode 100644 index 00000000000..fd852f32a2e --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/flux2_klein_4b_scaled_fp8_keys.py @@ -0,0 +1,66 @@ +"""Representative key layout of the official scaled-fp8 FLUX.2 Klein 4B checkpoint. + +Captured from `black-forest-labs/FLUX.2-klein-4b-fp8`. Kept alongside +`flux2_transformer_fp8mixed_keys` because that one comes from FLUX.2 **dev**, whose quantizer left +the fused `qkv` alone -- so it cannot exercise the one case that makes FLUX.2 harder than FLUX.1. + +Here `double_blocks.0.img_attn.qkv.weight` is a fused `[9216, 3072]` fp8 tensor carrying a single +**scalar** `weight_scale`. diffusers wants three separate projections, so the weight is chunked and +the scalar has to be *copied* to all three -- splitting it would be wrong, and leaving it on the +fused path is worse: `attach_fp8_scales` then matches nothing and the three weights stay quantized +but unscaled, off by 1/weight_scale with nothing logged. + +The layer flags arrive through the safetensors header (`_quantization_metadata`), which names +layers in the BFL scheme -- so they need the same one-to-many rename as the scales. + +Subsetting rule as for the sibling fixtures: block 0 of each stack plus every non-block key. +Values are `(shape, dtype)`. +""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "double_blocks.0.img_attn.norm.key_norm.scale": ([128], "BF16"), + "double_blocks.0.img_attn.norm.query_norm.scale": ([128], "BF16"), + "double_blocks.0.img_attn.proj.input_scale": ([], "F32"), + "double_blocks.0.img_attn.proj.weight": ([3072, 3072], "F8_E4M3"), + "double_blocks.0.img_attn.proj.weight_scale": ([], "F32"), + "double_blocks.0.img_attn.qkv.input_scale": ([], "F32"), + "double_blocks.0.img_attn.qkv.weight": ([9216, 3072], "F8_E4M3"), + "double_blocks.0.img_attn.qkv.weight_scale": ([], "F32"), + "double_blocks.0.img_mlp.0.input_scale": ([], "F32"), + "double_blocks.0.img_mlp.0.weight": ([18432, 3072], "F8_E4M3"), + "double_blocks.0.img_mlp.0.weight_scale": ([], "F32"), + "double_blocks.0.img_mlp.2.input_scale": ([], "F32"), + "double_blocks.0.img_mlp.2.weight": ([3072, 9216], "F8_E4M3"), + "double_blocks.0.img_mlp.2.weight_scale": ([], "F32"), + "double_blocks.0.txt_attn.norm.key_norm.scale": ([128], "BF16"), + "double_blocks.0.txt_attn.norm.query_norm.scale": ([128], "BF16"), + "double_blocks.0.txt_attn.proj.input_scale": ([], "F32"), + "double_blocks.0.txt_attn.proj.weight": ([3072, 3072], "F8_E4M3"), + "double_blocks.0.txt_attn.proj.weight_scale": ([], "F32"), + "double_blocks.0.txt_attn.qkv.input_scale": ([], "F32"), + "double_blocks.0.txt_attn.qkv.weight": ([9216, 3072], "F8_E4M3"), + "double_blocks.0.txt_attn.qkv.weight_scale": ([], "F32"), + "double_blocks.0.txt_mlp.0.input_scale": ([], "F32"), + "double_blocks.0.txt_mlp.0.weight": ([18432, 3072], "F8_E4M3"), + "double_blocks.0.txt_mlp.0.weight_scale": ([], "F32"), + "double_blocks.0.txt_mlp.2.input_scale": ([], "F32"), + "double_blocks.0.txt_mlp.2.weight": ([3072, 9216], "F8_E4M3"), + "double_blocks.0.txt_mlp.2.weight_scale": ([], "F32"), + "double_stream_modulation_img.lin.weight": ([18432, 3072], "BF16"), + "double_stream_modulation_txt.lin.weight": ([18432, 3072], "BF16"), + "final_layer.adaLN_modulation.1.weight": ([6144, 3072], "BF16"), + "final_layer.linear.weight": ([128, 3072], "BF16"), + "img_in.weight": ([3072, 128], "BF16"), + "single_blocks.0.linear1.input_scale": ([], "F32"), + "single_blocks.0.linear1.weight": ([27648, 3072], "F8_E4M3"), + "single_blocks.0.linear1.weight_scale": ([], "F32"), + "single_blocks.0.linear2.input_scale": ([], "F32"), + "single_blocks.0.linear2.weight": ([3072, 12288], "F8_E4M3"), + "single_blocks.0.linear2.weight_scale": ([], "F32"), + "single_blocks.0.norm.key_norm.scale": ([128], "BF16"), + "single_blocks.0.norm.query_norm.scale": ([128], "BF16"), + "single_stream_modulation.lin.weight": ([9216, 3072], "BF16"), + "time_in.in_layer.weight": ([3072, 256], "BF16"), + "time_in.out_layer.weight": ([3072, 3072], "BF16"), + "txt_in.weight": ([3072, 7680], "BF16"), +} diff --git a/tests/backend/model_manager/load/state_dicts/flux2_transformer_fp8mixed_keys.py b/tests/backend/model_manager/load/state_dicts/flux2_transformer_fp8mixed_keys.py new file mode 100644 index 00000000000..a5082d4bc95 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/flux2_transformer_fp8mixed_keys.py @@ -0,0 +1,56 @@ +"""Representative key layout of a *mixed* fp8 FLUX.2 checkpoint. + +Captured from Comfy-Org's `flux2_dev_fp8mixed.safetensors`. Two properties make it worth having +next to the all-fp8 fixtures: + +- It is *mixed*: only some Linears are quantized, and the rest stay bf16 in the same file. A + metadata filter that is too broad silently deletes those bf16 weights instead of only the scale + bookkeeping, which no all-fp8 fixture can catch. +- Every quantized Linear carries a calibrated `.input_scale` next to its `.weight_scale`. That key + has to be stripped as well, or `load_state_dict(..., strict=True)` rejects it. + +Same subsetting rule as the sibling fixtures: block 0 of each stack plus every non-block key. +Values are `(shape, dtype)`. +""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "double_blocks.0.img_attn.norm.key_norm.scale": ([128], "BF16"), + "double_blocks.0.img_attn.norm.query_norm.scale": ([128], "BF16"), + "double_blocks.0.img_attn.proj.weight": ([6144, 6144], "BF16"), + "double_blocks.0.img_attn.qkv.weight": ([18432, 6144], "BF16"), + "double_blocks.0.img_mlp.0.input_scale": ([], "F32"), + "double_blocks.0.img_mlp.0.weight": ([36864, 6144], "F8_E4M3"), + "double_blocks.0.img_mlp.0.weight_scale": ([], "F32"), + "double_blocks.0.img_mlp.2.input_scale": ([], "F32"), + "double_blocks.0.img_mlp.2.weight": ([6144, 18432], "F8_E4M3"), + "double_blocks.0.img_mlp.2.weight_scale": ([], "F32"), + "double_blocks.0.txt_attn.norm.key_norm.scale": ([128], "BF16"), + "double_blocks.0.txt_attn.norm.query_norm.scale": ([128], "BF16"), + "double_blocks.0.txt_attn.proj.weight": ([6144, 6144], "BF16"), + "double_blocks.0.txt_attn.qkv.weight": ([18432, 6144], "BF16"), + "double_blocks.0.txt_mlp.0.input_scale": ([], "F32"), + "double_blocks.0.txt_mlp.0.weight": ([36864, 6144], "F8_E4M3"), + "double_blocks.0.txt_mlp.0.weight_scale": ([], "F32"), + "double_blocks.0.txt_mlp.2.input_scale": ([], "F32"), + "double_blocks.0.txt_mlp.2.weight": ([6144, 18432], "F8_E4M3"), + "double_blocks.0.txt_mlp.2.weight_scale": ([], "F32"), + "double_stream_modulation_img.lin.weight": ([36864, 6144], "BF16"), + "double_stream_modulation_txt.lin.weight": ([36864, 6144], "BF16"), + "final_layer.adaLN_modulation.1.weight": ([12288, 6144], "BF16"), + "final_layer.linear.weight": ([128, 6144], "BF16"), + "guidance_in.in_layer.weight": ([6144, 256], "BF16"), + "guidance_in.out_layer.weight": ([6144, 6144], "BF16"), + "img_in.weight": ([6144, 128], "BF16"), + "single_blocks.0.linear1.input_scale": ([], "F32"), + "single_blocks.0.linear1.weight": ([55296, 6144], "F8_E4M3"), + "single_blocks.0.linear1.weight_scale": ([], "F32"), + "single_blocks.0.linear2.input_scale": ([], "F32"), + "single_blocks.0.linear2.weight": ([6144, 24576], "F8_E4M3"), + "single_blocks.0.linear2.weight_scale": ([], "F32"), + "single_blocks.0.norm.key_norm.scale": ([128], "BF16"), + "single_blocks.0.norm.query_norm.scale": ([128], "BF16"), + "single_stream_modulation.lin.weight": ([18432, 6144], "BF16"), + "time_in.in_layer.weight": ([6144, 256], "BF16"), + "time_in.out_layer.weight": ([6144, 6144], "BF16"), + "txt_in.weight": ([6144, 15360], "BF16"), +} diff --git a/tests/backend/model_manager/load/state_dicts/z_image_transformer_scaled_fp8_keys.py b/tests/backend/model_manager/load/state_dicts/z_image_transformer_scaled_fp8_keys.py new file mode 100644 index 00000000000..b971d339dc7 --- /dev/null +++ b/tests/backend/model_manager/load/state_dicts/z_image_transformer_scaled_fp8_keys.py @@ -0,0 +1,94 @@ +"""Representative key layout of a *scaled* fp8 Z-Image transformer checkpoint. + +Captured from `zImageTurboFP8Kijai_fp8ScaledE4m3fn.safetensors`. Unlike +`z_image_transformer_comfyui_keys`, which is a bf16 checkpoint, this one is ComfyUI "scaled fp8": +each quantized Linear carries an fp8 `.weight` plus a scalar `.scale_weight`, and the +file is marked with a stray `scaled_fp8` key. + +Two things make it worth capturing. It spells the scale `.scale_weight` rather than +`.weight_scale` -- the spelling several loaders used to ignore -- and its scales are all *above* +one (1.5 to 7.6 in the real file), so dropping them leaves each weight at a different fraction of +its true magnitude rather than uniformly wrong. + +Same subsetting rule as the sibling fixture: block 0 of each stack plus every non-block key. +Values are `(shape, dtype)`. +""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "cap_embedder.0.weight": ([2560], "F32"), + "cap_embedder.1.bias": ([3840], "F32"), + "cap_embedder.1.scale_weight": ([1], "F32"), + "cap_embedder.1.weight": ([3840, 2560], "F8_E4M3"), + "cap_pad_token": ([1, 3840], "F32"), + "context_refiner.0.attention.k_norm.weight": ([128], "F32"), + "context_refiner.0.attention.out.scale_weight": ([1], "F32"), + "context_refiner.0.attention.out.weight": ([3840, 3840], "F8_E4M3"), + "context_refiner.0.attention.q_norm.weight": ([128], "F32"), + "context_refiner.0.attention.qkv.scale_weight": ([1], "F32"), + "context_refiner.0.attention.qkv.weight": ([11520, 3840], "F8_E4M3"), + "context_refiner.0.attention_norm1.weight": ([3840], "F32"), + "context_refiner.0.attention_norm2.weight": ([3840], "F32"), + "context_refiner.0.feed_forward.w1.scale_weight": ([1], "F32"), + "context_refiner.0.feed_forward.w1.weight": ([10240, 3840], "F8_E4M3"), + "context_refiner.0.feed_forward.w2.scale_weight": ([1], "F32"), + "context_refiner.0.feed_forward.w2.weight": ([3840, 10240], "F8_E4M3"), + "context_refiner.0.feed_forward.w3.scale_weight": ([1], "F32"), + "context_refiner.0.feed_forward.w3.weight": ([10240, 3840], "F8_E4M3"), + "context_refiner.0.ffn_norm1.weight": ([3840], "F32"), + "context_refiner.0.ffn_norm2.weight": ([3840], "F32"), + "final_layer.adaLN_modulation.1.bias": ([3840], "F32"), + "final_layer.adaLN_modulation.1.scale_weight": ([1], "F32"), + "final_layer.adaLN_modulation.1.weight": ([3840, 256], "F8_E4M3"), + "final_layer.linear.bias": ([64], "F32"), + "final_layer.linear.scale_weight": ([1], "F32"), + "final_layer.linear.weight": ([64, 3840], "F8_E4M3"), + "layers.0.adaLN_modulation.0.bias": ([15360], "F32"), + "layers.0.adaLN_modulation.0.scale_weight": ([1], "F32"), + "layers.0.adaLN_modulation.0.weight": ([15360, 256], "F8_E4M3"), + "layers.0.attention.k_norm.weight": ([128], "F32"), + "layers.0.attention.out.scale_weight": ([1], "F32"), + "layers.0.attention.out.weight": ([3840, 3840], "F8_E4M3"), + "layers.0.attention.q_norm.weight": ([128], "F32"), + "layers.0.attention.qkv.scale_weight": ([1], "F32"), + "layers.0.attention.qkv.weight": ([11520, 3840], "F8_E4M3"), + "layers.0.attention_norm1.weight": ([3840], "F32"), + "layers.0.attention_norm2.weight": ([3840], "F32"), + "layers.0.feed_forward.w1.scale_weight": ([1], "F32"), + "layers.0.feed_forward.w1.weight": ([10240, 3840], "F8_E4M3"), + "layers.0.feed_forward.w2.scale_weight": ([1], "F32"), + "layers.0.feed_forward.w2.weight": ([3840, 10240], "F8_E4M3"), + "layers.0.feed_forward.w3.scale_weight": ([1], "F32"), + "layers.0.feed_forward.w3.weight": ([10240, 3840], "F8_E4M3"), + "layers.0.ffn_norm1.weight": ([3840], "F32"), + "layers.0.ffn_norm2.weight": ([3840], "F32"), + "noise_refiner.0.adaLN_modulation.0.bias": ([15360], "F32"), + "noise_refiner.0.adaLN_modulation.0.scale_weight": ([1], "F32"), + "noise_refiner.0.adaLN_modulation.0.weight": ([15360, 256], "F8_E4M3"), + "noise_refiner.0.attention.k_norm.weight": ([128], "F32"), + "noise_refiner.0.attention.out.scale_weight": ([1], "F32"), + "noise_refiner.0.attention.out.weight": ([3840, 3840], "F8_E4M3"), + "noise_refiner.0.attention.q_norm.weight": ([128], "F32"), + "noise_refiner.0.attention.qkv.scale_weight": ([1], "F32"), + "noise_refiner.0.attention.qkv.weight": ([11520, 3840], "F8_E4M3"), + "noise_refiner.0.attention_norm1.weight": ([3840], "F32"), + "noise_refiner.0.attention_norm2.weight": ([3840], "F32"), + "noise_refiner.0.feed_forward.w1.scale_weight": ([1], "F32"), + "noise_refiner.0.feed_forward.w1.weight": ([10240, 3840], "F8_E4M3"), + "noise_refiner.0.feed_forward.w2.scale_weight": ([1], "F32"), + "noise_refiner.0.feed_forward.w2.weight": ([3840, 10240], "F8_E4M3"), + "noise_refiner.0.feed_forward.w3.scale_weight": ([1], "F32"), + "noise_refiner.0.feed_forward.w3.weight": ([10240, 3840], "F8_E4M3"), + "noise_refiner.0.ffn_norm1.weight": ([3840], "F32"), + "noise_refiner.0.ffn_norm2.weight": ([3840], "F32"), + "scaled_fp8": ([2], "F8_E4M3"), + "t_embedder.mlp.0.bias": ([1024], "F32"), + "t_embedder.mlp.0.scale_weight": ([1], "F32"), + "t_embedder.mlp.0.weight": ([1024, 256], "F8_E4M3"), + "t_embedder.mlp.2.bias": ([256], "F32"), + "t_embedder.mlp.2.scale_weight": ([1], "F32"), + "t_embedder.mlp.2.weight": ([256, 1024], "F8_E4M3"), + "x_embedder.bias": ([3840], "F32"), + "x_embedder.scale_weight": ([1], "F32"), + "x_embedder.weight": ([3840, 64], "F8_E4M3"), + "x_pad_token": ([1, 3840], "F32"), +} diff --git a/tests/backend/model_manager/load/test_anima_scaled_fp8_keys.py b/tests/backend/model_manager/load/test_anima_scaled_fp8_keys.py new file mode 100644 index 00000000000..05f7c597e83 --- /dev/null +++ b/tests/backend/model_manager/load/test_anima_scaled_fp8_keys.py @@ -0,0 +1,131 @@ +"""A scaled-fp8 Anima checkpoint must load, and its layer flags must survive the prefix strip. + +Two separate failures are covered here. Before scaled-fp8 support reached `AnimaCheckpointModel`, +`_filter_non_model_keys` let the scale and marker keys through and the loader raised on them -- +500 unexpected keys on a plain scaled export, 749 on one that also ships `comfy_quant` markers, so +such a checkpoint did not load at all. + +The second is quieter: `_quantization_metadata` names its layers `net.`-prefixed, in the +checkpoint's own scheme, while the scales are read after the prefix has been stripped. Reading the +header without renaming matches nothing, and every `full_precision_matrix_mult` flag is dropped +without a word. +""" + +import json + +import torch + +from invokeai.backend.model_manager.load.model_loaders.anima import ( + _filter_non_model_keys, + _strip_anima_bundle_prefix, +) +from invokeai.backend.quantization.fp8_scaled import ( + FP8_DTYPE, + extract_comfy_quant_hints, + extract_fp8_scaled_layers, + is_scale_metadata_key, + strip_layer_path_prefix, +) +from tests.backend.model_manager.load.state_dicts.anima_transformer_scaled_fp8_keys import ( + layer_hints as header_hints, +) +from tests.backend.model_manager.load.state_dicts.anima_transformer_scaled_fp8_keys import ( + state_dict_keys as anima_keys, +) + +_DTYPES = {"F8_E4M3": FP8_DTYPE, "F32": torch.float32, "BF16": torch.bfloat16, "U8": torch.uint8} + + +def _build_state_dict() -> dict[str, torch.Tensor]: + """Rebuild the checkpoint. `comfy_quant` markers are real JSON blobs, not placeholders. + + The marker is the only transport some checkpoints have, so a test that fakes it would not + exercise the path that reads it. + """ + sd: dict[str, torch.Tensor] = {} + for key, (shape, dtype) in anima_keys.items(): + torch_dtype = _DTYPES[dtype] + if key.endswith(".comfy_quant"): + path = key[: -len(".comfy_quant")] + blob = header_hints.get(path, {"format": "float8_e4m3fn"}) + # Real JSON: the producer writes `true`/`false`, and `extract_comfy_quant_hints` + # parses the blob. Python's `False` would not parse and the flag would vanish. + sd[key] = torch.frombuffer(bytearray(json.dumps(blob).encode()), dtype=torch.uint8).clone() + elif torch_dtype is FP8_DTYPE: + sd[key] = torch.zeros(shape, dtype=torch.float32).to(FP8_DTYPE) + elif key.endswith((".weight_scale", ".input_scale")): + # 1.0 is the placeholder `_usable_input_scale` rejects, so it must not be used here. + sd[key] = torch.full(shape, 2.5, dtype=torch_dtype) + else: + sd[key] = torch.zeros(shape, dtype=torch_dtype) + return _filter_non_model_keys(_strip_anima_bundle_prefix(sd)) + + +def test_nothing_the_loader_would_reject_is_left_behind() -> None: + """`load_state_dict` must see only model tensors; the loader raises on anything else.""" + sd = _build_state_dict() + assert [k for k in sd if is_scale_metadata_key(k)], "fixture carries no side-channel keys" + + extract_fp8_scaled_layers(sd) + + assert [k for k in sd if is_scale_metadata_key(k)] == [] + + +def test_every_quantized_linear_is_recognized() -> None: + sd = _build_state_dict() + fp8_weights = {k for k, v in sd.items() if v.dtype is FP8_DTYPE and k.endswith(".weight")} + assert fp8_weights + + layers = extract_fp8_scaled_layers(sd, layer_hints=extract_comfy_quant_hints(sd)) + + assert {f"{path}.weight" for path in layers} == fp8_weights + + +def test_header_hints_are_renamed_to_the_stripped_paths() -> None: + """The header names layers `net.`-prefixed; the scales are keyed on the stripped paths. + + Without the rename the flags match nothing. This is the mistake that already cost a debugging + round on Krea-2, where the metadata was read against the wrong naming. + """ + assert all(name.startswith("net.") for name in header_hints), "fixture header is not prefixed" + + remapped = strip_layer_path_prefix(header_hints) + + assert all(not renamed.startswith("net.") for renamed in remapped) + + layers = extract_fp8_scaled_layers(_build_state_dict(), layer_hints=remapped) + + assert layers + assert set(remapped) >= set(layers), "a renamed hint no longer lines up with its layer" + + +def test_the_single_full_precision_flag_survives_both_transports() -> None: + """Exactly one layer is marked. It must come through whichever transport is read. + + Every other captured checkpoint marks either no layers or a large fraction, so a bug that + dropped a lone flag would go unnoticed there. + """ + marked = [name for name, hints in header_hints.items() if hints.get("full_precision_matrix_mult")] + assert len(marked) == 1, f"fixture should carry exactly one marked layer, has {len(marked)}" + + for hints in ( + strip_layer_path_prefix(header_hints), # header transport + extract_comfy_quant_hints(_build_state_dict()), # per-layer marker transport + ): + layers = extract_fp8_scaled_layers(_build_state_dict(), layer_hints=hints) + assert sum(1 for layer in layers.values() if layer.full_precision_matmul) == 1 + + +def test_only_the_full_precision_layer_lacks_an_input_scale() -> None: + """Every quantized Linear ships a calibrated `scale_input` -- except the marked one. + + That is coherent on the producer's side: a layer excluded from the fp8 matmul has no activation + scale to calibrate. It also means "all layers have an input scale" is the wrong invariant. + """ + marked = {n for n, h in strip_layer_path_prefix(header_hints).items() if h.get("full_precision_matrix_mult")} + + layers = extract_fp8_scaled_layers(_build_state_dict()) + + assert layers + without = {path for path, layer in layers.items() if layer.input_scale is None} + assert without == marked diff --git a/tests/backend/model_manager/load/test_flux1_scaled_fp8_keys.py b/tests/backend/model_manager/load/test_flux1_scaled_fp8_keys.py new file mode 100644 index 00000000000..5dba27a6729 --- /dev/null +++ b/tests/backend/model_manager/load/test_flux1_scaled_fp8_keys.py @@ -0,0 +1,143 @@ +"""The scaled-fp8 FLUX.1 key layout must be recognized, scales and all. + +Before scaled-fp8 support reached `FluxCheckpointModel`, this checkpoint did not load at all: the +loader knew nothing about `.scale_weight`/`.scale_input`, so all 629 scale and marker keys reached +`load_state_dict` as unexpected keys and it raised. + +The fixture is a real key layout; the tensors are synthetic because only shapes and dtypes matter +to the code under test. +""" + +import torch + +from invokeai.backend.model_manager.util.model_util import convert_bundle_to_flux_transformer_checkpoint +from invokeai.backend.quantization.fp8_scaled import ( + FP8_DTYPE, + can_stay_quantized, + extract_comfy_quant_hints, + extract_fp8_scaled_layers, + is_scale_metadata_key, +) +from tests.backend.model_manager.load.state_dicts.flux1_transformer_scaled_fp8_keys import ( + state_dict_keys as scaled_keys, +) + +_DTYPES = {"F8_E4M3": FP8_DTYPE, "F32": torch.float32, "BF16": torch.bfloat16} + + +def _build_state_dict() -> dict[str, torch.Tensor]: + sd: dict[str, torch.Tensor] = {} + for key, (shape, dtype) in scaled_keys.items(): + torch_dtype = _DTYPES[dtype] + if torch_dtype is FP8_DTYPE: + sd[key] = torch.zeros(shape, dtype=torch.float32).to(FP8_DTYPE) + elif key.endswith((".scale_weight", ".scale_input")): + # Real checkpoints carry calibrated scales; 1.0 is the placeholder value that + # `_usable_input_scale` deliberately rejects, so it must not be used here. + sd[key] = torch.full(shape, 2.5, dtype=torch_dtype) + else: + sd[key] = torch.zeros(shape, dtype=torch_dtype) + return sd + + +def test_every_quantized_linear_is_recognized() -> None: + """Each fp8 `.weight` in the fixture must come back as a scaled layer, not be left behind.""" + sd = _build_state_dict() + fp8_weights = {k for k, v in sd.items() if v.dtype is FP8_DTYPE and k.endswith(".weight")} + assert fp8_weights, "fixture carries no fp8 weights" + + layers = extract_fp8_scaled_layers(sd, layer_hints=extract_comfy_quant_hints(sd)) + + assert {f"{path}.weight" for path in layers} == fp8_weights + + +def test_the_fp8_biases_must_not_stay_quantized() -> None: + """This checkpoint quantizes the biases too -- 314 of them, alongside the 314 weights. + + An fp8 bias saves nothing usable and breaks inference: the value reaches the activations and + the next Linear receives an fp8 *input*, which dies in `x.abs()` with + `"abs_cuda" not implemented for 'Float8_e4m3fn'`. Only 2-D `nn.Linear.weight` may stay + quantized, which is what `can_stay_quantized` decides -- a bias is 1-D and fails it. + """ + fp8_biases = [k for k, (_, dtype) in scaled_keys.items() if dtype == "F8_E4M3" and k.endswith(".bias")] + assert fp8_biases, "fixture is expected to carry fp8 biases -- the real checkpoint does" + + sd = _build_state_dict() + for key in fp8_biases: + assert not can_stay_quantized(key, sd[key], None) + + +def test_the_second_spelling_is_read() -> None: + """This checkpoint spells the scales `.scale_weight`/`.scale_input`. + + Reading only `.weight_scale`/`.input_scale` would drop every scale here, leaving each weight + off by `1/scale` with nothing logged. + """ + assert any(k.endswith(".scale_weight") for k in scaled_keys) + assert not any(k.endswith(".weight_scale") for k in scaled_keys) + + layers = extract_fp8_scaled_layers(_build_state_dict()) + + assert layers + assert all(layer.weight_scale is not None for layer in layers.values()) + + +def test_the_calibrated_input_scales_survive() -> None: + """Every quantized Linear here ships a calibrated `scale_input`; none may be discarded. + + `_usable_input_scale` drops uncalibrated placeholders (exactly 1.0, non-finite, <= 0). A + calibrated value must pass, otherwise every forward pays an amax reduction it does not need. + """ + layers = extract_fp8_scaled_layers(_build_state_dict()) + + assert all(layer.input_scale is not None for layer in layers.values()) + + +def test_extraction_strips_every_scale_and_marker_key() -> None: + """What is left must load into a model that knows nothing about fp8.""" + sd = _build_state_dict() + assert "scaled_fp8" in sd, "fixture is missing the producer's marker key" + + extract_fp8_scaled_layers(sd) + + assert not [k for k in sd if is_scale_metadata_key(k)] + assert "scaled_fp8" not in sd + + +def test_qkv_stays_fused_with_one_scalar_scale() -> None: + """FLUX.1 needs no qkv split, unlike FLUX.2. + + InvokeAI's `Flux` implements the BFL layout, where `qkv` is a single Linear. The checkpoint + matches it, so the per-tensor scale attaches to exactly the module it was computed for. If this + ever changes, the scale would have to be *copied* to each split part rather than moved. + """ + qkv = [k for k in scaled_keys if k.endswith(".qkv.weight")] + assert qkv, "fixture carries no fused qkv weight" + + for key in qkv: + shape, dtype = scaled_keys[key] + assert dtype == "F8_E4M3" + assert shape[0] == 3 * shape[1], f"{key} is not a fused qkv: {shape}" + assert scaled_keys[key.replace(".weight", ".scale_weight")][0] == [] + + +def test_bundle_conversion_keeps_the_scales_at_full_precision() -> None: + """The bundle converter casts RMSNorm `.scale` to bf16; fp8 quantization scales must not follow. + + `.weight_scale` also ends in "scale". Folding an f32 quantization scale to bf16 leaves 8 + mantissa bits on a value every quantized weight is multiplied by. + """ + sd = { + "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale": torch.ones(128, dtype=torch.float32), + "model.diffusion_model.double_blocks.0.img_attn.qkv.weight_scale": torch.full((), 2.5, dtype=torch.float32), + "model.diffusion_model.double_blocks.0.img_attn.qkv.scale_weight": torch.full((), 2.5, dtype=torch.float32), + "model.diffusion_model.double_blocks.0.img_attn.qkv.weight": torch.zeros(9216, 3072, dtype=torch.float32).to( + FP8_DTYPE + ), + } + + converted = convert_bundle_to_flux_transformer_checkpoint(sd) + + assert converted["double_blocks.0.img_attn.norm.key_norm.scale"].dtype is torch.bfloat16 + assert converted["double_blocks.0.img_attn.qkv.weight_scale"].dtype is torch.float32 + assert converted["double_blocks.0.img_attn.qkv.scale_weight"].dtype is torch.float32 diff --git a/tests/backend/model_manager/load/test_flux2_fp8mixed_keys.py b/tests/backend/model_manager/load/test_flux2_fp8mixed_keys.py new file mode 100644 index 00000000000..1ab3165eb64 --- /dev/null +++ b/tests/backend/model_manager/load/test_flux2_fp8mixed_keys.py @@ -0,0 +1,120 @@ +"""A mixed fp8 FLUX.2 checkpoint must lose its scale bookkeeping and nothing else. + +`flux2_dev_fp8mixed` quantizes only some Linears; the rest stay bf16 in the same file. It also +carries a calibrated `.input_scale` beside every `.weight_scale`, and — the reason this fixture +earns its place — real learned parameters whose names end in `.scale` +(`img_attn.norm.query_norm.scale`). A metadata filter that matches on "scale" anywhere, or on a +bare `.scale` suffix, deletes those weights. An all-fp8 fixture cannot catch either mistake. +""" + +import torch + +from invokeai.backend.model_manager.load.model_loaders.flux import Flux2CheckpointModel +from invokeai.backend.quantization.fp8_scaled import FP8_DTYPE, is_scale_metadata_key +from tests.backend.model_manager.load.state_dicts.flux2_transformer_fp8mixed_keys import ( + state_dict_keys as mixed_keys, +) + +_DTYPES = {"F8_E4M3": FP8_DTYPE, "F32": torch.float32, "BF16": torch.bfloat16} + +SCALE_VALUE = 0.25 + + +def _mock_state_dict() -> dict[str, torch.Tensor]: + sd: dict[str, torch.Tensor] = {} + for key, (shape, dtype) in mixed_keys.items(): + if key.endswith((".weight_scale", ".input_scale")): + sd[key] = torch.tensor(SCALE_VALUE, dtype=torch.float32) + else: + sd[key] = torch.ones(shape, dtype=torch.float32).to(_DTYPES[dtype]) + return sd + + +def _quantized_weights() -> set[str]: + return {k[: -len(".weight_scale")] + ".weight" for k in mixed_keys if k.endswith(".weight_scale")} + + +def test_the_fixture_is_really_mixed_and_carries_input_scales() -> None: + """Guard the fixture: an all-fp8 or scale-free recapture makes the tests below vacuous.""" + dtypes = {dtype for _, dtype in mixed_keys.values()} + assert "F8_E4M3" in dtypes and "BF16" in dtypes, "fixture is no longer mixed" + assert any(k.endswith(".weight_scale") for k in mixed_keys) + assert any(k.endswith(".input_scale") for k in mixed_keys) + assert any(k.endswith(".scale") for k in mixed_keys), "no learned `.scale` parameter left to protect" + + +def test_learned_scale_parameters_are_not_treated_as_metadata() -> None: + learned = [k for k in mixed_keys if k.endswith(".scale")] + + assert learned + for key in learned: + assert not is_scale_metadata_key(key), f"{key} is a weight, not quantization bookkeeping" + + +def test_dequantization_strips_every_scale_key() -> None: + sd = Flux2CheckpointModel._dequantize_fp8_weights(None, _mock_state_dict()) + + assert not [k for k in sd if k.endswith((".weight_scale", ".scale_weight", ".input_scale"))] + + +def test_unquantized_tensors_survive_untouched() -> None: + before = _mock_state_dict() + quantized = _quantized_weights() + + after = Flux2CheckpointModel._dequantize_fp8_weights(None, dict(before)) + + for key, value in before.items(): + if key.endswith((".weight_scale", ".input_scale")) or key in quantized: + continue + assert key in after, f"{key} was deleted" + assert torch.equal(after[key], value), f"{key} was modified" + + +def test_the_scale_is_folded_into_the_quantized_weights() -> None: + sd = Flux2CheckpointModel._dequantize_fp8_weights(None, _mock_state_dict()) + + quantized = _quantized_weights() + assert quantized + for key in quantized: + # Each mock fp8 weight holds 1.0, so the folded result is exactly the scale. + assert torch.allclose(sd[key].float(), torch.full_like(sd[key].float(), SCALE_VALUE)) + assert sd[key].dtype is torch.bfloat16 + + +def test_the_scale_weight_spelling_is_folded_on_this_layout_too() -> None: + """The regression guard: same real layout, scales spelled the other way. + + Reading only `.weight_scale` while stripping both spellings left each quantized weight at its + raw fp8 codes — measured at a relative error of ~4700 on the real checkpoint, with nothing + logged. + """ + renamed = { + (k[: -len(".weight_scale")] + ".scale_weight" if k.endswith(".weight_scale") else k): v + for k, v in _mock_state_dict().items() + } + + sd = Flux2CheckpointModel._dequantize_fp8_weights(None, renamed) + + quantized = _quantized_weights() + assert quantized + for key in quantized: + assert torch.allclose(sd[key].float(), torch.full_like(sd[key].float(), SCALE_VALUE)) + assert not [k for k in sd if k.endswith((".weight_scale", ".scale_weight"))] + + +def test_a_per_channel_input_scale_is_stripped_by_name_not_by_rank() -> None: + """A vector-valued `input_scale` must go too. + + Scalar scales were removed incidentally, by the "0-dimensional tensors are metadata" branch + rather than by name. Anything with a shape slipped past it and reached + `load_state_dict(..., strict=True)` as an unexpected key. + """ + sd = _mock_state_dict() + per_channel = [k for k in sd if k.endswith(".input_scale")] + assert per_channel + for key in per_channel: + sd[key] = torch.full((8,), SCALE_VALUE, dtype=torch.float32) + + out = Flux2CheckpointModel._dequantize_fp8_weights(None, sd) + + assert not [k for k in out if k.endswith(".input_scale")] diff --git a/tests/backend/model_manager/load/test_flux2_scaled_fp8_keys.py b/tests/backend/model_manager/load/test_flux2_scaled_fp8_keys.py new file mode 100644 index 00000000000..1acd3c72758 --- /dev/null +++ b/tests/backend/model_manager/load/test_flux2_scaled_fp8_keys.py @@ -0,0 +1,111 @@ +"""A scaled-fp8 FLUX.2 checkpoint must survive the BFL -> diffusers rename with its scales. + +FLUX.2 is the hard case of the scaled-fp8 loaders: unlike FLUX.1 the keys are renamed, and the +fused `qkv` is split into three projections. Both steps can lose a scale silently -- the weights +then stay quantized but unscaled, off by `1/weight_scale`, with nothing logged. +""" + +import torch + +from invokeai.backend.model_manager.load.model_loaders.flux2_state_dict_utils import ( + convert_flux2_bfl_to_diffusers, + remap_flux2_layer_paths, +) +from invokeai.backend.quantization.fp8_scaled import FP8_DTYPE, extract_fp8_scaled_layers +from tests.backend.model_manager.load.state_dicts.flux2_klein_4b_scaled_fp8_keys import ( + state_dict_keys as klein_keys, +) + +_DTYPES = {"F8_E4M3": FP8_DTYPE, "F32": torch.float32, "BF16": torch.bfloat16} + + +def _build_state_dict() -> dict[str, torch.Tensor]: + sd: dict[str, torch.Tensor] = {} + for key, (shape, dtype) in klein_keys.items(): + torch_dtype = _DTYPES[dtype] + if torch_dtype is FP8_DTYPE: + sd[key] = torch.zeros(shape, dtype=torch.float32).to(FP8_DTYPE) + elif key.endswith((".weight_scale", ".input_scale")): + # 1.0 is the placeholder `_usable_input_scale` rejects, so it must not be used here. + sd[key] = torch.full(shape, 2.5, dtype=torch_dtype) + else: + sd[key] = torch.zeros(shape, dtype=torch_dtype) + return sd + + +def test_no_scale_is_left_on_a_bfl_path() -> None: + """Every scale must land next to the weight it belongs to, under its diffusers name.""" + converted = convert_flux2_bfl_to_diffusers(_build_state_dict()) + + scales = [k for k in converted if k.endswith((".weight_scale", ".input_scale"))] + assert scales, "fixture carries no scales" + orphans = [k for k in scales if k.rsplit(".", 1)[0] + ".weight" not in converted] + assert orphans == [] + assert [k for k in converted if k.startswith(("double_blocks.", "single_blocks."))] == [] + + +def test_the_fused_qkv_scalar_is_copied_to_all_three_projections() -> None: + """A per-tensor scale describes the whole fused tensor, so each third inherits it unchanged.""" + fused = [k for k in klein_keys if k.endswith(".qkv.weight")] + assert fused, "fixture carries no fused qkv weight" + + converted = convert_flux2_bfl_to_diffusers(_build_state_dict()) + + for group in (("to_q", "to_k", "to_v"), ("add_q_proj", "add_k_proj", "add_v_proj")): + for name in group: + key = f"transformer_blocks.0.attn.{name}.weight_scale" + assert key in converted, f"missing {key}" + assert converted[key].shape == torch.Size([]), "a scalar scale must not be split" + assert torch.equal(converted[key], torch.tensor(2.5)) + + +def test_a_scale_never_overwrites_its_own_weight() -> None: + """The block renames are substring tests, so `...proj.weight_scale` matches `...proj.weight`. + + Routed through the weight converter, the scale would be written to the weight's destination key + and replace it. The weight must still be the fp8 tensor afterwards, not an f32 scalar. + """ + converted = convert_flux2_bfl_to_diffusers(_build_state_dict()) + + weight = converted["transformer_blocks.0.attn.to_out.0.weight"] + assert weight.dtype is FP8_DTYPE + assert weight.dim() == 2 + + +def test_every_quantized_linear_is_recognized_after_the_rename() -> None: + sd = convert_flux2_bfl_to_diffusers(_build_state_dict()) + fp8_weights = {k for k, v in sd.items() if v.dtype is FP8_DTYPE and k.endswith(".weight")} + + layers = extract_fp8_scaled_layers(sd) + + assert {f"{path}.weight" for path in layers} == fp8_weights + assert all(layer.input_scale is not None for layer in layers.values()) + + +def test_layer_hints_are_renamed_one_to_many_for_the_fused_qkv() -> None: + """Hints name layers in the BFL scheme; the scales are read after the rename. + + Without remapping, `full_precision_matrix_mult` matches nothing and is silently ignored. + """ + mapping = remap_flux2_layer_paths(["double_blocks.0.img_attn.qkv", "double_blocks.0.img_attn.proj"]) + + assert mapping["double_blocks.0.img_attn.qkv"] == [ + "transformer_blocks.0.attn.to_q", + "transformer_blocks.0.attn.to_k", + "transformer_blocks.0.attn.to_v", + ] + assert mapping["double_blocks.0.img_attn.proj"] == ["transformer_blocks.0.attn.to_out.0"] + + +def test_a_per_channel_scale_is_split_like_the_weight() -> None: + """Not present in this checkpoint, but the split must not assume every scale is a scalar.""" + sd = { + "double_blocks.0.img_attn.qkv.weight": torch.zeros(9, 2, dtype=torch.float32).to(FP8_DTYPE), + "double_blocks.0.img_attn.qkv.weight_scale": torch.arange(9, dtype=torch.float32), + } + + converted = convert_flux2_bfl_to_diffusers(sd) + + assert torch.equal(converted["transformer_blocks.0.attn.to_q.weight_scale"], torch.arange(0, 3.0)) + assert torch.equal(converted["transformer_blocks.0.attn.to_k.weight_scale"], torch.arange(3, 6.0)) + assert torch.equal(converted["transformer_blocks.0.attn.to_v.weight_scale"], torch.arange(6, 9.0)) diff --git a/tests/backend/model_manager/load/test_flux2_state_dict_utils.py b/tests/backend/model_manager/load/test_flux2_state_dict_utils.py index 8425118af11..a5c3c5972dc 100644 --- a/tests/backend/model_manager/load/test_flux2_state_dict_utils.py +++ b/tests/backend/model_manager/load/test_flux2_state_dict_utils.py @@ -99,3 +99,109 @@ def test_swaps_the_two_halves(self): def test_leaves_malformed_tensor_untouched(self): weight = torch.ones(3) # odd length -> cannot be split assert torch.allclose(_flux2_swap_scale_shift(weight), weight) + + +class TestFlux2RawFp8Gate: + """`_dequantize_fp8_weights` runs before `cast_state_dict`, so anything it converts is gone. + + Its trailing loop used to cast *every* float8 tensor unconditionally, which meant nothing fp8 + ever reached `cast_state_dict` and the FLUX.2 half of the raw-fp8 path could not execute. + """ + + def _dequantize(self, sd, keep_fp8): + from invokeai.backend.model_manager.load.model_loaders.flux import Flux2CheckpointModel + + # The method does not touch `self`; calling it unbound avoids building a whole loader. + return Flux2CheckpointModel._dequantize_fp8_weights(None, sd, keep_fp8=keep_fp8) + + def test_raw_fp8_linear_weights_survive_when_kept(self): + sd = { + "double_blocks.0.img_attn.qkv.weight": torch.zeros(48, 16).to(torch.float8_e4m3fn), + "double_blocks.0.img_attn.qkv.bias": torch.zeros(48).to(torch.float8_e4m3fn), + "double_blocks.0.img_norm.scale": torch.ones(16).to(torch.float8_e4m3fn), + } + out = self._dequantize(sd, keep_fp8=True) + assert out["double_blocks.0.img_attn.qkv.weight"].dtype is torch.float8_e4m3fn + # 1-D tensors are never usable on the tensor cores and must not stay quantized. + assert out["double_blocks.0.img_attn.qkv.bias"].dtype is torch.bfloat16 + assert out["double_blocks.0.img_norm.scale"].dtype is torch.bfloat16 + + def test_everything_is_dequantized_when_not_kept(self): + sd = {"double_blocks.0.img_attn.qkv.weight": torch.zeros(48, 16).to(torch.float8_e4m3fn)} + out = self._dequantize(sd, keep_fp8=False) + assert out["double_blocks.0.img_attn.qkv.weight"].dtype is torch.bfloat16 + + def test_scaled_fp8_is_still_folded_even_when_keeping(self): + """A weight with a `weight_scale` is dequantized *with* its scale, as before — only + scale-less fp8 is kept, because only that is safe to hand to `_scaled_mm` unscaled.""" + sd = { + "double_blocks.0.img_attn.qkv.weight": torch.ones(48, 16).to(torch.float8_e4m3fn), + "double_blocks.0.img_attn.qkv.weight_scale": torch.tensor(4.0), + } + out = self._dequantize(sd, keep_fp8=True) + assert out["double_blocks.0.img_attn.qkv.weight"].dtype is torch.bfloat16 + assert torch.allclose(out["double_blocks.0.img_attn.qkv.weight"], torch.full((48, 16), 4.0).bfloat16()) + assert "double_blocks.0.img_attn.qkv.weight_scale" not in out + + +class TestAdaLnSwapIsMirroredOntoTheScale: + """`final_layer.adaLN_modulation.1.weight` has its two halves swapped (BFL vs diffusers order). + + A per-output-channel `weight_scale` has one entry per row, so it has to be swapped identically. + Copying it verbatim leaves rows 0..n/2 holding the original second half while still carrying the + first half's scale factors - every row scaled by another row's factor. It is the one converter + transform that reorders rows and is not the fused-qkv split. + """ + + def test_a_per_channel_scale_is_swapped_with_its_weight(self) -> None: + weight = torch.arange(24, dtype=torch.float32).reshape(6, 4) + sd = { + "final_layer.adaLN_modulation.1.weight": weight, + "final_layer.adaLN_modulation.1.weight_scale": torch.arange(1, 7, dtype=torch.float32), + } + + converted = convert_flux2_bfl_to_diffusers(sd) + + assert converted["norm_out.linear.weight"][:, 0].tolist() == [12, 16, 20, 0, 4, 8] + assert converted["norm_out.linear.weight_scale"].tolist() == [4, 5, 6, 1, 2, 3] + + def test_the_scale_weight_spelling_is_swapped_too(self) -> None: + sd = { + "final_layer.adaLN_modulation.1.weight": torch.arange(24, dtype=torch.float32).reshape(6, 4), + "final_layer.adaLN_modulation.1.scale_weight": torch.arange(1, 7, dtype=torch.float32), + } + + converted = convert_flux2_bfl_to_diffusers(sd) + + assert converted["norm_out.linear.scale_weight"].tolist() == [4, 5, 6, 1, 2, 3] + + def test_per_tensor_scales_and_markers_are_left_alone(self) -> None: + """They describe the whole layer, so reordering them would be wrong. + + The `input_scale` is per-tensor by construction (Ada rejects per-row activation scaling) and + `comfy_quant` is a JSON byte blob, not a vector. + """ + sd = { + "final_layer.adaLN_modulation.1.weight": torch.arange(24, dtype=torch.float32).reshape(6, 4), + "final_layer.adaLN_modulation.1.weight_scale": torch.tensor(0.5), + "final_layer.adaLN_modulation.1.input_scale": torch.tensor(0.25), + "final_layer.adaLN_modulation.1.comfy_quant": torch.tensor(list(b'{"a":1}'), dtype=torch.uint8), + } + + converted = convert_flux2_bfl_to_diffusers(sd) + + assert converted["norm_out.linear.weight_scale"].item() == 0.5 + assert converted["norm_out.linear.input_scale"].item() == 0.25 + assert bytes(converted["norm_out.linear.comfy_quant"].tolist()) == b'{"a":1}' + + def test_other_layers_are_not_reordered(self) -> None: + """Only this one key is row-permuted; a scale elsewhere must be copied verbatim.""" + sd = { + "final_layer.linear.weight": torch.arange(24, dtype=torch.float32).reshape(6, 4), + "final_layer.linear.weight_scale": torch.arange(1, 7, dtype=torch.float32), + } + + converted = convert_flux2_bfl_to_diffusers(sd) + + scale = next(v for k, v in converted.items() if k.endswith(".weight_scale")) + assert scale.tolist() == [1, 2, 3, 4, 5, 6] diff --git a/tests/backend/model_manager/load/test_fp8_hint_prefixes.py b/tests/backend/model_manager/load/test_fp8_hint_prefixes.py new file mode 100644 index 00000000000..3c93e7349a2 --- /dev/null +++ b/tests/backend/model_manager/load/test_fp8_hint_prefixes.py @@ -0,0 +1,92 @@ +"""Layer hints must survive each loader's own key rewrites. + +`_quantization_metadata` names its layers in the checkpoint's scheme, while the scales are read +after the loader has stripped whatever prefix that checkpoint carries. A hint whose name still +carries the prefix matches no layer, so `full_precision_matrix_mult` is silently ignored and the +producer's "do not multiply this one in fp8" instruction is disregarded -- the exact failure the +hint plumbing exists to prevent. It has now been found twice (FLUX/Krea-2/Z-Image/Anima in round 2, +Mistral in round 3), so each loader's own prefix list gets a pin. +""" + +import torch + +from invokeai.backend.model_manager.load.model_loaders.mistral_encoder import ( + MISTRAL_KEY_PREFIXES, + _strip_known_prefixes, +) +from invokeai.backend.quantization.fp8_scaled import ( + FP8_DTYPE, + TRANSFORMER_KEY_PREFIXES, + extract_fp8_scaled_layers, + strip_layer_path_prefix, +) + + +def _quantized(path: str) -> dict[str, torch.Tensor]: + return { + f"{path}.weight": torch.zeros(16, 16, dtype=torch.float32).to(FP8_DTYPE), + f"{path}.weight_scale": torch.tensor(2.0), + } + + +class TestMistralHintPrefixes: + """The Mistral loader strips wrapper prefixes the generic tuple does not know about.""" + + def test_its_own_prefixes_are_not_covered_by_the_generic_tuple(self) -> None: + """If this ever becomes false the extra argument at the call site is dead weight.""" + assert not set(MISTRAL_KEY_PREFIXES) & set(TRANSFORMER_KEY_PREFIXES) + + def test_a_language_model_prefixed_hint_reaches_its_layer(self) -> None: + """Multimodal Mistral3 redistributions prefix every key with `language_model.`. + + The sd loses that prefix in `_strip_known_prefixes`; the header does not. Passing only the + generic tuple leaves the hint keyed on `language_model.model.layers.0.…` while the layer is + `model.layers.0.…`, and the flag is dropped in silence. + """ + sd = _strip_known_prefixes(_quantized("language_model.model.layers.0.self_attn.q_proj")) + header = {"language_model.model.layers.0.self_attn.q_proj": {"full_precision_matrix_mult": True}} + + hints = strip_layer_path_prefix(header, prefixes=(*MISTRAL_KEY_PREFIXES, *TRANSFORMER_KEY_PREFIXES)) + layers = extract_fp8_scaled_layers(sd, layer_hints=hints) + + assert set(layers) == {"model.layers.0.self_attn.q_proj"} + assert layers["model.layers.0.self_attn.q_proj"].full_precision_matmul is True + + def test_the_generic_tuple_alone_would_have_lost_the_flag(self) -> None: + """Pins the failure mode itself, so the fix cannot be reverted unnoticed.""" + sd = _strip_known_prefixes(_quantized("language_model.model.layers.0.self_attn.q_proj")) + header = {"language_model.model.layers.0.self_attn.q_proj": {"full_precision_matrix_mult": True}} + + hints = strip_layer_path_prefix(header) + layers = extract_fp8_scaled_layers(sd, layer_hints=hints) + + assert layers["model.layers.0.self_attn.q_proj"].full_precision_matmul is False + + +class TestFluxHintPrefixes: + """FLUX.1's hint plumbing had no test at all. + + ComfyUI FLUX.1 redistributions prefix their keys with `model.diffusion_model.`, which the + bundle conversion strips before the scales are read. + """ + + def test_a_prefixed_hint_lands_on_the_renamed_layer(self) -> None: + sd = { + k[len("model.diffusion_model.") :]: v + for k, v in _quantized("model.diffusion_model.double_blocks.0.img_attn.qkv").items() + } + header = {"model.diffusion_model.double_blocks.0.img_attn.qkv": {"full_precision_matrix_mult": True}} + + layers = extract_fp8_scaled_layers(sd, layer_hints=strip_layer_path_prefix(header)) + + assert set(layers) == {"double_blocks.0.img_attn.qkv"} + assert layers["double_blocks.0.img_attn.qkv"].full_precision_matmul is True + + def test_an_unprefixed_hint_is_passed_through_unchanged(self) -> None: + """A partially-prefixed header must not have its plain names dropped.""" + sd = _quantized("double_blocks.0.img_attn.qkv") + header = {"double_blocks.0.img_attn.qkv": {"full_precision_matrix_mult": True}} + + layers = extract_fp8_scaled_layers(sd, layer_hints=strip_layer_path_prefix(header)) + + assert layers["double_blocks.0.img_attn.qkv"].full_precision_matmul is True diff --git a/tests/backend/model_manager/load/test_krea2_loader_boundaries.py b/tests/backend/model_manager/load/test_krea2_loader_boundaries.py index ba618ab03b1..90eed6cbdbc 100644 --- a/tests/backend/model_manager/load/test_krea2_loader_boundaries.py +++ b/tests/backend/model_manager/load/test_krea2_loader_boundaries.py @@ -36,6 +36,8 @@ def test_single_file_loader_constructs_and_materializes_model(monkeypatch, tmp_p ram_cache = SimpleNamespace(make_room=MagicMock()) loader = object.__new__(Krea2CheckpointModel) loader._ram_cache = ram_cache + # ModelLoader.__init__ always sets a logger; this test bypasses __init__, so supply one. + loader._logger = MagicMock() loader._apply_fp8_layerwise_casting = lambda model, _config, _submodel: model monkeypatch.setattr(diffusers, "Krea2Transformer2DModel", _TinyKrea2Transformer, raising=False) diff --git a/tests/backend/model_manager/load/test_krea2_state_dict_utils.py b/tests/backend/model_manager/load/test_krea2_state_dict_utils.py index b5abded5d51..d9058268bb4 100644 --- a/tests/backend/model_manager/load/test_krea2_state_dict_utils.py +++ b/tests/backend/model_manager/load/test_krea2_state_dict_utils.py @@ -17,13 +17,19 @@ from invokeai.backend.model_manager.load.model_loaders.krea2 import ( KREA2_TRANSFORMER_CONFIG, _convert_krea2_native_to_diffusers, - _dequantize_scaled_fp8, _is_native_krea2_format, _normalize_qwen3vl_rope_config, _reject_incomplete_load, + _remap_native_layer_paths, _remap_qwen3vl_singlefile_keys, _strip_comfyui_prefix, ) +from invokeai.backend.quantization.fp8_scaled import ( + FP8_DTYPE, + detach_layer_sidechannel, + extract_fp8_scaled_layers, + reattach_layer_sidechannel, +) class TestNormalizeQwen3vlRopeConfig: @@ -84,48 +90,6 @@ def test_false_for_diffusers_keys(self, key: str) -> None: assert _is_native_krea2_format({key: torch.zeros(1)}) is False -class TestDequantizeScaledFp8: - def test_folds_scale_into_weight_and_drops_scale_key(self) -> None: - sd = { - "layer.weight": torch.tensor([2.0, 4.0]), - "layer.weight_scale": torch.tensor(0.5), - } - out = _dequantize_scaled_fp8(sd, torch.bfloat16) - assert "layer.weight_scale" not in out - assert torch.allclose(out["layer.weight"].float(), torch.tensor([1.0, 2.0])) - - def test_result_is_stored_in_the_compute_dtype_not_float32(self) -> None: - """The whole model must never be materialized in float32. - - The multiply runs in float32 for precision, but holding every dequantized weight there - costs 4 bytes per parameter: Krea-2's ~12 GB fp8 checkpoint peaked at ~50 GB of RAM before - the caller's later bf16 cast, which swaps a 32 GB machine during a cold load. - """ - sd = { - "layer.weight": torch.tensor([2.0, 4.0]), - "layer.weight_scale": torch.tensor(0.5), - } - assert _dequantize_scaled_fp8(dict(sd), torch.bfloat16)["layer.weight"].dtype is torch.bfloat16 - assert _dequantize_scaled_fp8(dict(sd), torch.float16)["layer.weight"].dtype is torch.float16 - - def test_dtype_is_required(self) -> None: - """No implicit bfloat16 fallback: on a float16-only device that would cost an extra rounding step.""" - with pytest.raises(TypeError): - _dequantize_scaled_fp8({"layer.weight": torch.tensor([2.0])}) # type: ignore[call-arg] - - def test_noop_without_scale_keys(self) -> None: - sd = {"layer.weight": torch.tensor([2.0, 4.0])} - out = _dequantize_scaled_fp8(sd, torch.bfloat16) - assert out is sd - - def test_orphan_scale_key_is_dropped(self) -> None: - # A scale key with no matching weight is simply removed (nothing to multiply). - sd = {"other.weight": torch.tensor([1.0]), "layer.weight_scale": torch.tensor(0.5)} - out = _dequantize_scaled_fp8(sd, torch.bfloat16) - assert "layer.weight_scale" not in out - assert "other.weight" in out - - class TestConvertKrea2NativeToDiffusers: def test_top_level_module_renames(self) -> None: sd = { @@ -376,3 +340,83 @@ def test_scale_shift_tables_match_real_module_dims(self) -> None: for name in table_keys: if name.startswith(("transformer_blocks.", "text_fusion.")): assert expected[name][0] == 6, f"{name} expected 6 modulation rows, got {expected[name]}" + + +class TestNativeConversionCarriesTheQuantizationSideChannel: + """Regression guard for scales orphaned by the native -> diffusers rename. + + `_convert_krea2_native_to_diffusers` renames `.weight` keys by substring *and* renames five + more by whole-key equality. Neither carries a sibling scale: `.scale_weight` and `.input_scale` + do not contain `.weight`, and `last.linear.weight_scale` is not equal to `last.linear.weight`. + An orphaned scale is then dropped by `extract_fp8_scaled_layers` (no fp8 weight sits at the old + path any more) and the weight is left quantized but unscaled - off by exactly 1/weight_scale, + with nothing logged, because the layer never enters `fp8_layers` for + `warn_on_unattached_scales` to see. + """ + + @staticmethod + def _convert(sd: dict) -> dict: + detached = detach_layer_sidechannel(sd) + converted = _convert_krea2_native_to_diffusers(sd) + orphaned = reattach_layer_sidechannel(converted, detached, _remap_native_layer_paths(set(detached))) + return converted, orphaned + + def test_every_spelling_and_rename_style_survives(self) -> None: + fp8 = torch.ones(4, 4).to(FP8_DTYPE) + sd = { + # substring rename, `.weight_scale` spelling + "blocks.0.attn.wq.weight": fp8.clone(), + "blocks.0.attn.wq.weight_scale": torch.tensor(0.5), + "blocks.0.attn.wq.input_scale": torch.tensor(0.25), + # substring rename, `.scale_weight` spelling + "blocks.0.attn.wk.weight": fp8.clone(), + "blocks.0.attn.wk.scale_weight": torch.tensor(0.5), + # whole-key-equality rename + "last.linear.weight": fp8.clone(), + "last.linear.weight_scale": torch.tensor(0.5), + } + + converted, orphaned = self._convert(sd) + layers = extract_fp8_scaled_layers(converted) + + assert set(layers) == { + "transformer_blocks.0.attn.to_q", + "transformer_blocks.0.attn.to_k", + "final_layer.linear", + } + assert orphaned == [] + # The calibrated activation scale must come across too, or every forward pays an amax pass. + assert layers["transformer_blocks.0.attn.to_q"].input_scale is not None + + def test_no_fp8_weight_is_left_without_its_scale(self) -> None: + """The property that actually matters, stated directly.""" + fp8 = torch.ones(4, 4).to(FP8_DTYPE) + sd = { + "blocks.0.attn.wk.weight": fp8.clone(), + "blocks.0.attn.wk.scale_weight": torch.tensor(0.5), + "last.linear.weight": fp8.clone(), + "last.linear.weight_scale": torch.tensor(0.5), + } + + converted, _ = self._convert(sd) + layers = extract_fp8_scaled_layers(converted) + + unscaled = [ + key + for key, value in converted.items() + if key.endswith(".weight") + and getattr(value, "dtype", None) is FP8_DTYPE + and key[: -len(".weight")] not in layers + ] + assert unscaled == [] + + def test_a_scale_on_a_dropped_module_is_reported(self) -> None: + """`last.down`/`last.up` have no diffusers counterpart, so their scales go too - loudly.""" + sd = { + "last.down.weight": torch.ones(4, 4).to(FP8_DTYPE), + "last.down.weight_scale": torch.tensor(0.5), + } + + _, orphaned = self._convert(sd) + + assert orphaned == ["last.down"] 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 cd87018aff7..30957c4a3a6 100644 --- a/tests/backend/model_manager/load/test_load_default_fp8.py +++ b/tests/backend/model_manager/load/test_load_default_fp8.py @@ -35,6 +35,7 @@ apply_custom_layers_to_model, ) from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType +from invokeai.backend.util.fp8 import FP8_COMPUTE_DTYPE_ATTR def _make_loader(device: str = "cuda") -> ModelLoader: @@ -748,3 +749,89 @@ def __init__(self): assert model.time_embedder.weight.dtype == torch.bfloat16 assert model.attn.weight.dtype == torch.float8_e4m3fn + + +class TestAlreadyFp8StorageGuard: + """FP8 storage must not run over weights that are already fp8 and headed for the tensor cores. + + Layerwise casting installs a pre-hook that restores the compute dtype before every forward. On + a scaled-fp8 checkpoint that hook does two things at once: `_can_use_fp8_matmul` no longer sees + an fp8 weight, so the matmul silently falls back, and the hook upcasts the weight *without* + applying its `weight_scale`, i.e. a weight off by `1/weight_scale`. Neither is visible in the + output of a successful generation, so the guard is load-bearing and needs a test that fails if + it is reverted. + """ + + def _model(self) -> torch.nn.Module: + model = torch.nn.Sequential(torch.nn.Linear(16, 32)) + model[0].weight = torch.nn.Parameter( + torch.zeros(32, 16, dtype=torch.float32).to(torch.float8_e4m3fn), requires_grad=False + ) + return model + + def test_a_model_whose_weights_are_already_fp8_is_left_alone(self) -> None: + loader = _make_loader("cuda") + model = self._model() + + with patch("invokeai.backend.model_manager.load.load_default.should_keep_fp8_weights", return_value=True): + result = loader._apply_fp8_layerwise_casting(model, _make_config(ModelType.Main, fp8=True)) + + assert result is model + # No compute-dtype marker means `_apply_fp8_to_nn_module` never ran over it. + assert getattr(model, FP8_COMPUTE_DTYPE_ATTR, None) is None + assert model[0].weight.dtype is torch.float8_e4m3fn + + def test_a_full_precision_model_is_still_cast(self) -> None: + """The guard must key on the weights, not merely on fp8_compute being enabled.""" + loader = _make_loader("cuda") + model = torch.nn.Sequential(torch.nn.Linear(16, 32)) + + with patch("invokeai.backend.model_manager.load.load_default.should_keep_fp8_weights", return_value=True): + loader._apply_fp8_layerwise_casting(model, _make_config(ModelType.Main, fp8=True)) + + assert model[0].weight.dtype is torch.float8_e4m3fn + assert getattr(model, FP8_COMPUTE_DTYPE_ATTR, None) is not None + + def test_ordinary_storage_casting_is_unaffected_when_the_matmul_is_unavailable(self) -> None: + """The guard must not become a blanket "skip fp8 storage" once fp8_compute is off. + + Only the already-fp8 case is protected. A full-precision model still gets the storage cast, + which is the entire point of the toggle on a card without the fp8 matmul. + + (The mirror case -- an already-fp8 model with the matmul *off* -- is unreachable: every + loader either folds the scales or casts raw fp8 away when `should_keep_fp8_weights` is + False, so nothing fp8 survives to reach this method. `set_fp8_compute_dtype` refuses it + outright rather than recording float8 as a compute dtype.) + """ + loader = _make_loader("cuda") + model = torch.nn.Sequential(torch.nn.Linear(16, 32)) + + with patch("invokeai.backend.model_manager.load.load_default.should_keep_fp8_weights", return_value=False): + loader._apply_fp8_layerwise_casting(model, _make_config(ModelType.Main, fp8=True)) + + assert model[0].weight.dtype is torch.float8_e4m3fn + + +class TestApplyFp8SkipCallback: + """The `skip=` callback keeps scaled-fp8 layers out of the storage cast. + + Its one caller (the Qwen3-VL encoder) casts the *unquantized* remainder of a partly-quantized + checkpoint to fp8 storage while leaving the scaled layers on the matmul path. Without the + callback those layers would be cast without their scale. + """ + + def test_a_module_the_callback_rejects_is_not_cast(self) -> None: + model = torch.nn.Sequential() + model.add_module("keep", torch.nn.Linear(16, 32)) + model.add_module("cast", torch.nn.Linear(16, 32)) + model.keep.weight_scale = torch.tensor(2.0) + + ModelLoader._apply_fp8_to_nn_module( + model, + storage_dtype=torch.float8_e4m3fn, + compute_dtype=torch.bfloat16, + skip=lambda _name, module: getattr(module, "weight_scale", None) is not None, + ) + + assert model.keep.weight.dtype is not torch.float8_e4m3fn + assert model.cast.weight.dtype is torch.float8_e4m3fn diff --git a/tests/backend/model_manager/load/test_scaled_fp8_fold_axis.py b/tests/backend/model_manager/load/test_scaled_fp8_fold_axis.py new file mode 100644 index 00000000000..2f08cb7a9e3 --- /dev/null +++ b/tests/backend/model_manager/load/test_scaled_fp8_fold_axis.py @@ -0,0 +1,92 @@ +"""Per-call-site guards for the scale axis used when folding a scaled-fp8 weight. + +A per-output-channel `weight_scale` is 1-D of length `out_features`, and `(out, in) * (out,)` +broadcasts on the *last* axis — so a bare multiply scales input channels instead of output +channels. That is a shape error on a non-square weight (loud) and a silently wrong weight on a +square one (not loud at all). Both loaders below used to do the bare multiply; they now go through +`expand_weight_scale`. + +`test_fp8_scaled.py` covers the helper itself. These pin the two call sites, so reverting either +one back to its local loop fails a test rather than passing CI. +""" + +from unittest.mock import MagicMock + +import pytest +import torch + +from invokeai.backend.model_manager.load.model_loaders.mistral_encoder import _drop_quantization_metadata +from invokeai.backend.model_manager.load.model_loaders.z_image import _fold_comfy_scaled_weights + + +def _per_channel_case(out_features: int = 4, in_features: int = 2): + """A non-square weight with a per-output-channel scale: rows 0..n-1 scaled by 1..n.""" + weight = torch.ones(out_features, in_features).to(torch.float8_e4m3fn) + scale = torch.arange(1, out_features + 1, dtype=torch.float32) + expected = torch.arange(1, out_features + 1, dtype=torch.bfloat16).reshape(-1, 1).expand(-1, in_features) + return weight, scale, expected + + +class TestMistralEncoderFold: + def test_per_channel_scale_multiplies_rows(self) -> None: + weight, scale, expected = _per_channel_case() + sd = {"layer.weight": weight, "layer.weight_scale": scale} + + _drop_quantization_metadata(sd, MagicMock(), target_dtype=torch.bfloat16) + + assert torch.equal(sd["layer.weight"], expected) + assert "layer.weight_scale" not in sd + + def test_square_weight_is_not_silently_transposed(self) -> None: + """The dangerous case: a square weight broadcasts happily on the wrong axis.""" + sd = { + "layer.weight": torch.ones(3, 3).to(torch.float8_e4m3fn), + "layer.weight_scale": torch.tensor([1.0, 2.0, 3.0]), + } + + _drop_quantization_metadata(sd, MagicMock(), target_dtype=torch.bfloat16) + + rows = torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0], [3.0, 3.0, 3.0]], dtype=torch.bfloat16) + assert torch.equal(sd["layer.weight"], rows), "scale must vary down the rows, not across them" + + +class TestZImageQwen3EncoderFold: + """The Z-Image Qwen3 encoder's single-file loader folds the same way.""" + + def _fold(self, sd: dict, dtype: torch.dtype = torch.bfloat16) -> dict: + # The real call site. It was a loop inside `_load_from_singlefile` -- unreachable without a + # checkpoint and a transformers model -- so it was extracted; re-inlining it as a local + # multiply fails these tests instead of passing CI. + _fold_comfy_scaled_weights(sd, dtype) + return sd + + def test_per_channel_scale_multiplies_rows(self) -> None: + weight, scale, expected = _per_channel_case() + sd = self._fold({"layer.weight": weight, "layer.weight_scale": scale}) + assert torch.equal(sd["layer.weight"], expected) + + def test_non_square_weight_does_not_raise(self) -> None: + weight, scale, _ = _per_channel_case(out_features=6, in_features=2) + sd = self._fold({"layer.weight": weight, "layer.weight_scale": scale}) + assert sd["layer.weight"].shape == (6, 2) + + def test_block_wise_scale_is_expanded_rather_than_rejected(self) -> None: + sd = self._fold( + { + "layer.weight": torch.ones(4, 2).to(torch.float8_e4m3fn), + "layer.weight_scale": torch.tensor([[1.0], [2.0]]), # one entry per 2-row block + } + ) + assert torch.equal( + sd["layer.weight"], + torch.tensor([[1.0, 1.0], [1.0, 1.0], [2.0, 2.0], [2.0, 2.0]], dtype=torch.bfloat16), + ) + + def test_a_scale_matching_neither_layout_is_reported(self) -> None: + with pytest.raises(ValueError, match="neither per-tensor nor per-output-channel"): + self._fold( + { + "layer.weight": torch.ones(4, 2).to(torch.float8_e4m3fn), + "layer.weight_scale": torch.full((3,), 2.0), + } + ) 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 index 631dc4a48fe..baf2594d93e 100644 --- a/tests/backend/model_manager/load/test_z_image_fp8_wiring.py +++ b/tests/backend/model_manager/load/test_z_image_fp8_wiring.py @@ -45,6 +45,10 @@ def _prepare_loader(monkeypatch, tmp_path, state_dict: dict[str, torch.Tensor]): loader = object.__new__(ZImageCheckpointModel) loader._ram_cache = SimpleNamespace(make_room=MagicMock()) + # The scaled-fp8 path reads the safetensors header (logging if it cannot) and asks the device + # whether the fp8 matmul is usable, so the stub needs both of those too. + loader._logger = MagicMock() + loader._torch_device = torch.device("cpu") return loader, config diff --git a/tests/backend/model_manager/load/test_z_image_scaled_fp8_keys.py b/tests/backend/model_manager/load/test_z_image_scaled_fp8_keys.py new file mode 100644 index 00000000000..2665255e6d7 --- /dev/null +++ b/tests/backend/model_manager/load/test_z_image_scaled_fp8_keys.py @@ -0,0 +1,70 @@ +"""The scaled-fp8 Z-Image key layout must be recognized, scales and all. + +A ComfyUI "scaled fp8" checkpoint stores each quantized Linear as an fp8 `.weight` plus a +`.scale_weight`. Before the scaled-fp8 handling landed, the Z-Image loader deleted those +scale keys in its `keys_to_remove` filter and cast the weight — so every quantized weight was off +by a factor of `weight_scale`, with nothing logged. On the captured checkpoint the scales run from +1.5 to 7.6, which leaves each weight at a *different* fraction of its true magnitude. + +The fixture is a real key layout; the tensors are synthetic because only shapes and dtypes matter +to the code under test. +""" + +import torch + +from invokeai.backend.quantization.fp8_scaled import ( + FP8_DTYPE, + extract_comfy_quant_hints, + extract_fp8_scaled_layers, +) +from tests.backend.model_manager.load.state_dicts.z_image_transformer_scaled_fp8_keys import ( + state_dict_keys as scaled_keys, +) + +_DTYPES = {"F8_E4M3": FP8_DTYPE, "F32": torch.float32, "BF16": torch.bfloat16} + + +def _mock_state_dict(scale_value: float = 4.0) -> dict[str, torch.Tensor]: + sd: dict[str, torch.Tensor] = {} + for key, (shape, dtype) in scaled_keys.items(): + if key.endswith(".scale_weight"): + sd[key] = torch.full(shape, scale_value, dtype=torch.float32) + else: + sd[key] = torch.ones(shape, dtype=torch.float32).to(_DTYPES[dtype]) + return sd + + +def test_the_fixture_really_is_scaled_fp8() -> None: + """Guard the fixture itself: a bf16 recapture would make every assertion below vacuous.""" + assert any(k.endswith(".scale_weight") for k in scaled_keys) + assert not any(k.endswith(".weight_scale") for k in scaled_keys) + assert any(dtype == "F8_E4M3" for _, dtype in scaled_keys.values()) + assert "scaled_fp8" in scaled_keys + + +def test_every_scale_weight_is_recognized_as_a_scaled_layer() -> None: + sd = _mock_state_dict() + expected = {k[: -len(".scale_weight")] for k in scaled_keys if k.endswith(".scale_weight")} + + layers = extract_fp8_scaled_layers(sd, layer_hints=extract_comfy_quant_hints(sd)) + + assert set(layers) == expected, "a .scale_weight spelling was not recognized" + + +def test_extraction_consumes_the_scale_keys() -> None: + # Anything left behind reaches `load_state_dict(..., strict=True)` as an unexpected key. + sd = _mock_state_dict() + + extract_fp8_scaled_layers(sd, layer_hints=extract_comfy_quant_hints(sd)) + + assert not [k for k in sd if k.endswith((".scale_weight", ".weight_scale"))] + + +def test_the_scale_is_carried_through_not_discarded() -> None: + sd = _mock_state_dict(scale_value=4.0) + + layers = extract_fp8_scaled_layers(sd, layer_hints=extract_comfy_quant_hints(sd)) + + assert layers + for layer in layers.values(): + assert layer.weight_scale.float().reshape(-1)[0].item() == 4.0 diff --git a/tests/backend/model_manager/load/test_z_image_state_dict_utils.py b/tests/backend/model_manager/load/test_z_image_state_dict_utils.py index 37e113d8570..42b43c70cee 100644 --- a/tests/backend/model_manager/load/test_z_image_state_dict_utils.py +++ b/tests/backend/model_manager/load/test_z_image_state_dict_utils.py @@ -1,8 +1,12 @@ """Unit tests for the Z-Image GGUF/ComfyUI -> diffusers state-dict converter.""" +import pytest import torch -from invokeai.backend.model_manager.load.model_loaders.z_image import _convert_z_image_gguf_to_diffusers +from invokeai.backend.model_manager.load.model_loaders.z_image import ( + _convert_z_image_gguf_to_diffusers, + _remap_z_image_layer_paths, +) from tests.backend.model_manager.load.state_dicts.utils import keys_to_mock_state_dict from tests.backend.model_manager.load.state_dicts.z_image_transformer_comfyui_keys import ( state_dict_keys as z_image_keys, @@ -64,3 +68,83 @@ def test_qkv_split_preserves_values(self): assert torch.allclose(out["blk.attention.to_q.weight"], qkv[0:2]) assert torch.allclose(out["blk.attention.to_k.weight"], qkv[2:4]) assert torch.allclose(out["blk.attention.to_v.weight"], qkv[4:6]) + + +class TestQkvQuantizationSideChannel: + """A scaled-fp8 checkpoint puts a `scale_weight` next to the fused `qkv.weight`. + + Left on `...attention.qkv`, the recovered scale is keyed on a module path the diffusers model + does not have, so `attach_fp8_scales` finds nothing and the three split weights stay quantized + but *unscaled* — off by 1/weight_scale, with no error anywhere. + """ + + def test_per_tensor_scale_reaches_all_three_projections(self): + out = _convert_z_image_gguf_to_diffusers( + { + "blk.attention.qkv.weight": torch.arange(12, dtype=torch.float32).reshape(6, 2), + "blk.attention.qkv.scale_weight": torch.tensor(0.25), + } + ) + assert not any(".attention.qkv." in k for k in out) + for name in ("to_q", "to_k", "to_v"): + assert torch.equal(out[f"blk.attention.{name}.scale_weight"], torch.tensor(0.25)) + + def test_per_channel_scale_is_split_like_the_weight(self): + out = _convert_z_image_gguf_to_diffusers( + { + "blk.attention.qkv.weight": torch.arange(12, dtype=torch.float32).reshape(6, 2), + "blk.attention.qkv.weight_scale": torch.arange(6, dtype=torch.float32), + } + ) + assert torch.equal(out["blk.attention.to_q.weight_scale"], torch.tensor([0.0, 1.0])) + assert torch.equal(out["blk.attention.to_k.weight_scale"], torch.tensor([2.0, 3.0])) + assert torch.equal(out["blk.attention.to_v.weight_scale"], torch.tensor([4.0, 5.0])) + + def test_marker_blob_is_copied_not_split(self): + # `.comfy_quant` is a 1-D JSON byte string describing the layer, not a per-channel vector. + blob = torch.frombuffer(b'{"format":"float8_e4m3fn"}', dtype=torch.uint8).clone() + out = _convert_z_image_gguf_to_diffusers( + { + "blk.attention.qkv.weight": torch.arange(12, dtype=torch.float32).reshape(6, 2), + "blk.attention.qkv.comfy_quant": blob, + } + ) + for name in ("to_q", "to_k", "to_v"): + assert torch.equal(out[f"blk.attention.{name}.comfy_quant"], blob) + + def test_unknown_suffix_is_left_alone(self): + out = _convert_z_image_gguf_to_diffusers( + { + "blk.attention.qkv.weight": torch.arange(12, dtype=torch.float32).reshape(6, 2), + "blk.attention.qkv.something_else": torch.tensor(1.0), + } + ) + assert "blk.attention.qkv.something_else" in out + + def test_undivisible_scale_is_rejected_rather_than_mis_split(self): + with pytest.raises(ValueError, match="Cannot split fused QKV quantization data"): + _convert_z_image_gguf_to_diffusers( + { + "blk.attention.qkv.weight": torch.arange(12, dtype=torch.float32).reshape(6, 2), + "blk.attention.qkv.weight_scale": torch.arange(4, dtype=torch.float32), + } + ) + + +class TestMetadataPathRemap: + """`_quantization_metadata` names layers in the checkpoint's scheme; the scales are recovered + after the rename, so the per-layer hints have to follow the same route.""" + + def test_renamed_layers_map_one_to_one(self): + mapping = _remap_z_image_layer_paths(["x_embedder", "final_layer.linear", "layers.0.attention.out"]) + assert mapping["x_embedder"] == ["all_x_embedder.2-1"] + assert mapping["final_layer.linear"] == ["all_final_layer.2-1.linear"] + assert mapping["layers.0.attention.out"] == ["layers.0.attention.to_out.0"] + + def test_fused_qkv_maps_to_all_three_projections(self): + mapping = _remap_z_image_layer_paths(["layers.0.attention.qkv"]) + assert mapping["layers.0.attention.qkv"] == [ + "layers.0.attention.to_q", + "layers.0.attention.to_k", + "layers.0.attention.to_v", + ] diff --git a/tests/backend/quantization/test_fp8_scaled.py b/tests/backend/quantization/test_fp8_scaled.py new file mode 100644 index 00000000000..13a5a432504 --- /dev/null +++ b/tests/backend/quantization/test_fp8_scaled.py @@ -0,0 +1,1292 @@ +import contextlib +import logging +from unittest import mock + +import pytest +import torch + +from invokeai.backend.quantization.fp8_scaled import ( + FP8_DTYPE, + Fp8ScaledLayer, + attach_fp8_scales, + cast_state_dict, + count_fp8_weights, + dequantize_fp8_scaled, + dequantize_weight, + detach_layer_sidechannel, + device_supports_fp8_matmul, + expand_weight_scale, + extract_comfy_quant_hints, + extract_fp8_scaled_layers, + is_matmul_usable_scale, + is_scale_metadata_key, + iter_weight_scale_pairs, + parse_quantization_metadata, + predict_cast_state_dict_size, + reattach_layer_sidechannel, + reset_fp8_matmul_support_cache, + scaled_mm_linear, + set_fp8_matmul_enabled, + set_full_precision_hints_respected, + split_fp8_scaled_layers, + strip_layer_path_prefix, + warn_on_unattached_scales, +) + +cuda_fp8 = pytest.mark.skipif( + not (torch.cuda.is_available() and device_supports_fp8_matmul(torch.device("cuda"))), + reason="requires a CUDA device with fp8 tensor cores (SM 8.9+)", +) + + +def _fp8_weight(out_f: int, in_f: int, per_channel: bool = False): + w = torch.randn(out_f, in_f, dtype=torch.bfloat16) * 0.02 + if per_channel: + scale = (w.abs().amax(dim=1) / torch.finfo(FP8_DTYPE).max).float().clamp(min=1e-12) + q = (w / scale.reshape(-1, 1)).to(FP8_DTYPE) + else: + scale = (w.abs().amax() / torch.finfo(FP8_DTYPE).max).float().clamp(min=1e-12) + q = (w / scale).to(FP8_DTYPE) + return q, scale + + +class TestExtract: + def test_pops_scales_and_keys_layer_by_path(self): + q, scale = _fp8_weight(32, 16) + sd = {"blk.0.lin.weight": q, "blk.0.lin.weight_scale": scale, "blk.0.lin.bias": torch.zeros(32)} + layers = extract_fp8_scaled_layers(sd) + assert set(layers) == {"blk.0.lin"} + assert "blk.0.lin.weight_scale" not in sd, "scale keys must be removed so the sd loads cleanly" + assert sd["blk.0.lin.weight"].dtype == FP8_DTYPE, "weights must stay quantized" + + def test_accepts_scale_weight_suffix(self): + q, scale = _fp8_weight(32, 16) + sd = {"lin.weight": q, "lin.scale_weight": scale} + assert set(extract_fp8_scaled_layers(sd)) == {"lin"} + + def test_input_scale_is_captured(self): + q, scale = _fp8_weight(32, 16) + sd = {"lin.weight": q, "lin.weight_scale": scale, "lin.input_scale": torch.tensor(0.5)} + assert extract_fp8_scaled_layers(sd)["lin"].input_scale == pytest.approx(0.5) + + def test_scale_without_fp8_weight_is_dropped(self): + """A scale applied to an already-dequantized weight would corrupt it.""" + sd = {"lin.weight": torch.randn(32, 16, dtype=torch.bfloat16), "lin.weight_scale": torch.tensor(2.0)} + assert extract_fp8_scaled_layers(sd) == {} + assert "lin.weight_scale" not in sd + + def test_strips_stray_marker_keys(self): + q, scale = _fp8_weight(32, 16) + sd = { + "lin.weight": q, + "lin.weight_scale": scale, + "scaled_fp8": torch.tensor(0.0), + "comfy_quant_x": torch.tensor(1), + } + extract_fp8_scaled_layers(sd) + assert set(sd) == {"lin.weight"} + + def test_full_precision_hint_from_metadata(self): + q, scale = _fp8_weight(32, 16) + sd = {"a.weight": q, "a.weight_scale": scale, "b.weight": q.clone(), "b.weight_scale": scale.clone()} + meta = {"_quantization_metadata": '{"layers": {"a": {"full_precision_matrix_mult": true}, "b": {}}}'} + layers = extract_fp8_scaled_layers(sd, meta) + assert layers["a"].full_precision_matmul is True + assert layers["b"].full_precision_matmul is False + + def test_malformed_metadata_is_ignored(self): + assert parse_quantization_metadata({"_quantization_metadata": "not json"}) == {} + assert parse_quantization_metadata(None) == {} + + def test_layer_hints_override_metadata(self): + q, scale = _fp8_weight(32, 16) + sd = {"renamed.lin.weight": q, "renamed.lin.weight_scale": scale} + # Metadata uses the pre-rename path, so only the explicitly remapped hints can match. + meta = {"_quantization_metadata": '{"layers": {"native.lin": {"full_precision_matrix_mult": true}}}'} + layers = extract_fp8_scaled_layers( + dict(sd), metadata=meta, layer_hints={"renamed.lin": {"full_precision_matrix_mult": True}} + ) + assert layers["renamed.lin"].full_precision_matmul is True + + # Without the remap the flag silently matches nothing - the regression this guards against. + layers = extract_fp8_scaled_layers(dict(sd), metadata=meta) + assert layers["renamed.lin"].full_precision_matmul is False + + +def _comfy_quant_blob(full_precision: bool): + """The per-layer marker some ComfyUI exports write instead of the header entry.""" + payload = f'{{"format": "float8_e4m3fn", "full_precision_matrix_mult": {"true" if full_precision else "false"}}}' + return torch.tensor(list(payload.encode("utf-8")), dtype=torch.uint8) + + +class TestRawFp8: + """Checkpoints that ship fp8 weights with no weight_scale at all.""" + + def test_cast_state_dict_preserves_fp8_when_asked(self): + q, _ = _fp8_weight(32, 16) + sd = {"lin.weight": q, "lin.bias": torch.zeros(32), "norm.weight": torch.ones(16)} + kept = cast_state_dict(sd, torch.bfloat16, keep_fp8=True) + assert kept == 1 + assert sd["lin.weight"].dtype is FP8_DTYPE, "raw fp8 must survive the load" + assert sd["lin.bias"].dtype is torch.bfloat16 + assert sd["norm.weight"].dtype is torch.bfloat16 + + def test_cast_state_dict_dequantizes_when_matmul_unavailable(self): + """Without the matmul, staying quantized costs a dequantize per forward for no gain.""" + q, _ = _fp8_weight(32, 16) + sd = {"lin.weight": q} + assert cast_state_dict(sd, torch.bfloat16, keep_fp8=False) == 0 + assert sd["lin.weight"].dtype is torch.bfloat16 + + def test_e5m2_is_never_preserved(self): + """scaled_mm cannot take e5m2 as the weight operand on Ada, so keeping it buys nothing.""" + sd = {"lin.weight": torch.zeros(32, 16).to(torch.float8_e5m2)} + assert cast_state_dict(sd, torch.bfloat16, keep_fp8=True) == 0 + assert sd["lin.weight"].dtype is torch.bfloat16 + + def test_only_linear_weights_stay_quantized(self): + """The regression an end-to-end run caught: checkpoints exist that quantize *everything*. + + A Z-Image checkpoint had 243 of its 453 fp8 tensors 1-D — biases, norm weights, a learned + pad token. Keeping those quantized saves nothing usable and breaks inference: the fp8 value + flows into the activations and the next Linear receives an fp8 *input*, which dies in + `x.abs()` with `"abs_cuda" not implemented for 'Float8_e4m3fn'`. + """ + model = torch.nn.Sequential() + model.add_module("lin", torch.nn.Linear(16, 32)) + model.add_module("norm", torch.nn.LayerNorm(32)) + q, _ = _fp8_weight(32, 16) + sd = { + "lin.weight": q, # keep: a Linear weight the matmul can use + "lin.bias": torch.zeros(32).to(FP8_DTYPE), # dequantize: bias + "norm.weight": torch.ones(32).to(FP8_DTYPE), # dequantize: 1-D norm + "pad_token": torch.zeros(1, 32).to(FP8_DTYPE), # dequantize: not a module weight + } + assert cast_state_dict(sd, torch.bfloat16, keep_fp8=True, model=model) == 1 + assert sd["lin.weight"].dtype is FP8_DTYPE + for key in ("lin.bias", "norm.weight", "pad_token"): + assert sd[key].dtype is torch.bfloat16, f"{key} must not stay fp8" + + def test_skip_patterns_dequantize_named_modules(self): + """Modules a model marks precision-sensitive must be dequantized even if they are Linears. + + Z-Image declares `_skip_layerwise_casting_patterns = ["t_embedder", "cap_embedder"]`, and + `TimestepEmbedder.forward` casts its activations to `self.mlp[0].weight.dtype` — an fp8 + weight there makes the activations fp8. + """ + model = torch.nn.Sequential() + model.add_module("t_embedder", torch.nn.Linear(16, 32)) + model.add_module("blocks", torch.nn.Linear(16, 32)) + q, _ = _fp8_weight(32, 16) + sd = {"t_embedder.weight": q, "blocks.weight": q.clone()} + assert cast_state_dict(sd, torch.bfloat16, keep_fp8=True, model=model, skip_patterns=["t_embedder"]) == 1 + assert sd["t_embedder.weight"].dtype is torch.bfloat16 + assert sd["blocks.weight"].dtype is FP8_DTYPE + + def test_count_fp8_weights(self): + model = torch.nn.Sequential(torch.nn.Linear(16, 32, bias=False), torch.nn.Linear(32, 16, bias=False)) + assert count_fp8_weights(model) == 0 + model[0].weight = torch.nn.Parameter(_fp8_weight(32, 16)[0], requires_grad=False) + assert count_fp8_weights(model) == 1 + + @cuda_fp8 + def test_fp8_weight_without_scale_uses_the_tensor_cores(self): + """The runtime already supports scale-less fp8 — `weight_scale` is optional in + `scaled_mm_linear`, and `_can_use_fp8_matmul` only requires the fp8 dtype. This pins that + down so a future change cannot quietly make a scale mandatory.""" + from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import ( + apply_custom_layers_to_model, + ) + + torch.manual_seed(0) + net = torch.nn.Sequential(torch.nn.Linear(64, 64, bias=False)).cuda() + net[0].weight.data = (net[0].weight.data * 0.02).to(FP8_DTYPE) + apply_custom_layers_to_model(net) + lin = net[0] + x = torch.randn(1, 32, 64, device="cuda", dtype=torch.bfloat16) + + set_fp8_matmul_enabled(True) + try: + assert getattr(lin, "weight_scale", None) is None + assert lin._can_use_fp8_matmul(x) is True + out = net(x) + finally: + set_fp8_matmul_enabled(None) + + reference = torch.nn.functional.linear(x, dequantize_weight(lin.weight, None, x.dtype)) + rel = ((out.float() - reference.float()).norm() / reference.float().norm()).item() + assert torch.isfinite(out).all() + assert rel < 0.1, f"unit-scaled fp8 matmul drifted too far from bf16: {rel:.4f}" + + +class TestInputScale: + def test_calibrated_scale_is_kept(self): + q, scale = _fp8_weight(32, 16) + sd = {"lin.weight": q, "lin.weight_scale": scale, "lin.input_scale": torch.tensor(0.017)} + layer = extract_fp8_scaled_layers(sd)["lin"] + assert layer.input_scale is not None + assert pytest.approx(layer.input_scale.item()) == 0.017 + assert "lin.input_scale" not in sd + + def test_scale_input_spelling_is_accepted(self): + """`.scale_input` is the other spelling in the wild; ignoring it discards the calibration.""" + q, scale = _fp8_weight(32, 16) + sd = {"lin.weight": q, "lin.scale_weight": scale, "lin.scale_input": torch.tensor(0.017)} + layer = extract_fp8_scaled_layers(sd)["lin"] + assert pytest.approx(layer.input_scale.item()) == 0.017 + assert not [k for k in sd if "scale_input" in k], "must be popped, not left for load_state_dict" + + @pytest.mark.parametrize( + "value", [1.0, 0.0, -0.5, float("nan"), float("inf")], ids=["placeholder", "zero", "negative", "nan", "inf"] + ) + def test_unusable_scales_fall_back_to_dynamic(self, value: float): + """A 1.0 input_scale is an uncalibrated placeholder. Using it means *no* activation scaling, + so everything above the fp8 max saturates -- strictly worse than the per-forward amax it + would replace. Zero/negative/non-finite cannot be a divisor at all.""" + q, scale = _fp8_weight(32, 16) + sd = {"lin.weight": q, "lin.weight_scale": scale, "lin.input_scale": torch.tensor(value)} + assert extract_fp8_scaled_layers(sd)["lin"].input_scale is None + + @cuda_fp8 + def test_placeholder_scale_saturates_activations_above_the_fp8_range(self): + """End-to-end cost of trusting a 1.0 placeholder. + + fp8_e4m3 is a floating-point format, so a scale factor does not buy relative precision the + way it would for int8 — for activations inside +/-448 both paths are equivalent. The damage + appears only above the representable range, where an unscaled cast clamps hard while the + dynamic amax scale maps the whole tensor into range. Transformer activations do reach that + regime, which is why a calibrated input_scale exists at all. + """ + torch.manual_seed(0) + q, scale = _fp8_weight(64, 64) + q, scale = q.cuda(), scale.cuda() + x = (torch.randn(32, 64, dtype=torch.bfloat16, device="cuda") * 2000.0).unsqueeze(0) + assert x.abs().max() > 448, "test must exercise the saturating regime" + reference = torch.nn.functional.linear(x, dequantize_weight(q, scale, x.dtype)) + + def err(got: torch.Tensor) -> float: + return ((got.float() - reference.float()).norm() / reference.float().norm()).item() + + dynamic = scaled_mm_linear(x, q, scale, None, input_scale=None) + unscaled = scaled_mm_linear(x, q, scale, None, input_scale=torch.tensor(1.0, device="cuda")) + assert err(dynamic) < err(unscaled) / 2, f"dynamic={err(dynamic):.4f} unscaled={err(unscaled):.4f}" + + +class TestComfyQuantHints: + def test_reads_per_layer_markers_and_pops_them(self): + sd = { + "blk.0.lin.weight": _fp8_weight(32, 16)[0], + "blk.0.lin.comfy_quant": _comfy_quant_blob(True), + "blk.1.lin.comfy_quant": _comfy_quant_blob(False), + } + hints = extract_comfy_quant_hints(sd) + assert hints["blk.0.lin"]["full_precision_matrix_mult"] is True + assert hints["blk.1.lin"]["full_precision_matrix_mult"] is False + assert not [k for k in sd if "comfy_quant" in k], "markers must not reach load_state_dict" + + def test_full_precision_flag_reaches_the_layer(self): + """The regression this guards: a checkpoint carrying the flags *only* in per-layer markers + had them silently ignored, so layers the producer marked unsafe were multiplied in fp8.""" + q, scale = _fp8_weight(32, 16) + sd = {"blk.0.lin.weight": q, "blk.0.lin.weight_scale": scale, "blk.0.lin.comfy_quant": _comfy_quant_blob(True)} + hints = extract_comfy_quant_hints(sd) + layers = extract_fp8_scaled_layers(sd, layer_hints=hints) + assert layers["blk.0.lin"].full_precision_matmul is True + + # Reading only the header (no hints) is what used to happen - the flag matches nothing. + sd2 = {"blk.0.lin.weight": q, "blk.0.lin.weight_scale": scale, "blk.0.lin.comfy_quant": _comfy_quant_blob(True)} + assert extract_fp8_scaled_layers(sd2).get("blk.0.lin").full_precision_matmul is False + + def test_nul_padded_and_malformed_blobs(self): + padded = torch.cat([_comfy_quant_blob(True), torch.zeros(8, dtype=torch.uint8)]) + sd = { + "a.comfy_quant": padded, + "b.comfy_quant": torch.tensor(list(b"not json"), dtype=torch.uint8), + } + hints = extract_comfy_quant_hints(sd) + assert hints["a"]["full_precision_matrix_mult"] is True + # A malformed marker is a lost hint, never a failed load. + assert "b" not in hints + assert not sd + + +class TestQwen3VLKeyRemap: + def test_scale_keys_and_hint_paths_land_on_the_same_module(self): + """attach_fp8_scales resolves hint paths against the *model*, so the state-dict remap and the + hint remap must agree - otherwise every recovered scale silently matches nothing.""" + from invokeai.backend.model_manager.load.model_loaders.krea2 import ( + _qwen3vl_target_key, + _remap_qwen3vl_singlefile_keys, + ) + + q, scale = _fp8_weight(32, 16) + sd = _remap_qwen3vl_singlefile_keys( + { + "model.layers.0.mlp.down_proj.weight": q, + "model.layers.0.mlp.down_proj.weight_scale": scale, + "model.visual.blocks.0.attn.qkv.weight": torch.zeros(16, 16), + } + ) + assert "language_model.layers.0.mlp.down_proj.weight_scale" in sd + assert "visual.blocks.0.attn.qkv.weight" in sd + + hints = {_qwen3vl_target_key("model.layers.0.mlp.down_proj"): {"full_precision_matrix_mult": True}} + layers = extract_fp8_scaled_layers(sd, layer_hints=hints) + assert set(layers) == {"language_model.layers.0.mlp.down_proj"} + assert layers["language_model.layers.0.mlp.down_proj"].full_precision_matmul is True + + +class TestFullPrecisionHintToggle: + def test_marker_is_applied_by_default(self): + q, scale = _fp8_weight(32, 16) + module = torch.nn.Linear(16, 32, bias=False) + module.weight = torch.nn.Parameter(q, requires_grad=False) + model = torch.nn.Sequential() + model.add_module("lin", module) + attach_fp8_scales(model, {"lin": Fp8ScaledLayer(weight_scale=scale, full_precision_matmul=True)}) + assert module._fp8_full_precision_matmul is True + + def test_marker_suppressed_when_hints_are_off(self): + """Turning the hints off must reach the module flag CustomLinear reads, not just the parse.""" + q, scale = _fp8_weight(32, 16) + module = torch.nn.Linear(16, 32, bias=False) + module.weight = torch.nn.Parameter(q, requires_grad=False) + model = torch.nn.Sequential() + model.add_module("lin", module) + set_full_precision_hints_respected(False) + try: + attach_fp8_scales(model, {"lin": Fp8ScaledLayer(weight_scale=scale, full_precision_matmul=True)}) + finally: + set_full_precision_hints_respected(None) + assert module._fp8_full_precision_matmul is False + + +class TestKrea2MetadataRemap: + def test_native_layer_paths_are_remapped_like_the_state_dict(self): + """The quantization metadata names layers natively; the scales are keyed after renaming.""" + from invokeai.backend.model_manager.load.model_loaders.krea2 import _remap_native_layer_paths + + mapping = _remap_native_layer_paths( + ["blocks.0.attn.wq", "blocks.0.attn.wo", "blocks.3.mlp.down", "txtfusion.refiner_blocks.1.attn.wk"] + ) + assert mapping["blocks.0.attn.wq"] == "transformer_blocks.0.attn.to_q" + assert mapping["blocks.0.attn.wo"] == "transformer_blocks.0.attn.to_out.0" + assert mapping["blocks.3.mlp.down"] == "transformer_blocks.3.ff.down" + assert mapping["txtfusion.refiner_blocks.1.attn.wk"] == "text_fusion.refiner_blocks.1.attn.to_k" + + +class TestSplitScaledLayers: + """A scaled layer must never reach `cast_state_dict` still holding an unapplied scale.""" + + def _model(self): + model = torch.nn.Sequential() + model.add_module("time_embed", torch.nn.Sequential()) + model.time_embed.add_module("linear_2", torch.nn.Linear(16, 32)) + model.add_module("attn", torch.nn.Linear(16, 32)) + return model + + def test_skip_pattern_layer_is_dequantized_with_its_scale(self): + """The Krea-2 regression: `time_embed` matches the model's skip patterns, so the cast would + turn its fp8 weight into bf16 *codes* — the scale silently dropped, the weight off by + 1/weight_scale, and `attach_fp8_scales` unable to repair it afterwards.""" + model = self._model() + w = torch.randn(32, 16, dtype=torch.bfloat16) * 0.02 + scale = (w.abs().amax() / torch.finfo(FP8_DTYPE).max).float().clamp(min=1e-12) + q = (w / scale).to(FP8_DTYPE) + sd = {"time_embed.linear_2.weight": q, "attn.weight": q.clone()} + layers = { + "time_embed.linear_2": Fp8ScaledLayer(weight_scale=scale), + "attn": Fp8ScaledLayer(weight_scale=scale), + } + + remaining = split_fp8_scaled_layers( + sd, layers, torch.bfloat16, model=model, skip_patterns=["time_embed", "norm"] + ) + + assert set(remaining) == {"attn"}, "only the layer that can stay quantized is left to attach" + assert sd["time_embed.linear_2.weight"].dtype is torch.bfloat16 + rel = ((sd["time_embed.linear_2.weight"].float() - w.float()).norm() / w.float().norm()).item() + assert rel < 0.05, f"the scale was not applied on the way down: rel-err {rel:.4f}" + + # And the survivor is untouched by the subsequent cast, so its scale still attaches. + assert cast_state_dict(sd, torch.bfloat16, keep_fp8=True, model=model, skip_patterns=["time_embed"]) == 1 + model.load_state_dict(sd, assign=True, strict=False) + assert attach_fp8_scales(model, remaining) == 1 + assert torch.equal(model.attn.weight_scale, scale) + + def test_non_linear_scaled_weight_is_dequantized_with_its_scale(self): + """Same failure without any skip pattern: a scaled tensor that is not an nn.Linear weight + fails `_is_fp8_matmul_weight`, so the cast would strip its scale.""" + model = self._model() + q, scale = _fp8_weight(32, 16) + sd = {"pad_token.weight": q} + remaining = split_fp8_scaled_layers( + sd, {"pad_token": Fp8ScaledLayer(weight_scale=scale)}, torch.bfloat16, model=model + ) + assert remaining == {} + assert torch.equal(sd["pad_token.weight"], dequantize_weight(q, scale, torch.bfloat16)) + + def test_warns_when_a_scale_finds_no_module(self): + logger = logging.getLogger("fp8-test") + with mock.patch.object(logger, "warning") as warn: + warn_on_unattached_scales(logger, "Krea-2", 1, {"a": object(), "b": object()}) + assert warn.call_count == 1 + assert "1 of 2" in warn.call_args[0][0] + + def test_silent_when_every_scale_landed(self): + logger = logging.getLogger("fp8-test") + with mock.patch.object(logger, "warning") as warn: + warn_on_unattached_scales(logger, "Krea-2", 2, {"a": object(), "b": object()}) + assert warn.call_count == 0 + + +class TestPredictedSize: + def test_matches_what_cast_state_dict_actually_leaves(self): + """`make_room` reserves against this, so a mismatch is a silent under-reservation.""" + model = torch.nn.Sequential() + model.add_module("lin", torch.nn.Linear(16, 32)) + model.add_module("t_embedder", torch.nn.Linear(16, 32)) + model.add_module("norm", torch.nn.LayerNorm(32)) + q, _ = _fp8_weight(32, 16) + sd = { + "lin.weight": q, + "lin.bias": torch.zeros(32).to(FP8_DTYPE), + "t_embedder.weight": q.clone(), + "norm.weight": torch.ones(32).to(FP8_DTYPE), + "pad_token": torch.zeros(1, 32).to(FP8_DTYPE), + } + kwargs = {"model": model, "skip_patterns": ["t_embedder"]} + # What the loaders used to reserve: 1 byte/element for *every* fp8 tensor. + naive = sum(t.nelement() * (t.element_size() if t.dtype is FP8_DTYPE else 2) for t in sd.values()) + + predicted = predict_cast_state_dict_size(sd, torch.bfloat16, keep_fp8=True, **kwargs) + cast_state_dict(sd, torch.bfloat16, keep_fp8=True, **kwargs) + actual = sum(t.nelement() * t.element_size() for t in sd.values()) + assert predicted == actual + + # Only `lin.weight` stays 1 byte/element; the rest lands at 2, so the old sum under-counted. + assert naive < actual + + def test_all_dequantized_when_keep_fp8_is_off(self): + q, _ = _fp8_weight(32, 16) + sd = {"lin.weight": q} + assert predict_cast_state_dict_size(sd, torch.bfloat16, keep_fp8=False) == q.nelement() * 2 + + +class TestDeviceSupport: + def test_probe_decides_rather_than_the_capability_number(self): + """ROCm reports the gfx arch from `get_device_capability`, so RDNA3 (gfx1100 -> (11, 0)) + passes a `>= (8, 9)` test and then raises on every forward. The probe is what decides.""" + reset_fp8_matmul_support_cache() + with mock.patch("torch.cuda.is_available", return_value=True): + with mock.patch("torch.cuda.get_device_capability", return_value=(11, 0)): + with mock.patch( + "invokeai.backend.quantization.fp8_scaled._probe_fp8_matmul", return_value=False + ) as probe: + assert device_supports_fp8_matmul(torch.device("cuda", 0)) is False + assert probe.called + reset_fp8_matmul_support_cache() + + def test_probe_result_is_cached_per_device(self): + reset_fp8_matmul_support_cache() + with mock.patch("torch.cuda.is_available", return_value=True): + with mock.patch("torch.cuda.get_device_capability", return_value=(8, 9)): + with mock.patch( + "invokeai.backend.quantization.fp8_scaled._probe_fp8_matmul", return_value=True + ) as probe: + assert device_supports_fp8_matmul(torch.device("cuda", 0)) is True + assert device_supports_fp8_matmul(torch.device("cuda", 0)) is True + assert probe.call_count == 1 + reset_fp8_matmul_support_cache() + + def test_non_cuda_never_probes(self): + reset_fp8_matmul_support_cache() + with mock.patch("invokeai.backend.quantization.fp8_scaled._probe_fp8_matmul") as probe: + assert device_supports_fp8_matmul(torch.device("cpu")) is False + assert not probe.called + + +class TestDequantize: + @pytest.mark.parametrize("per_channel", [False, True]) + def test_roundtrip_close_to_original(self, per_channel: bool): + w = torch.randn(64, 32, dtype=torch.bfloat16) * 0.02 + scale_src = w.abs().amax(dim=1) if per_channel else w.abs().amax() + scale = (scale_src / torch.finfo(FP8_DTYPE).max).float().clamp(min=1e-12) + q = (w / (scale.reshape(-1, 1) if per_channel else scale)).to(FP8_DTYPE) + + out = dequantize_weight(q, scale, torch.bfloat16) + assert out.dtype == torch.bfloat16 + assert ((out.float() - w.float()).norm() / w.float().norm()).item() < 0.05 + + def test_missing_scale_is_a_plain_cast(self): + """fp8_storage layerwise casting produces scale-free fp8 weights.""" + q = torch.randn(8, 8, dtype=torch.bfloat16).to(FP8_DTYPE) + assert torch.equal(dequantize_weight(q, None, torch.bfloat16), q.to(torch.bfloat16)) + + def test_state_dict_dequantization_matches_helper(self): + q, scale = _fp8_weight(64, 32) + sd = {"lin.weight": q} + layers = {"lin": Fp8ScaledLayer(weight_scale=scale)} + dequantize_fp8_scaled(sd, layers) + assert torch.equal(sd["lin.weight"], dequantize_weight(q, scale, torch.bfloat16)) + + +class TestAttach: + def test_registers_non_persistent_buffers(self): + lin = torch.nn.Linear(16, 32, bias=False) + q, scale = _fp8_weight(32, 16) + lin.weight = torch.nn.Parameter(q, requires_grad=False) + model = torch.nn.Sequential(lin) + + assert attach_fp8_scales(model, {"0": Fp8ScaledLayer(weight_scale=scale, full_precision_matmul=True)}) == 1 + assert torch.equal(model[0].weight_scale, scale) + assert model[0]._fp8_full_precision_matmul is True + # Re-saving a model whose scales landed in state_dict() would double-scale on reload. + assert "0.weight_scale" not in model.state_dict() + + def test_skips_non_fp8_modules(self): + model = torch.nn.Sequential(torch.nn.Linear(16, 32)) + assert attach_fp8_scales(model, {"0": Fp8ScaledLayer(weight_scale=torch.tensor(1.0))}) == 0 + + +@cuda_fp8 +class TestScaledMm: + @pytest.mark.parametrize("per_channel", [False, True]) + @pytest.mark.parametrize("tokens", [64, 100]) # 100 is deliberately not a multiple of 16 + def test_matches_reference_linear(self, per_channel: bool, tokens: int): + dev = torch.device("cuda") + q, scale = _fp8_weight(256, 128, per_channel) + q, scale = q.to(dev), scale.to(dev) + w_ref = dequantize_weight(q, scale, torch.bfloat16) + x = torch.randn(tokens, 128, device=dev, dtype=torch.bfloat16) + + got = scaled_mm_linear(x, q, scale) + expected = torch.nn.functional.linear(x, w_ref) + + assert got.shape == expected.shape + rel = ((got.float() - expected.float()).norm() / expected.float().norm()).item() + assert rel < 0.08, f"relative error {rel:.4f} too high" + + def test_preserves_leading_dims_and_bias(self): + dev = torch.device("cuda") + q, scale = _fp8_weight(64, 32) + q, scale = q.to(dev), scale.to(dev) + bias = torch.randn(64, device=dev, dtype=torch.bfloat16) + x = torch.randn(2, 48, 32, device=dev, dtype=torch.bfloat16) + + got = scaled_mm_linear(x, q, scale, bias) + expected = torch.nn.functional.linear(x, dequantize_weight(q, scale, torch.bfloat16), bias) + assert got.shape == (2, 48, 64) + assert ((got.float() - expected.float()).norm() / expected.float().norm()).item() < 0.08 + + def test_static_input_scale_path(self): + dev = torch.device("cuda") + q, scale = _fp8_weight(64, 32) + q, scale = q.to(dev), scale.to(dev) + x = torch.randn(32, 32, device=dev, dtype=torch.bfloat16) + static = (x.abs().amax() / torch.finfo(FP8_DTYPE).max).float() + + got = scaled_mm_linear(x, q, scale, input_scale=static) + expected = torch.nn.functional.linear(x, dequantize_weight(q, scale, torch.bfloat16)) + assert ((got.float() - expected.float()).norm() / expected.float().norm()).item() < 0.08 + + +class TestCustomLinearIntegration: + """The fp8 matmul must be opt-in and must degrade to the dequantized path, never raise.""" + + @pytest.fixture(autouse=True) + def _restore_matmul_override(self): + """Clear the override after each test, rather than leaving it pinned to ``False``. + + These tests reset with ``set_fp8_matmul_enabled(False)``, which is not the neutral state -- + it is an explicit "off" that outlives the class and silently overrides the config for + anything that runs later in the same process. + """ + yield + set_fp8_matmul_enabled(None) + + def _module(self, device: torch.device, in_f=64, out_f=128, device_autocasting: bool = False): + from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import ( + apply_custom_layers_to_model, + ) + + lin = torch.nn.Linear(in_f, out_f, bias=False) + q, scale = _fp8_weight(out_f, in_f) + lin.weight = torch.nn.Parameter(q.to(device), requires_grad=False) + lin.register_buffer("weight_scale", scale.to(device), persistent=False) + model = torch.nn.Sequential(lin).to(device) + apply_custom_layers_to_model(model, device_autocasting_enabled=device_autocasting) + return model + + @cuda_fp8 + @pytest.mark.parametrize("device_autocasting", [False, True]) + def test_fp8_path_runs_regardless_of_device_autocasting(self, device_autocasting: bool): + """`apply_custom_layers_to_model` leaves autocasting off for fully-resident models. + + A fp8 check that only lives in `_autocast_forward` is therefore skipped in the common case: + `forward` falls through to the dtype-mismatch branch and silently dequantizes instead. This + regression cost the entire speedup while every other test stayed green. + """ + dev = torch.device("cuda") + model = self._module(dev, device_autocasting=device_autocasting) + x = torch.randn(32, 64, device=dev, dtype=torch.bfloat16) + custom = model[0] + + calls = [] + original = custom._maybe_fp8_forward + custom._maybe_fp8_forward = lambda inp: (calls.append(1), original(inp))[1] + + set_fp8_matmul_enabled(True) + try: + out = model(x) + finally: + set_fp8_matmul_enabled(False) + custom._maybe_fp8_forward = original + + assert calls, "the fp8 branch was never consulted" + assert out.shape == (32, 128) + # And it must actually have taken the fp8 path, not just been asked. + w_ref = dequantize_weight(custom.weight, custom.weight_scale, torch.bfloat16) + assert not torch.equal(out, torch.nn.functional.linear(x, w_ref)), ( + "output is bit-identical to the dequantized path, so _scaled_mm did not run" + ) + + def test_disabled_by_default_uses_scaled_dequant(self): + """With the matmul off, the weight must still be dequantized *with* its scale.""" + set_fp8_matmul_enabled(False) + dev = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = self._module(dev) + x = torch.randn(32, 64, device=dev, dtype=torch.bfloat16) + + got = model(x) + w_ref = dequantize_weight(model[0].weight, model[0].weight_scale, torch.bfloat16) + expected = torch.nn.functional.linear(x, w_ref) + assert torch.equal(got, expected) + + @cuda_fp8 + def test_full_precision_flag_forces_fallback(self): + set_fp8_matmul_enabled(True) + try: + dev = torch.device("cuda") + model = self._module(dev) + model[0]._fp8_full_precision_matmul = True + x = torch.randn(32, 64, device=dev, dtype=torch.bfloat16) + + got = model(x) + w_ref = dequantize_weight(model[0].weight, model[0].weight_scale, torch.bfloat16) + assert torch.equal(got, torch.nn.functional.linear(x, w_ref)) + finally: + set_fp8_matmul_enabled(False) + + @cuda_fp8 + def test_sidecar_patched_fp8_layer_uses_the_fp8_matmul(self): + """LoRAs on fp8 weights are force-routed to sidecar patching (layer_patcher), and the + sidecar wrapper calls `_autocast_forward`. The fp8 branch must survive that route, and the + LoRA residual must still be added on top.""" + from invokeai.backend.patches.layers.lora_layer import LoRALayer + + dev = torch.device("cuda") + model = self._module(dev) + custom = model[0] + rank = 8 + lora = LoRALayer( + up=torch.randn(custom.out_features, rank, device=dev, dtype=torch.bfloat16) * 0.05, + mid=None, + down=torch.randn(rank, custom.in_features, device=dev, dtype=torch.bfloat16) * 0.05, + alpha=float(rank), + bias=None, + ) + custom.add_patch(lora, 1.0) + x = torch.randn(32, 64, device=dev, dtype=torch.bfloat16) + + set_fp8_matmul_enabled(True) + try: + patched = model(x) + finally: + set_fp8_matmul_enabled(False) + unpatched_ref = torch.nn.functional.linear( + x, dequantize_weight(custom.weight, custom.weight_scale, torch.bfloat16) + ) + + assert patched.shape == (32, 128) + # The patch must actually change the output... + assert not torch.allclose(patched, unpatched_ref, atol=1e-2) + # ...and the base contribution must still be roughly the fp8 linear, not garbage. + residual = (lora.get_weight(1.0) * lora.scale()).to(torch.bfloat16) + expected = unpatched_ref + torch.nn.functional.linear(x, residual) + rel = ((patched.float() - expected.float()).norm() / expected.float().norm()).item() + assert rel < 0.1, f"sidecar-patched fp8 output diverges by {rel:.4f}" + + @cuda_fp8 + def test_unaligned_features_fall_back(self): + set_fp8_matmul_enabled(True) + try: + dev = torch.device("cuda") + model = self._module(dev, in_f=60, out_f=120) # not multiples of 16 + x = torch.randn(16, 60, device=dev, dtype=torch.bfloat16) + assert model(x).shape == (16, 120) # must not raise + finally: + set_fp8_matmul_enabled(False) + + @cuda_fp8 + def test_enabled_path_agrees_with_fallback(self): + dev = torch.device("cuda") + model = self._module(dev) + x = torch.randn(48, 64, device=dev, dtype=torch.bfloat16) + + set_fp8_matmul_enabled(False) + baseline = model(x) + set_fp8_matmul_enabled(True) + try: + fp8 = model(x) + finally: + set_fp8_matmul_enabled(False) + + rel = ((fp8.float() - baseline.float()).norm() / baseline.float().norm()).item() + assert rel < 0.1, f"fp8 matmul diverges from the dequantized path by {rel:.4f}" + + +class TestScaleSpellingHelpers: + """Both spellings of the weight scale must be handled everywhere. + + Reading only `.weight_scale` is the mistake that keeps recurring in the per-loader dequant + helpers. Depending on what the loader strips afterwards it either deletes a `.scale_weight` + without applying it — leaving the weight off by `1/weight_scale`, silently — or leaves the key + behind for `load_state_dict(..., strict=True)` to reject. + """ + + @pytest.mark.parametrize("spelling", [".weight_scale", ".scale_weight"]) + def test_pairs_are_found_in_either_spelling(self, spelling: str) -> None: + sd = {"blk.weight": torch.ones(2, 2), f"blk{spelling}": torch.tensor(0.5)} + + assert list(iter_weight_scale_pairs(sd)) == [("blk.weight", f"blk{spelling}")] + + def test_a_scale_without_its_weight_is_not_paired(self) -> None: + # Pairing it would invent a weight key the checkpoint never had. + sd = {"other.weight": torch.ones(1), "blk.weight_scale": torch.tensor(0.5)} + + assert list(iter_weight_scale_pairs(sd)) == [] + + def test_non_string_keys_are_ignored(self) -> None: + # `.pt`/`.ckpt` sources can carry int keys; `endswith` would raise on them. + assert list(iter_weight_scale_pairs({0: torch.ones(1)})) == [] + + @pytest.mark.parametrize( + "key", + [".weight_scale", ".scale_weight", ".input_scale", ".scale_input"], + ) + def test_metadata_keys_are_recognized_in_either_spelling(self, key: str) -> None: + assert is_scale_metadata_key(f"blk{key}") + + @pytest.mark.parametrize("key", ["blk.comfy_quant", "scaled_fp8"]) + def test_marker_keys_are_recognized(self, key: str) -> None: + assert is_scale_metadata_key(key) + + @pytest.mark.parametrize("key", ["blk.weight", "blk.bias", "norm.scale", 0]) + def test_model_tensors_are_not_mistaken_for_metadata(self, key: object) -> None: + # `norm.scale` is a real learned parameter in several architectures - stripping it would + # delete weights, and it is why this cannot just match "scale" anywhere in the key. + assert not is_scale_metadata_key(key) + + +def _fp8(rows: int = 4, cols: int = 4, value: float = 1.0) -> torch.Tensor: + return (torch.ones(rows, cols) * value).to(FP8_DTYPE) + + +class _Linear(torch.nn.Module): + def __init__(self, out_features: int = 4, in_features: int = 4) -> None: + super().__init__() + self.lin = torch.nn.Linear(in_features, out_features, bias=False) + + +class TestE5m2ScaleRecovery: + """`float8_e5m2` may not stay quantized, but it must not lose its scale on the way to bf16.""" + + def test_the_scale_is_recovered_and_folded(self) -> None: + # Gating extraction on e4m3fn alone popped the scale key and then dropped it, so + # `cast_state_dict` did a plain `.to(bf16)` and the weight came out off by 1/weight_scale. + sd = {"lin.weight": (torch.ones(4, 4) * 2).to(torch.float8_e5m2), "lin.weight_scale": torch.tensor(0.25)} + + layers = extract_fp8_scaled_layers(sd) + assert "lin" in layers + + dequantize_fp8_scaled(sd, layers, torch.bfloat16) + assert torch.allclose(sd["lin.weight"].float(), torch.full((4, 4), 0.5)) + + def test_it_is_never_left_quantized(self) -> None: + """`scaled_mm_linear` cannot take e5m2 as the weight operand, so it must be widened.""" + sd = {"lin.weight": (torch.ones(4, 4) * 2).to(torch.float8_e5m2), "lin.weight_scale": torch.tensor(0.25)} + + kept = split_fp8_scaled_layers(sd, extract_fp8_scaled_layers(sd), torch.bfloat16, model=_Linear()) + + assert kept == {} + assert sd["lin.weight"].dtype is torch.bfloat16 + assert torch.allclose(sd["lin.weight"].float(), torch.full((4, 4), 0.5)) + + +class TestBlockWiseScale: + """A 2-D ``weight_scale`` has one entry per *block* of weight elements, not per row.""" + + def test_the_layout_survives_extraction(self) -> None: + # Flattening it destroys the block geometry, and the multiply then fails on shape. + sd = {"b.weight": _fp8(64, 128), "b.weight_scale": torch.full((64, 2), 0.5)} + + layers = extract_fp8_scaled_layers(sd) + + assert tuple(layers["b"].weight_scale.shape) == (64, 2) + + def test_it_is_expanded_rather_than_raising(self) -> None: + # Before: `RuntimeError: The size of tensor a (64) must match the size of tensor b (128)`, + # i.e. with fp8_compute off - the default - the model failed to load at all. + sd = {"b.weight": _fp8(64, 128), "b.weight_scale": torch.full((64, 2), 0.5)} + + dequantize_fp8_scaled(sd, extract_fp8_scaled_layers(sd), torch.bfloat16) + + assert sd["b.weight"].shape == (64, 128) + assert torch.allclose(sd["b.weight"].float(), torch.full((64, 128), 0.5)) + + def test_it_is_never_left_quantized(self) -> None: + """`scaled_mm_linear` can apply a per-tensor or per-row scale and nothing else. + + Left quantized, the mismatch surfaces mid-generation inside the kernel instead. + Dequantizing here - before `predict_cast_state_dict_size` runs - also keeps the RAM + reservation honest. + """ + sd = {"lin.weight": _fp8(4, 4), "lin.weight_scale": torch.full((4, 2), 0.5)} + + kept = split_fp8_scaled_layers(sd, extract_fp8_scaled_layers(sd), torch.bfloat16, model=_Linear()) + + assert kept == {} + assert sd["lin.weight"].dtype is torch.bfloat16 + + def test_a_per_row_scale_still_stays_quantized(self) -> None: + sd = {"lin.weight": _fp8(4, 4), "lin.weight_scale": torch.full((4,), 0.5)} + + kept = split_fp8_scaled_layers(sd, extract_fp8_scaled_layers(sd), torch.bfloat16, model=_Linear()) + + assert set(kept) == {"lin"} + assert sd["lin.weight"].dtype is FP8_DTYPE + + +class TestSidechannelDetachReattach: + """Key converters rename ``.weight``; the side channel has to be carried across separately.""" + + def test_entries_follow_their_module_to_the_new_path(self) -> None: + sd = { + "old.linear.weight": _fp8(), + "old.linear.weight_scale": torch.tensor(0.5), + "old.linear.input_scale": torch.tensor(0.25), + } + + detached = detach_layer_sidechannel(sd) + assert list(sd) == ["old.linear.weight"], "scales must be out of the converter's way" + + converted = {"new.linear.weight": sd["old.linear.weight"]} + orphaned = reattach_layer_sidechannel(converted, detached, {"old.linear": "new.linear"}) + + assert orphaned == [] + layers = extract_fp8_scaled_layers(converted) + assert set(layers) == {"new.linear"} + assert layers["new.linear"].input_scale is not None + + def test_a_module_the_converter_drops_is_reported_not_swallowed(self) -> None: + """A silently dropped scale is exactly the failure this pair exists to prevent.""" + sd = {"gone.weight": _fp8(), "gone.weight_scale": torch.tensor(0.5)} + + detached = detach_layer_sidechannel(sd) + orphaned = reattach_layer_sidechannel({}, detached, {}) + + assert orphaned == ["gone"] + + def test_both_scale_spellings_are_detached(self) -> None: + sd = { + "a.weight": _fp8(), + "a.scale_weight": torch.tensor(0.5), + "b.weight": _fp8(), + "b.weight_scale": torch.tensor(0.5), + } + + detached = detach_layer_sidechannel(sd) + + assert set(detached) == {"a", "b"} + + +class TestStripLayerPathPrefix: + """`_quantization_metadata` is read from the file header, so its names keep the prefix.""" + + def test_the_checkpoint_prefix_is_removed(self) -> None: + hints = {"model.diffusion_model.blocks.0.attn.wq": {"full_precision_matrix_mult": True}} + + assert strip_layer_path_prefix(hints) == {"blocks.0.attn.wq": {"full_precision_matrix_mult": True}} + + def test_unprefixed_names_are_passed_through_not_dropped(self) -> None: + """Running the names through a prefix *filter* truncated a partially-prefixed header, and + the strict-zip that read the result back aborted the whole load with a ValueError.""" + hints = {"net.blocks.0.attn.q_proj": {}, "final_layer.linear": {}} + + assert set(strip_layer_path_prefix(hints)) == {"blocks.0.attn.q_proj", "final_layer.linear"} + + +class TestNonFloatTensorsAreNotCast: + """Integer payloads are not weights; casting them to the compute dtype corrupts them.""" + + def test_cast_state_dict_leaves_them_alone(self) -> None: + sd = {"w.weight": torch.ones(2, 2), "ids": torch.arange(4, dtype=torch.int64)} + + cast_state_dict(sd, torch.bfloat16, keep_fp8=False) + + assert sd["w.weight"].dtype is torch.bfloat16 + assert sd["ids"].dtype is torch.int64 + + def test_the_size_prediction_agrees(self) -> None: + sd = {"ids": torch.arange(4, dtype=torch.int64)} + + assert predict_cast_state_dict_size(sd, torch.bfloat16, keep_fp8=False) == 4 * 8 + + +@contextlib.contextmanager +def _probe_on_cpu(): + """Let `_probe_fp8_matmul` run its allocations on CPU so the caching policy can be tested. + + The probe allocates real tensors on the device *before* it reaches `torch._scaled_mm`. On a + CPU-only runner that allocation raises first, the mocked matmul never runs, and the test + measures the wrong thing — it passed on a CUDA box and failed on CI for exactly this reason. + float8 tensors are allocatable on CPU, so dropping the device argument is enough. + """ + real_zeros, real_ones = torch.zeros, torch.ones + + def on_cpu(real): + def alloc(*args, **kwargs): + kwargs.pop("device", None) + return real(*args, **kwargs) + + return alloc + + with ( + mock.patch("torch.cuda.is_available", return_value=True), + mock.patch("torch.cuda.get_device_capability", return_value=(8, 9)), + mock.patch("torch.cuda.current_device", return_value=0), + mock.patch("torch.zeros", on_cpu(real_zeros)), + mock.patch("torch.ones", on_cpu(real_ones)), + ): + yield + + +class TestProbeFailureCaching: + """The probe runs during a model load, i.e. under real VRAM pressure.""" + + def test_an_allocation_failure_is_not_cached(self) -> None: + reset_fp8_matmul_support_cache() + device = torch.device("cuda", 0) + with _probe_on_cpu(): + with mock.patch("torch._scaled_mm", side_effect=torch.OutOfMemoryError("transient")): + assert device_supports_fp8_matmul(device) is False + # A momentary OOM must not disable fp8 for the rest of the process. + with mock.patch("torch._scaled_mm", return_value=torch.zeros(1)): + assert device_supports_fp8_matmul(device) is True + reset_fp8_matmul_support_cache() + + def test_a_genuine_unsupported_op_is_cached(self) -> None: + reset_fp8_matmul_support_cache() + device = torch.device("cuda", 0) + with _probe_on_cpu(): + with mock.patch("torch._scaled_mm", side_effect=RuntimeError("not supported on this device")): + assert device_supports_fp8_matmul(device) is False + # No second probe: the answer cannot change at runtime. + with mock.patch("torch._scaled_mm", side_effect=AssertionError("must not be probed again")): + assert device_supports_fp8_matmul(device) is False + reset_fp8_matmul_support_cache() + + +class TestMxfp8IsRefused: + """MXFP8 block scales are refused rather than guessed at. + + Established against a real pair of checkpoints: the MXFP8 and the scaled-fp8 build of + `krea2TurboOfficialComfy` share all 174 bf16 tensors bit-for-bit, so the scaled build is an + exact reference. Decoding the uint8 exponents as `2**(v-127)` and expanding them 32-wide + reaches a correlation of only 0.60 against it and generates a pure-noise image; the measured + per-block scale has no monotonic relation to the byte (112 and 116 give the same true scale), + which points at a swizzled scale layout. + + Refusing matters because the block-wise expansion is what makes such a file *loadable*: without + a guard it produces garbage silently, which is strictly worse than the shape error it used to + raise. + """ + + def test_a_uint8_block_scale_is_rejected_with_an_actionable_message(self) -> None: + sd = { + "lin.weight": torch.full((4, 64), 2.0).to(FP8_DTYPE), + "lin.weight_scale": torch.full((4, 2), 125, dtype=torch.uint8), + } + + with pytest.raises(NotImplementedError) as excinfo: + extract_fp8_scaled_layers(sd) + + message = str(excinfo.value) + assert "lin" in message, "the failing layer must be named" + assert "MXFP8" in message + assert "noise" in message, "say what happens if it were loaded anyway" + + def test_float_scales_are_unaffected(self) -> None: + """The guard keys off the dtype, so ordinary scaled-fp8 checkpoints must still load.""" + sd = { + "lin.weight": torch.full((4, 64), 2.0).to(FP8_DTYPE), + "lin.weight_scale": torch.full((4, 2), 0.25), + } + + layers = extract_fp8_scaled_layers(sd) + + assert tuple(layers["lin"].weight_scale.shape) == (4, 2) + + +class TestNonScalarInputScale: + """A multi-element ``input_scale`` is dropped, not raised on. + + `scaled_mm_linear` scales activations by one value, so there is nothing to do with a + per-channel or per-block activation scale -- but reshaping it to a scalar unconditionally + turned such a checkpoint into a shape error at load time. Every other malformed side-channel in + this module is either skipped or refused with an actionable message; this one was neither. + """ + + def _sd(self, input_scale: torch.Tensor) -> dict[str, torch.Tensor]: + return { + "blocks.0.attn.weight": torch.zeros(16, 16, dtype=torch.float32).to(FP8_DTYPE), + "blocks.0.attn.weight_scale": torch.tensor(2.0), + "blocks.0.attn.input_scale": input_scale, + } + + def test_a_per_channel_input_scale_does_not_abort_the_load(self) -> None: + layers = extract_fp8_scaled_layers(self._sd(torch.full((16,), 0.5))) + + assert set(layers) == {"blocks.0.attn"} + # Dropped, so the forward falls back to the dynamic amax path -- which is correct. + assert layers["blocks.0.attn"].input_scale is None + # The weight scale is unaffected by the unusable activation scale. + assert layers["blocks.0.attn"].weight_scale is not None + + def test_a_scalar_input_scale_is_still_kept(self) -> None: + layers = extract_fp8_scaled_layers(self._sd(torch.tensor(0.5))) + + assert layers["blocks.0.attn"].input_scale is not None + + +class TestMatmulUsableScale: + """`is_matmul_usable_scale` decides what may reach `_scaled_mm` un-dequantized.""" + + def test_a_per_tensor_scale_is_usable(self) -> None: + assert is_matmul_usable_scale(torch.zeros(32, 16), torch.tensor(2.0)) is True + + def test_a_per_output_channel_scale_is_usable(self) -> None: + assert is_matmul_usable_scale(torch.zeros(32, 16), torch.full((32,), 2.0)) is True + + def test_a_one_dimensional_scale_of_the_wrong_length_is_not(self) -> None: + """The branch that keeps a mislabelled scale out of the kernel. + + A 1-D scale whose length is not the row count is neither per-tensor nor + per-output-channel. Letting it through raises inside `_scaled_mm` mid-generation instead of + dequantizing the layer up front, which is the whole point of asking. + """ + assert is_matmul_usable_scale(torch.zeros(32, 16), torch.full((16,), 2.0)) is False + + def test_a_block_wise_scale_is_not(self) -> None: + assert is_matmul_usable_scale(torch.zeros(32, 16), torch.full((4, 2), 2.0)) is False + + +class TestReattachToNonWeightDestination: + """The reattach guard must not assume every module stores its parameter as ``weight``. + + A producer that quantizes norms -- the "quantizes everything" class this module documents -- + writes ``.scale``. Testing for ``.weight`` rejects such a destination for the + wrong reason and drops a scale it could have placed, with only a log line to show for it. + """ + + def test_a_destination_whose_tensor_is_not_named_weight_still_receives_its_scale(self) -> None: + sd = {"blocks.0.norm_q.scale": torch.zeros(16, dtype=torch.float32).to(FP8_DTYPE)} + detached = {"blocks.0.qnorm": [(".weight_scale", torch.tensor(2.0))]} + + orphaned = reattach_layer_sidechannel(sd, detached, {"blocks.0.qnorm": "blocks.0.norm_q"}) + + assert orphaned == [] + assert "blocks.0.norm_q.weight_scale" in sd + + def test_a_destination_absent_from_the_state_dict_is_still_reported(self) -> None: + sd = {"blocks.0.attn.weight": torch.zeros(4, 4)} + detached = {"blocks.0.dropped": [(".weight_scale", torch.tensor(2.0))]} + + orphaned = reattach_layer_sidechannel(sd, detached, {}) + + assert orphaned == ["blocks.0.dropped"] + assert not [k for k in sd if k.endswith(".weight_scale")] + + +class TestProbeTransientFailures: + """Only a definitively unsupported device may be cached; everything else re-probes. + + Listing the transient wordings is the wrong way round: an OOM and a cuBLAS workspace failure + are two of the ways a loaded machine can fail this call, and any wording not on the list would + be cached as permanent -- reproducing the failure the OOM branch was written to avoid. + """ + + def test_an_unrecognized_runtime_error_is_not_cached(self) -> None: + reset_fp8_matmul_support_cache() + device = torch.device("cuda", 0) + with _probe_on_cpu(): + with mock.patch("torch._scaled_mm", side_effect=RuntimeError("CUDA driver reset")): + assert device_supports_fp8_matmul(device) is False + # Inconclusive, so the next load re-probes rather than the process losing fp8. + with mock.patch("torch._scaled_mm", return_value=torch.zeros(1)): + assert device_supports_fp8_matmul(device) is True + reset_fp8_matmul_support_cache() + + def test_a_capability_error_is_cached(self) -> None: + reset_fp8_matmul_support_cache() + device = torch.device("cuda", 0) + message = "torch._scaled_mm is only supported on CUDA devices with compute capability >= 8.9" + with _probe_on_cpu(): + with mock.patch("torch._scaled_mm", side_effect=RuntimeError(message)): + assert device_supports_fp8_matmul(device) is False + with mock.patch("torch._scaled_mm", side_effect=AssertionError("must not be probed again")): + assert device_supports_fp8_matmul(device) is False + reset_fp8_matmul_support_cache() + + +class TestExpandWeightScaleAxis: + """A per-output-channel scale must scale rows, not columns. + + `(out, in) * (out,)` broadcasts on the *last* axis, so a bare multiply scales input channels -- + wrong on a square weight, a shape error on any other. Both legacy folds used to do exactly that. + """ + + def test_a_per_channel_scale_lines_up_with_the_rows(self) -> None: + weight = torch.ones(3, 2) + scale = torch.tensor([1.0, 2.0, 3.0]) + + result = weight * expand_weight_scale(weight, scale) + + assert torch.equal(result, torch.tensor([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]])) + + def test_a_non_square_weight_no_longer_raises(self) -> None: + weight = torch.ones(4, 2) + + # Without the expansion this is a broadcast error, not merely a wrong number. + assert (weight * expand_weight_scale(weight, torch.arange(4, dtype=torch.float32))).shape == (4, 2) + + +class TestPredictionIsSplitAware: + """`make_room` runs before `split_fp8_scaled_layers`, so the prediction has to account for what + the split widens — not just for what the cast keeps. + + `can_stay_quantized` knows nothing about the scale *layout*. A block-wise-scaled 2-D + `nn.Linear.weight` passes it, gets predicted at 1 byte/element, and is then dequantized to + `dtype` by the split at 2 — the reservation is half the truth, and on a real checkpoint using + that layout the shortfall is the whole quantized-Linear set. + """ + + def _model(self): + model = torch.nn.Sequential() + model.add_module("lin", torch.nn.Linear(16, 32)) + return model + + @pytest.mark.parametrize( + "label, scale, stays_fp8", + [ + ("per-tensor", torch.tensor(2.0), True), + ("per-output-channel", torch.full((32,), 2.0), True), + ("block-wise", torch.full((8, 2), 2.0), False), + ], + ) + def test_prediction_matches_the_bytes_left_after_split_and_cast(self, label, scale, stays_fp8): + model = self._model() + sd = {"lin.weight": torch.zeros(32, 16).to(FP8_DTYPE), "lin.weight_scale": scale} + layers = extract_fp8_scaled_layers(sd) + + predicted = predict_cast_state_dict_size(sd, torch.bfloat16, keep_fp8=True, model=model, scaled_layers=layers) + remaining = split_fp8_scaled_layers(sd, layers, torch.bfloat16, model=model) + cast_state_dict(sd, torch.bfloat16, keep_fp8=True, model=model) + actual = sum(t.nelement() * t.element_size() for t in sd.values()) + + assert predicted == actual, f"{label}: reserved {predicted}, needed {actual}" + assert bool(remaining) is stays_fp8 + assert (sd["lin.weight"].dtype is FP8_DTYPE) is stays_fp8 + + def test_without_the_mapping_the_block_wise_layer_is_under_counted(self): + """The regression itself: omitting `scaled_layers` reverts to the weaker predicate.""" + model = self._model() + sd = {"lin.weight": torch.zeros(32, 16).to(FP8_DTYPE), "lin.weight_scale": torch.full((8, 2), 2.0)} + layers = extract_fp8_scaled_layers(sd) + + blind = predict_cast_state_dict_size(sd, torch.bfloat16, keep_fp8=True, model=model) + aware = predict_cast_state_dict_size(sd, torch.bfloat16, keep_fp8=True, model=model, scaled_layers=layers) + assert aware == 2 * blind + + def test_raw_fp8_without_any_scale_is_unaffected(self): + """A weight with no `weight_scale` has no layout to reject; it stays fp8 either way.""" + model = self._model() + sd = {"lin.weight": torch.zeros(32, 16).to(FP8_DTYPE)} + with_map = predict_cast_state_dict_size(sd, torch.bfloat16, keep_fp8=True, model=model, scaled_layers={}) + without = predict_cast_state_dict_size(sd, torch.bfloat16, keep_fp8=True, model=model) + assert with_map == without == 32 * 16 + + +class TestQwen3VlStyleSplit: + """Pins the helper contract the Qwen3-VL encoder switched to. + + Its old per-key `if dtype is not FP8_DTYPE` loop kept *every* fp8 tensor quantized. The three + things that got wrong are all visible without any `transformers` dependency, against a stub + module tree. + """ + + def _model(self): + model = torch.nn.Module() + model.add_module("norm", torch.nn.LayerNorm(32)) + model.add_module("lin", torch.nn.Linear(16, 32)) + model.add_module("plain", torch.nn.Linear(16, 32)) + return model + + def test_scaled_norm_is_folded_and_block_wise_linear_is_dropped(self): + model = self._model() + norm_scale = torch.tensor(4.0) + sd = { + # A 1-D fp8 norm carrying a scale: not a matmul weight, so it must be folded, not kept. + "norm.weight": torch.ones(32).to(FP8_DTYPE), + "norm.weight_scale": norm_scale, + # A Linear whose scale layout `scaled_mm` cannot apply: folded too. + "lin.weight": torch.ones(32, 16).to(FP8_DTYPE), + "lin.weight_scale": torch.full((8, 2), 2.0), + # A Linear the matmul can take: survives, and is what `attach_fp8_scales` gets. + "plain.weight": torch.ones(32, 16).to(FP8_DTYPE), + "plain.weight_scale": torch.tensor(3.0), + } + layers = extract_fp8_scaled_layers(sd) + assert set(layers) == {"norm", "lin", "plain"} + + remaining = split_fp8_scaled_layers(sd, layers, torch.bfloat16, model=model) + cast_state_dict(sd, torch.bfloat16, keep_fp8=True, model=model) + + assert set(remaining) == {"plain"}, "only the matmul-usable Linear may keep its scale" + # The norm was dequantized *with* its scale applied — the old loop left it as raw fp8 codes. + assert sd["norm.weight"].dtype is torch.bfloat16 + assert torch.allclose(sd["norm.weight"], torch.full((32,), 4.0, dtype=torch.bfloat16)) + assert sd["lin.weight"].dtype is torch.bfloat16 + assert sd["plain.weight"].dtype is FP8_DTYPE + + def test_e5m2_scaled_weight_keeps_its_scale_on_the_way_down(self): + """Only e4m3fn may stay quantized, so an e5m2 weight must be folded rather than cast bare.""" + model = self._model() + sd = { + "plain.weight": torch.ones(32, 16).to(torch.float8_e5m2), + "plain.weight_scale": torch.tensor(3.0), + } + layers = extract_fp8_scaled_layers(sd) + remaining = split_fp8_scaled_layers(sd, layers, torch.bfloat16, model=model) + + assert remaining == {} + assert torch.allclose(sd["plain.weight"], torch.full((32, 16), 3.0, dtype=torch.bfloat16)) + + +class TestMalformedWeightScale: + def test_a_1d_scale_of_the_wrong_length_is_reported_not_broadcast(self): + """Neither per-tensor nor per-output-channel. Left to torch, the multiply raises "size of + tensor a (32) must match tensor b (7)" from inside the fold, naming neither layer nor file.""" + with pytest.raises(ValueError, match="neither per-tensor nor per-output-channel"): + expand_weight_scale(torch.ones(32, 16), torch.full((7,), 2.0))