From 8a9358e2fd15b0af4e07ca0d5f2c8286bdccf87a Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Tue, 1 Sep 2026 03:31:38 +0200 Subject: [PATCH 1/2] feat(quantization): share the int8_convrot scheme and teach Krea-2 to read it `int8_tensorwise` is a ComfyUI-wide scheme, not one architecture's format, but the implementation lived in `backend/minimax_h3/`. Krea-2's loaders could not reach it, so `*_int8_convrot.safetensors` loaded as confident nonsense: `_dequantize_scaled_fp8` keys off `.weight_scale` alone, scaled the int8 weights and never un-rotated them. No error, no NaN -- correlation against a correct decode is 0.06. Moves `int8_convrot.py` to `backend/quantization/` beside the other schemes (gguf, sdnq, bnb) and wires both Krea-2 loaders to it. The weights stay int8-resident: each quantized Linear becomes an `Int8ConvrotLinear` that dequantizes and derotates per forward, as MiniMax H3 already does. That is 12.3 GB instead of 24 GB, on every platform and without the fp8-storage opt-in -- which is off by default and unavailable outside CUDA/XPU. Ordering is load-bearing three times over, and none of it fails loudly: - before the fp8 fold, which would scale an int8 weight without un-rotating it; - before the native->diffusers key conversion, which renames `.attn.wq.weight` by substring, carrying `.weight_scale` along but orphaning `.comfy_quant`; - before the encoder's fp8 detection, which answers yes to any `.weight_scale` and would keep an int8 encoder "fp8-resident" over weights that were never fp8. Each is pinned by a test that shows the damage rather than asserting the order. Three things the real checkpoints forced. `last.linear` was renamed by exact match per suffix while everything else uses prefix slicing, so `last.linear.weight_scale` kept its old name -- invisible until now, because both the fp8 path and a dense int8 decode consume the scales before the rename. One Qwen3-VL repack ships 337 `input_scale` activation scales, under a spelling the existing filter did not match. And `model_is_quantized` checked only the config format, so LoRA would have been written directly into int8 buffers; it is now `requires_sidecar_patching()`, which consults the module tree and can be tested on its own. Verified against four real checkpoints: 430/430 key coverage for the int8 and fp8 Krea-2 builds and 357/337 swapped layers for the two Qwen3-VL encoders, all with no missing, extra or orphaned tensors; weights at corr 0.997-0.9999 against an independent quantization of the same model (0.06-0.10 without the un-rotation); and a generation that produces the same photograph as the fp8 build. Co-Authored-By: Claude Opus 5 (1M context) --- invokeai/app/invocations/krea2_denoise.py | 20 +- .../app/invocations/minimax_h3_denoise.py | 2 +- .../custom_int8_convrot_linear.py | 2 +- .../torch_module_autocast.py | 2 +- .../model_manager/load/model_loaders/krea2.py | 198 +++++++++++++++--- .../load/model_loaders/minimax_h3.py | 4 +- .../minimax_h3_state_dict_utils.py | 2 +- .../int8_convrot.py | 63 +++++- tests/app/invocations/test_krea2_denoise.py | 41 +++- .../test_text_encoder_checkpoint.py | 2 +- .../load/test_krea2_loader_boundaries.py | 153 +++++++++++++- .../load/test_krea2_state_dict_utils.py | 158 ++++++++++++++ .../test_int8_convrot.py | 2 +- .../test_int8_convrot_lora_sidecar.py | 10 +- .../test_int8_convrot_state_dict.py | 97 +++++++++ 15 files changed, 707 insertions(+), 49 deletions(-) rename invokeai/backend/{minimax_h3 => quantization}/int8_convrot.py (69%) rename tests/backend/{minimax_h3 => quantization}/test_int8_convrot.py (99%) rename tests/backend/{minimax_h3 => quantization}/test_int8_convrot_lora_sidecar.py (98%) create mode 100644 tests/backend/quantization/test_int8_convrot_state_dict.py diff --git a/invokeai/app/invocations/krea2_denoise.py b/invokeai/app/invocations/krea2_denoise.py index 466568d5050..70e74d668d7 100644 --- a/invokeai/app/invocations/krea2_denoise.py +++ b/invokeai/app/invocations/krea2_denoise.py @@ -2,7 +2,7 @@ import math from contextlib import ExitStack from pathlib import Path -from typing import Callable, Iterator, Optional +from typing import Any, Callable, Iterator, Optional import torch import torchvision.transforms as tv_transforms @@ -46,6 +46,7 @@ from invokeai.backend.patches.layer_patcher import LayerPatcher, PatchSpec from invokeai.backend.patches.lora_conversions.krea2_lora_constants import KREA2_LORA_TRANSFORMER_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw +from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Krea2ConditioningInfo @@ -55,6 +56,20 @@ KREA2_LATENT_CHANNELS = 16 +def requires_sidecar_patching(transformer: Any, model_format: ModelFormat) -> bool: + """Whether LoRA has to be applied as a sidecar rather than written into the weights. + + The format alone does not answer this. A plain ``checkpoint`` Krea-2 may still be an + ``int8_tensorwise`` build, whose Linears the loader replaced with ``Int8ConvrotLinear`` -- + those hold their weights as int8 buffers, which a direct patch cannot write into (and which + could not represent the patched values anyway, the rotation having mixed 256 of them). So the + loaded module tree is consulted, not just the config. + """ + if model_format in (ModelFormat.GGUFQuantized,): + return True + return any(isinstance(module, Int8ConvrotLinear) for module in transformer.modules()) + + @invocation( "krea2_denoise", title="Denoise - Krea-2", @@ -409,6 +424,7 @@ def _run_diffusion(self, context: InvocationContext): ) transformer_config = context.models.get_config(self.transformer.transformer) + # Refined against the loaded module tree below, once the transformer is in hand. model_is_quantized = transformer_config.format in (ModelFormat.GGUFQuantized,) num_train_timesteps = scheduler.config.num_train_timesteps @@ -439,6 +455,8 @@ def _run_diffusion(self, context: InvocationContext): # SDPA for enable_gqa=True, which PyTorch only supports on the math backend — that materializes the # full O(seq^2) score matrix (~5.7 GB per attention at 1280x720, ~40 GB at 2560x1440) and OOMs. Swap # in a memory-efficient processor that expands the KV heads and uses the O(seq) SDPA kernel instead. + model_is_quantized = requires_sidecar_patching(transformer, transformer_config.format) + regional_prompting_state = Krea2RegionalPromptingState() transformer.set_attn_processor(build_krea2_attention_processors(transformer, regional_prompting_state)) # The processors remain installed on the cached transformer after this invocation. Do not let them diff --git a/invokeai/app/invocations/minimax_h3_denoise.py b/invokeai/app/invocations/minimax_h3_denoise.py index 1f75bd6f64f..29d27993337 100644 --- a/invokeai/app/invocations/minimax_h3_denoise.py +++ b/invokeai/app/invocations/minimax_h3_denoise.py @@ -33,7 +33,6 @@ from invokeai.app.services.session_processor.session_processor_common import CanceledException from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.minimax_h3.denoise import denoise -from invokeai.backend.minimax_h3.int8_convrot import Int8ConvrotLinear from invokeai.backend.minimax_h3.packing import ( MINIMAX_H3_CANVAS_MULTIPLE, MINIMAX_H3_FPS, @@ -72,6 +71,7 @@ is_minimax_h3_adaln_layer_path, ) from invokeai.backend.patches.model_patch_raw import ModelPatchRaw +from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState from invokeai.backend.stable_diffusion.diffusion.conditioning_data import MiniMaxH3ConditioningInfo from invokeai.backend.util.devices import TorchDevice diff --git a/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_int8_convrot_linear.py b/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_int8_convrot_linear.py index bc8899a6411..3cac78415b0 100644 --- a/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_int8_convrot_linear.py +++ b/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/custom_modules/custom_int8_convrot_linear.py @@ -1,6 +1,5 @@ import torch -from invokeai.backend.minimax_h3.int8_convrot import Int8ConvrotLinear from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.cast_to_device import cast_to_device from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_linear import ( autocast_linear_forward_sidecar_patches, @@ -9,6 +8,7 @@ CustomModuleMixin, ) from invokeai.backend.patches.layers.param_shape_utils import get_param_shape +from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear class CustomInt8ConvrotLinear(Int8ConvrotLinear, CustomModuleMixin): diff --git a/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/torch_module_autocast.py b/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/torch_module_autocast.py index 5b96ad86d18..c01d7893106 100644 --- a/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/torch_module_autocast.py +++ b/invokeai/backend/model_manager/load/model_cache/torch_module_autocast/torch_module_autocast.py @@ -4,7 +4,6 @@ from diffusers.models.normalization import RMSNorm as DiffusersRMSNorm from invokeai.backend.flux.modules.layers import RMSNorm as FluxRMSNorm -from invokeai.backend.minimax_h3.int8_convrot import Int8ConvrotLinear from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_conv1d import ( CustomConv1d, ) @@ -35,6 +34,7 @@ from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.custom_modules.custom_module_mixin import ( CustomModuleMixin, ) +from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear AUTOCAST_MODULE_TYPE_MAPPING: dict[type[torch.nn.Module], type[torch.nn.Module]] = { torch.nn.Linear: CustomLinear, diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index 774d3a9203c..51156a7d607 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/invokeai/backend/model_manager/load/model_loaders/krea2.py @@ -26,6 +26,12 @@ ) from invokeai.backend.model_manager.util.qwen3_vl import normalize_qwen3vl_rope_config from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader +from invokeai.backend.quantization.int8_convrot import ( + CONVROT_GROUP_SIZE, + Int8ConvrotLinear, + check_int8_scale_layout, + extract_int8_convrot_markers, +) from invokeai.backend.util.devices import TorchDevice if TYPE_CHECKING: @@ -115,6 +121,11 @@ def _dequantize_scaled_fp8(sd: dict[str, Any], dtype: "torch.dtype") -> dict[str out = dict(sd) for scale_key in scale_keys: weight_key = scale_key.replace(".weight_scale", ".weight") + if weight_key in out and not out[weight_key].is_floating_point(): + # An int8_tensorwise layer. Scaling it here without un-rotating it produces a state dict + # that loads cleanly and generates noise, which is the failure the marker path exists to + # prevent -- so leave both the weight and its scale for that path to consume. + continue 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() @@ -124,7 +135,26 @@ def _dequantize_scaled_fp8(sd: dict[str, Any], dtype: "torch.dtype") -> dict[str return out -def _convert_krea2_native_to_diffusers(sd: dict[str, Any]) -> dict[str, Any]: +# The original final-block up/down projections have no counterpart in the diffusers +# ``Krea2FinalLayer`` (a clean AdaLN + linear). Named once because two steps act on them: the +# converter drops them, and the single-file loader drops them *before* dequantizing so it never +# spends work - or trips over an exotic scale layout - on tensors that are about to be discarded. +DISCARDED_NATIVE_FINAL_KEYS = ("last.down", "last.up") + + +def _drop_discarded_native_final_layers(sd: dict[str, Any]) -> dict[str, Any]: + """Remove the dropped final-block projections together with their quantization metadata.""" + doomed = { + f"{path}{suffix}" + for path in DISCARDED_NATIVE_FINAL_KEYS + for suffix in (".weight", ".weight_scale", ".comfy_quant") + } + if not doomed & set(sd): + return sd + return {k: v for k, v in sd.items() if k not in doomed} + + +def _convert_krea2_native_to_diffusers(sd: dict[str, Any], *, key_map: dict[str, str] | None = None) -> dict[str, Any]: """Convert a native/ComfyUI-format Krea-2 state dict (e.g. GGUF) to diffusers Krea2Transformer2DModel keys. Top-level module renames:: @@ -158,7 +188,7 @@ def _convert_krea2_native_to_diffusers(sd: dict[str, Any]) -> dict[str, Any]: _put_unique_key(new_sd, key, value, source=key, source_of=source_of, what="Krea-2 checkpoint") continue # Drop original-only final-block projections (no diffusers equivalent). - if key in ("last.down.weight", "last.up.weight"): + if key in tuple(f"{p}.weight" for p in DISCARDED_NATIVE_FINAL_KEYS): continue k = key @@ -181,10 +211,11 @@ def _convert_krea2_native_to_diffusers(sd: dict[str, Any]) -> dict[str, Any]: k = "txt_in.linear_1." + k[len("txtmlp.1.") :] elif k.startswith("txtmlp.3."): k = "txt_in.linear_2." + k[len("txtmlp.3.") :] - elif k == "last.linear.weight": - k = "final_layer.linear.weight" - elif k == "last.linear.bias": - k = "final_layer.linear.bias" + elif k.startswith("last.linear."): + # Prefix rather than an exact match per suffix: a quantized build carries + # `last.linear.weight_scale` too, and an exact rule leaves it behind under the old name + # -- which the loader only notices as a missing scale, well after the rename. + k = "final_layer.linear." + k[len("last.linear.") :] elif k == "last.norm.scale": k = "final_layer.norm.weight" elif k == "last.modulation.lin": @@ -215,6 +246,8 @@ def _convert_krea2_native_to_diffusers(sd: dict[str, Any]) -> dict[str, Any]: value = torch.as_tensor(_to_plain_tensor(value)).reshape(6, -1) _put_unique_key(new_sd, k, value, source=key, source_of=source_of, what="Krea-2 checkpoint") + if key_map is not None: + key_map[key] = k return new_sd @@ -316,9 +349,15 @@ def _load_model( class Krea2CheckpointModel(ModelLoader): """Class to load Krea-2 transformer models from single-file checkpoints (safetensors). - Handles plain bf16/fp16 checkpoints as well as ComfyUI 'scaled fp8' checkpoints (fp8 weight + - ``.weight_scale``), and both the diffusers and native/ComfyUI key naming. Apply the fp8-storage - setting to keep the (large) transformer fp8-resident; otherwise it loads in full precision. + Handles plain bf16/fp16 checkpoints, ComfyUI 'scaled fp8' checkpoints (fp8 weight + + ``.weight_scale``) and ComfyUI 'int8_tensorwise' checkpoints (int8 weight + per-output-channel + ``.weight_scale`` + a ``.comfy_quant`` marker, optionally convrot-rotated), in both the diffusers + and native/ComfyUI key naming. Apply the fp8-storage setting to keep the (large) transformer + fp8-resident; otherwise it loads in full precision. + + The int8 build is decoded to dense weights rather than kept int8-resident: on this model that + costs nothing, because fp8 storage and int8 storage measure the same 12.0 GiB, and the reason to + keep int8 resident is int8 *compute*, which needs a kernel InvokeAI does not have yet. """ def _load_model( @@ -348,21 +387,37 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: sd = load_file(model_path) sd = _strip_comfyui_prefix(sd) + # Discard what the key conversion below would discard anyway, before anything is spent on + # it. One repack quantizes `last.up` with a blockwise scale grid this decode does not + # implement; refusing a tensor that is on its way to the bin would be an odd way to fail. + sd = _drop_discarded_native_final_layers(sd) + # ComfyUI 'int8_tensorwise' checkpoints (Comfy-Org's *_int8_convrot.safetensors): dequantize + # and un-rotate the marked layers. This must run BEFORE the fp8 path, which keys off + # `.weight_scale` alone and would happily scale an int8 weight without ever un-rotating it -- + # producing a state dict that loads cleanly and generates noise. It must also run BEFORE the + # key conversion below, which renames `.attn.wq.weight` by substring and so carries + # `.weight_scale` along but NOT `.comfy_quant`; decoding here keeps marker and weight paired. + int8_markers = extract_int8_convrot_markers(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) + sd = _drop_unconsumed_quantization_sidecars(sd) # Native/ComfyUI key naming → diffusers Krea2Transformer2DModel keys. + key_map: dict[str, str] = {} if _is_native_krea2_format(sd): - sd = _convert_krea2_native_to_diffusers(sd) + sd = _convert_krea2_native_to_diffusers(sd, key_map=key_map) + quantized = _resolve_quantized_module_paths(int8_markers, key_map) with accelerate.init_empty_weights(): model = Krea2Transformer2DModel(**KREA2_TRANSFORMER_CONFIG) - new_sd_size = sum(ten.nelement() * model_dtype.itemsize for ten in sd.values()) + # int8 payloads are one byte, not two: reserving the compute dtype's width for them would + # ask the cache to free ~12 GB that this load never uses. + new_sd_size = sum(ten.nelement() * max(ten.element_size(), 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) + _cast_unquantized(sd, model_dtype, quantized) + _swap_in_int8_linears(model, sd, quantized) model.load_state_dict(sd, assign=True, strict=False) _reject_incomplete_load(model, what="Krea-2 single-file checkpoint") @@ -471,7 +526,7 @@ def _load_model( ) -def _remap_qwen3vl_singlefile_keys(sd: dict[str, Any]) -> dict[str, Any]: +def _remap_qwen3vl_singlefile_keys(sd: dict[str, Any], *, key_map: dict[str, str] | None = None) -> dict[str, Any]: """Remap ComfyUI single-file Qwen3-VL keys to the transformers ``Qwen3VLModel`` layout. ComfyUI/native layout uses a single ``model.`` prefix for both towers; transformers splits them: @@ -488,13 +543,96 @@ def _remap_qwen3vl_singlefile_keys(sd: dict[str, Any]) -> dict[str, Any]: 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) + new_key = key 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) + new_key = "language_model." + key + _put_unique_key(out, new_key, v, source=k, source_of=source_of, what=what) + if key_map is not None: + key_map[k] = new_key return out +def _drop_unconsumed_quantization_sidecars(sd: dict[str, Any]) -> dict[str, Any]: + """Remove quantization metadata no loader here consumes. + + - ``.comfy_quant`` markers whose format was handled elsewhere, or not at all. + - ``.input_scale`` / ``.scale_input``: activation scales for W8A8 inference. This code + dequantizes the weight and computes in bf16, so there is nothing to apply them to. (Both + spellings appear in the wild; one Qwen3-VL repack ships 337 of the former.) + + `load_state_dict(strict=False)` would ignore them, but they are still cast and still counted + against the RAM reservation - and a loader that later switches to strict would fail on them. + """ + return { + k: v + for k, v in sd.items() + if not (isinstance(k, str) and (k.endswith(".comfy_quant") or "input_scale" in k or "scale_input" in k)) + } + + +def _resolve_quantized_module_paths( + markers: dict[str, dict[str, Any]], key_map: dict[str, str] +) -> dict[str, dict[str, Any]]: + """Re-key markers from the checkpoint's names to the built model's names. + + The markers are read in the checkpoint's key space, because that is the only place where a + marker and its weight are reliably paired: the key conversions rename `.weight` (and, by + substring, `.weight_scale`) but leave `.comfy_quant` behind on the old name. Following the + weight's own rename is therefore the only mapping that cannot drift from the conversion. + """ + resolved: dict[str, dict[str, Any]] = {} + for path, marker in markers.items(): + weight_key = key_map.get(f"{path}.weight", f"{path}.weight") + resolved[weight_key[: -len(".weight")]] = marker + return resolved + + +def _swap_in_int8_linears(model: Any, sd: dict[str, Any], quantized: dict[str, dict[str, Any]]) -> None: + """Replace each quantized ``nn.Linear`` with an ``Int8ConvrotLinear`` sized from the state dict. + + The weights stay int8 and rotated as stored; the layer dequantizes and derotates per forward. + That keeps a 12 GB checkpoint at 12 GB resident instead of the ~24 GB a dense decode would + produce, on every platform and without the fp8-storage opt-in (which is off by default and + unavailable outside CUDA/XPU). + + Its persistent buffers are named ``weight``/``weight_scale`` -- the checkpoint's own spelling -- + so the ``load_state_dict`` that follows assigns the quantized tensors straight into them. + """ + for path, marker in quantized.items(): + weight, scale = sd.get(f"{path}.weight"), sd.get(f"{path}.weight_scale") + if weight is None or scale is None: + raise ValueError( + f"'{path}' is marked int8_tensorwise but is missing its " + f"{'weight' if weight is None else 'weight_scale'}." + ) + check_int8_scale_layout(path, weight, scale) + parent_path, _, attribute = path.rpartition(".") + setattr( + model.get_submodule(parent_path) if parent_path else model, + attribute, + Int8ConvrotLinear( + weight=weight, + weight_scale=scale, + convrot=bool(marker.get("convrot", False)), + bias=sd.get(f"{path}.bias"), + group_size=int(marker.get("convrot_groupsize", CONVROT_GROUP_SIZE)), + ), + ) + + +def _cast_unquantized(sd: dict[str, Any], dtype: "torch.dtype", quantized: dict[str, dict[str, Any]]) -> None: + """Cast the dense tensors to the compute dtype, leaving the quantized payloads alone. + + An int8 weight cast to bf16 is no longer int8, and its float32 scale is what + ``Int8ConvrotLinear`` multiplies by -- both have to reach ``load_state_dict`` as stored. + """ + pinned = {key for path in quantized for key in (f"{path}.weight", f"{path}.weight_scale")} + for key in sd: + if key not in pinned: + sd[key] = sd[key].to(dtype) + + def _reject_incomplete_load(model: Any, *, what: str) -> None: """Raise if a ``load_state_dict(strict=False)`` left required tensors on the meta device. @@ -575,27 +713,35 @@ 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)) + # ComfyUI 'int8_tensorwise' encoders, decoded first for the same two reasons as the + # transformer above -- and for a third one here: an int8 layer's `.weight_scale` would + # otherwise make the fp8 detection below answer yes, and the encoder would be kept + # "fp8-resident" over weights that were never fp8. + int8_markers = extract_int8_convrot_markers(sd) # 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( - getattr(t, "dtype", None) in (torch.float8_e4m3fn, torch.float8_e5m2) for t in sd.values() - ) + # An int8 layer's scale is not evidence of fp8, so those are excluded -- otherwise an int8 + # encoder would be kept "fp8-resident" over weights that were never fp8. + source_is_fp8 = any( + isinstance(k, str) and k.endswith(".weight_scale") and k[: -len(".weight_scale")] not in int8_markers + for k in sd + ) 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) + sd = _drop_unconsumed_quantization_sidecars(sd) + key_map: dict[str, str] = {} + sd = _remap_qwen3vl_singlefile_keys(sd, key_map=key_map) + quantized = _resolve_quantized_module_paths(int8_markers, key_map) 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()) + new_sd_size = sum(ten.nelement() * max(ten.element_size(), 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) + _cast_unquantized(sd, model_dtype, quantized) + _swap_in_int8_linears(model, sd, quantized) model.load_state_dict(sd, assign=True, strict=False) _reject_incomplete_load(model, what="Qwen3-VL encoder checkpoint") diff --git a/invokeai/backend/model_manager/load/model_loaders/minimax_h3.py b/invokeai/backend/model_manager/load/model_loaders/minimax_h3.py index 3f5e8a7c343..8bf94492cc8 100644 --- a/invokeai/backend/model_manager/load/model_loaders/minimax_h3.py +++ b/invokeai/backend/model_manager/load/model_loaders/minimax_h3.py @@ -173,7 +173,6 @@ def _load_transformer_from_singlefile(self, config: Main_Checkpoint_MiniMaxH3_Co from invokeai.backend.minimax_h3.contiguous_attention import ( patch_minimax_h3_attention_contiguous_qkv, ) - from invokeai.backend.minimax_h3.int8_convrot import Int8ConvrotLinear from invokeai.backend.minimax_h3.transformer_minimax_h3_pruned import ( MiniMaxH3PrunedTransformer3DModel, set_curve_modulation_dtype, @@ -182,6 +181,7 @@ def _load_transformer_from_singlefile(self, config: Main_Checkpoint_MiniMaxH3_Co convert_minimax_h3_checkpoint_to_diffusers, read_comfy_quant_markers, ) + from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear model_path = Path(config.path) @@ -317,12 +317,12 @@ def _load_text_encoder_from_singlefile(self, config: AnyModelConfig) -> AnyModel from safetensors.torch import load_file from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration - from invokeai.backend.minimax_h3.int8_convrot import Int8ConvrotLinear from invokeai.backend.minimax_h3.text_conditioning import MINIMAX_H3_TEXT_ENCODER_LAYER from invokeai.backend.model_manager.load.model_loaders.minimax_h3_state_dict_utils import ( convert_minimax_h3_text_encoder_checkpoint, read_comfy_quant_markers, ) + from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear model_path = Path(config.path) diff --git a/invokeai/backend/model_manager/load/model_loaders/minimax_h3_state_dict_utils.py b/invokeai/backend/model_manager/load/model_loaders/minimax_h3_state_dict_utils.py index e492064e5b5..42069a70069 100644 --- a/invokeai/backend/model_manager/load/model_loaders/minimax_h3_state_dict_utils.py +++ b/invokeai/backend/model_manager/load/model_loaders/minimax_h3_state_dict_utils.py @@ -34,7 +34,7 @@ import torch -from invokeai.backend.minimax_h3.int8_convrot import parse_comfy_quant_marker +from invokeai.backend.quantization.int8_convrot import parse_comfy_quant_marker def read_comfy_quant_markers(path: Path) -> dict[str, dict[str, Any]]: diff --git a/invokeai/backend/minimax_h3/int8_convrot.py b/invokeai/backend/quantization/int8_convrot.py similarity index 69% rename from invokeai/backend/minimax_h3/int8_convrot.py rename to invokeai/backend/quantization/int8_convrot.py index 378cbf4b991..032e4c3f571 100644 --- a/invokeai/backend/minimax_h3/int8_convrot.py +++ b/invokeai/backend/quantization/int8_convrot.py @@ -1,7 +1,12 @@ """Runtime support for Comfy "int8_tensorwise + convrot" quantized linears. -The Comfy-Org single-file H3 transformers store their four big per-block linears -(qkv/out/fc1/fc2) as symmetric per-output-channel int8: +``int8_tensorwise`` is a ComfyUI-wide scheme, not one architecture's format: the same +spelling appears on MiniMax H3 (the first consumer here), Krea-2, and whatever Comfy-Org +publishes next. It therefore lives beside the other schemes in ``backend/quantization`` +rather than in an architecture package, so a second architecture costs a call site and not +a second copy of the mathematics. + +A quantized linear stores its weight as symmetric per-output-channel int8: - ``.weight``: int8 ``[out, in]`` - ``.weight_scale``: float32 ``[out, 1]`` @@ -23,11 +28,12 @@ weight and scale resident (4.6x smaller than bf16) and materializes the dequantized, derotated bf16 weight per forward call. The derotation is a ``[out, in/256, 256] @ [256, 256]`` matmul — a rounding error next to the -transformer forward itself — and the transient bf16 weight (<= ~310 MB for the -largest layer) lives inside the denoise node's working-memory reservation. +transformer forward itself — and the transient bf16 weight (<= ~310 MB for H3's +largest layer) has to fit inside the calling node's working-memory reservation. """ import json +from typing import Any import torch import torch.nn.functional as F @@ -36,6 +42,8 @@ _HADAMARD_SEED = ((1, 1, 1, -1), (1, 1, -1, 1), (1, -1, 1, 1), (-1, 1, 1, 1)) +INT8_TENSORWISE_FORMAT = "int8_tensorwise" + def build_regular_hadamard(size: int, dtype: torch.dtype = torch.float32) -> torch.Tensor: """Normalized regular Hadamard matrix of a power-of-4 size (CPU tensor).""" @@ -83,9 +91,9 @@ class Int8ConvrotLinear(torch.nn.Module): ``AUTOCAST_MODULE_TYPE_MAPPING``), which enables sidecar LoRA patches and lets a partial load leave some int8 buffers on the CPU — ``forward``'s per-call ``.to(device)`` then streams them (at half the bf16 byte count) instead of failing outright. Fully-resident - operation (~20 GiB free VRAM for the pruned transformer) remains the intended regime; - streamed layers pay a per-forward PCIe cost, and the diffusers-folder bf16 model is still - the better citizen on small cards. + operation remains the intended regime (~20 GiB free VRAM for H3's pruned transformer); + streamed layers pay a per-forward PCIe cost, and an unquantized model is still the better + citizen on small cards. """ def __init__( @@ -142,3 +150,44 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def extra_repr(self) -> str: return f"in_features={self.in_features}, out_features={self.out_features}, convrot={self.convrot}" + + +def extract_int8_convrot_markers(sd: dict[str, Any]) -> dict[str, dict[str, Any]]: + """Pop every ``int8_tensorwise`` marker out of ``sd``, keyed by the layer path it names. + + Markers for other formats are left in place, together with their weights and scales: ComfyUI's + fp8_scaled repacks share this key layout and belong to the fp8 path. Deciding per marker rather + than per file is also what keeps mixed-precision checkpoints correct -- one Krea-2 build leaves + 40 weights in bf16 with no marker at all, and both Qwen3-VL encoders leave over 200. + + Use this when the loader intends to keep the weights int8 and swap in + :class:`Int8ConvrotLinear`, which is what every loader here does. + """ + markers = {} + for key in [k for k in sd if isinstance(k, str) and k.endswith(".comfy_quant")]: + marker = parse_comfy_quant_marker(sd[key]) + if marker.get("format") != INT8_TENSORWISE_FORMAT: + continue + markers[key[: -len(".comfy_quant")]] = marker + del sd[key] + return markers + + +def check_int8_scale_layout(path: str, weight: torch.Tensor, scale: torch.Tensor) -> None: + """Refuse a scale granularity this decode does not implement. + + Two layouts are supported, because they are the two this dequantization is correct for: + per-output-channel (``[out, 1]`` or ``[out]``) and per-tensor (a scalar). Some repacks emit a + blockwise grid instead - a 6144x6144 weight with a ``[48, 48]`` scale is a 128x128 block grid - + which needs a different multiply. Left to broadcasting that either raises somewhere less + informative or, for an unlucky shape, silently scales the wrong axis. + """ + rows = weight.shape[0] if weight.dim() else 1 + if scale.dim() == 0 or tuple(scale.shape) in {(1,), (1, 1)}: + return + if tuple(scale.shape) in {(rows,), (rows, 1)}: + return + raise ValueError( + f"'{path}' has a {tuple(scale.shape)} scale for a {tuple(weight.shape)} weight, which is " + "neither per-output-channel nor per-tensor. Blockwise scale grids are not implemented." + ) diff --git a/tests/app/invocations/test_krea2_denoise.py b/tests/app/invocations/test_krea2_denoise.py index 670bac2d67d..4bbde4cffe8 100644 --- a/tests/app/invocations/test_krea2_denoise.py +++ b/tests/app/invocations/test_krea2_denoise.py @@ -6,9 +6,14 @@ import torch from invokeai.app.invocations.fields import DenoiseMaskField, Krea2ConditioningField, LatentsField, TensorField -from invokeai.app.invocations.krea2_denoise import KREA2_LATENT_CHANNELS, Krea2DenoiseInvocation +from invokeai.app.invocations.krea2_denoise import ( + KREA2_LATENT_CHANNELS, + Krea2DenoiseInvocation, + requires_sidecar_patching, +) from invokeai.app.invocations.model import ModelIdentifierField, TransformerField from invokeai.backend.model_manager.taxonomy import BaseModelType, Krea2VariantType, ModelFormat, ModelType +from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, Krea2ConditioningInfo @@ -320,6 +325,12 @@ def set_attn_processor(self, processor) -> None: # The real Krea2Transformer2DModel exposes this; denoise swaps in a memory-efficient attention processor. self.installed_processors = processor + def modules(self): + # Denoise walks the module tree to decide whether LoRA has to be applied as a sidecar: an + # int8_tensorwise build's Linears hold their weights as buffers, which a direct patch + # cannot write into. This double stands in for an unquantized model. + return iter(()) + def __call__(self, *, hidden_states, encoder_hidden_states, **_kwargs): self.conditioning_values.append(float(encoder_hidden_states.mean())) # The real transformer concatenates [text, image] before attention, so this is the sequence length a @@ -548,6 +559,9 @@ class _PositionIdChecker: def set_attn_processor(self, processor) -> None: pass + def modules(self): + return iter(()) + def __call__(self, *, hidden_states, encoder_hidden_states, position_ids, **_kwargs): text_len = encoder_hidden_states.shape[1] pos_len = position_ids.shape[0] @@ -757,3 +771,28 @@ def test_regional_attention_memory_includes_masks_build_scratch_and_dtype_sized_ assert Krea2DenoiseInvocation._regional_attention_mask_bytes(positive, negative, torch.bfloat16) == 500 assert Krea2DenoiseInvocation._regional_attention_mask_bytes(positive, negative, torch.float32) == 740 assert Krea2DenoiseInvocation._regional_attention_mask_bytes(positive, None, torch.bfloat16) == 400 + + +class TestRequiresSidecarPatching: + """LoRA cannot be written into an int8 buffer, and the config format does not say when one + is present: an `int8_tensorwise` Krea-2 is a plain `checkpoint` as far as the config knows.""" + + class _Tree: + def __init__(self, *modules) -> None: + self._modules = modules + + def modules(self): + return iter(self._modules) + + def test_an_unquantized_checkpoint_is_patched_directly(self) -> None: + assert not requires_sidecar_patching(self._Tree(torch.nn.Linear(2, 2)), ModelFormat.Checkpoint) + + def test_a_gguf_model_is_still_recognised_by_its_format(self) -> None: + assert requires_sidecar_patching(self._Tree(), ModelFormat.GGUFQuantized) + + def test_an_int8_convrot_checkpoint_is_recognised_by_its_modules(self) -> None: + int8_linear = Int8ConvrotLinear( + weight=torch.zeros(4, 4, dtype=torch.int8), weight_scale=torch.ones(4, 1), convrot=False + ) + tree = self._Tree(torch.nn.Linear(2, 2), int8_linear) + assert requires_sidecar_patching(tree, ModelFormat.Checkpoint) diff --git a/tests/backend/minimax_h3/test_text_encoder_checkpoint.py b/tests/backend/minimax_h3/test_text_encoder_checkpoint.py index e00fabfeb8c..dc4de8de887 100644 --- a/tests/backend/minimax_h3/test_text_encoder_checkpoint.py +++ b/tests/backend/minimax_h3/test_text_encoder_checkpoint.py @@ -202,8 +202,8 @@ def test_te_converted_keys_match_real_model_exactly() -> None: from transformers import Qwen3VLForConditionalGeneration from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig - from invokeai.backend.minimax_h3.int8_convrot import Int8ConvrotLinear from invokeai.backend.model_manager.util.qwen3_vl import normalize_qwen3vl_rope_config + from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear from tests.model_identification.stripped_model_on_disk import StrippedModelOnDisk fixture_dir = ( 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..2446039d2e4 100644 --- a/tests/backend/model_manager/load/test_krea2_loader_boundaries.py +++ b/tests/backend/model_manager/load/test_krea2_loader_boundaries.py @@ -8,14 +8,19 @@ Main_Diffusers_Krea2_Config, Main_GGUF_Krea2_Config, ) -from invokeai.backend.model_manager.configs.qwen3_vl_encoder import Qwen3VLEncoder_Qwen3VLEncoder_Config +from invokeai.backend.model_manager.configs.qwen3_vl_encoder import ( + Qwen3VLEncoder_Checkpoint_Config, + Qwen3VLEncoder_Qwen3VLEncoder_Config, +) from invokeai.backend.model_manager.load.model_loaders.krea2 import ( Krea2CheckpointModel, Krea2DiffusersModel, Krea2GGUFCheckpointModel, + Qwen3VLEncoderCheckpointLoader, Qwen3VLEncoderLoader, ) from invokeai.backend.model_manager.taxonomy import Krea2VariantType, SubModelType +from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear class _TinyKrea2Transformer(torch.nn.Module): @@ -57,6 +62,73 @@ def test_single_file_loader_constructs_and_materializes_model(monkeypatch, tmp_p ram_cache.make_room.assert_called_once() +def test_single_file_loader_decodes_an_int8_convrot_checkpoint(monkeypatch, tmp_path) -> None: + """The decode has to be *wired in*, not merely available. + + Every other test of this feature exercises the helpers directly, so deleting the call from + `_load_from_singlefile` would leave them all green while the loader silently produced a + scaled-but-still-rotated model. This one drives the loader itself and checks the weight that + actually reaches the module. + """ + import json + + import diffusers + import safetensors.torch + + from invokeai.backend.quantization.int8_convrot import CONVROT_GROUP_SIZE, build_regular_hadamard + + torch.manual_seed(0) + original = torch.randn(4, CONVROT_GROUP_SIZE) + hadamard = build_regular_hadamard(CONVROT_GROUP_SIZE) + rotated = (original.view(4, 1, CONVROT_GROUP_SIZE) @ hadamard.T).view(4, CONVROT_GROUP_SIZE) + scale = rotated.abs().amax(dim=1, keepdim=True) / 127.0 + marker = json.dumps({"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": CONVROT_GROUP_SIZE}) + # A marker is always `.comfy_quant`; a quantized weight is a submodule's, never the + # model's own, so the fixture mirrors that. + state_dict = { + "proj.weight": torch.clamp(torch.round(rotated / scale), -128, 127).to(torch.int8), + "proj.weight_scale": scale.to(torch.float32), + "proj.comfy_quant": torch.frombuffer(bytearray(marker.encode("utf-8")), dtype=torch.uint8), + } + + class _TinyInt8Krea2Transformer(torch.nn.Module): + def __init__(self, **_kwargs) -> None: + super().__init__() + self.proj = torch.nn.Linear(CONVROT_GROUP_SIZE, 4, bias=False) + + checkpoint_path = tmp_path / "krea2_int8_convrot.safetensors" + checkpoint_path.touch() + config = Main_Checkpoint_Krea2_Config.model_construct( + path=str(checkpoint_path), variant=Krea2VariantType.Turbo, fp8_storage=None + ) + loader = object.__new__(Krea2CheckpointModel) + loader._ram_cache = SimpleNamespace(make_room=MagicMock()) + loader._apply_fp8_layerwise_casting = lambda model, _config, _submodel: model + + monkeypatch.setattr(diffusers, "Krea2Transformer2DModel", _TinyInt8Krea2Transformer, raising=False) + monkeypatch.setattr(safetensors.torch, "load_file", lambda _path: state_dict) + monkeypatch.setattr( + "invokeai.backend.model_manager.load.model_loaders.krea2.TorchDevice.choose_torch_device", + lambda: torch.device("cpu"), + ) + monkeypatch.setattr( + "invokeai.backend.model_manager.load.model_loaders.krea2.TorchDevice.choose_bfloat16_safe_dtype", + lambda _device: torch.float32, + ) + + model = loader._load_from_singlefile(config) + + # The weight stays quantized: that is the whole point of the swap, and it is what keeps a + # 12 GB checkpoint at 12 GB instead of the ~24 GB a dense decode would produce. + assert isinstance(model.proj, Int8ConvrotLinear) + assert model.proj.weight.dtype is torch.int8 + + # And it still computes the un-rotated weight, which the fp8 path alone would never produce. + dequantized = model.proj._dequantized_weight(torch.device("cpu"), torch.float32).flatten() + assert torch.corrcoef(torch.stack([dequantized, original.flatten()]))[0, 1] > 0.999 + assert torch.corrcoef(torch.stack([dequantized, rotated.flatten()]))[0, 1].abs() < 0.2 + + def test_diffusers_loader_reaches_transformer_from_pretrained(monkeypatch, tmp_path) -> None: config = Main_Diffusers_Krea2_Config.model_construct(path=str(tmp_path), repo_variant=None) loader = object.__new__(Krea2DiffusersModel) @@ -110,6 +182,85 @@ def test_gguf_loader_constructs_and_materializes_model(monkeypatch, tmp_path) -> assert torch.equal(model.weight, torch.ones(2, 2)) +def test_checkpoint_encoder_loader_decodes_int8_and_does_not_call_it_fp8(monkeypatch, tmp_path) -> None: + """The encoder path carries the same hazard as the transformer, plus one of its own. + + `source_is_fp8` answers yes to any `.weight_scale` key. Left to itself an int8 encoder would + be scaled, never un-rotated, and then kept "fp8-resident" over weights that were never fp8. + Decoding first removes the int8 scales, so the detection sees what it was written to see. + """ + import json + + import transformers + from safetensors import torch as safetensors_torch + + from invokeai.backend.quantization.int8_convrot import CONVROT_GROUP_SIZE, build_regular_hadamard + + torch.manual_seed(0) + original = torch.randn(4, CONVROT_GROUP_SIZE) + hadamard = build_regular_hadamard(CONVROT_GROUP_SIZE) + rotated = (original.view(4, 1, CONVROT_GROUP_SIZE) @ hadamard.T).view(4, CONVROT_GROUP_SIZE) + scale = rotated.abs().amax(dim=1, keepdim=True) / 127.0 + marker = json.dumps({"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": CONVROT_GROUP_SIZE}) + state_dict = { + "proj.weight": torch.clamp(torch.round(rotated / scale), -128, 127).to(torch.int8), + "proj.weight_scale": scale.to(torch.float32), + "proj.comfy_quant": torch.frombuffer(bytearray(marker.encode("utf-8")), dtype=torch.uint8), + } + + class _TinyLanguageModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.proj = torch.nn.Linear(CONVROT_GROUP_SIZE, 4, bias=False) + + class _TinyEncoder(torch.nn.Module): + """`_remap_qwen3vl_singlefile_keys` routes bare keys under `language_model.`, so the + fixture keeps the checkpoint's spelling and lets the remap do its job.""" + + def __init__(self) -> None: + super().__init__() + self.language_model = _TinyLanguageModel() + + @classmethod + def _from_config(cls, _config): + return cls() + + checkpoint_path = tmp_path / "qwen3vl_int8_convrot.safetensors" + checkpoint_path.touch() + config = Qwen3VLEncoder_Checkpoint_Config.model_construct(path=str(checkpoint_path), name="encoder") + + fp8_calls: list = [] + loader = object.__new__(Qwen3VLEncoderCheckpointLoader) + loader._ram_cache = SimpleNamespace(make_room=MagicMock()) + loader._torch_device = torch.device("cpu") + loader._logger = MagicMock() + loader._load_hf_config = lambda: SimpleNamespace() + loader._apply_fp8_to_nn_module = lambda *a, **k: fp8_calls.append(a) + + monkeypatch.setattr(transformers, "Qwen3VLModel", _TinyEncoder, raising=False) + monkeypatch.setattr(safetensors_torch, "load_file", lambda _path: state_dict) + monkeypatch.setattr( + "invokeai.backend.model_manager.load.model_loaders.krea2.TorchDevice.choose_torch_device", + lambda: torch.device("cpu"), + ) + monkeypatch.setattr( + "invokeai.backend.model_manager.load.model_loaders.krea2.TorchDevice.choose_bfloat16_safe_dtype", + lambda _device: torch.float32, + ) + monkeypatch.setattr( + "invokeai.backend.model_manager.load.model_loaders.krea2._device_supports_fp8_storage", + lambda _device, _logger: True, + ) + + model = loader._load_text_encoder(config) + + assert isinstance(model.language_model.proj, Int8ConvrotLinear) + assert model.language_model.proj.weight.dtype is torch.int8 + dequantized = model.language_model.proj._dequantized_weight(torch.device("cpu"), torch.float32).flatten() + assert torch.corrcoef(torch.stack([dequantized, original.flatten()]))[0, 1] > 0.999 + assert fp8_calls == [], "an int8 encoder must not be treated as an fp8 one" + + def test_directory_encoder_loader_reaches_transformers_from_pretrained(monkeypatch, tmp_path) -> None: import transformers 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..8598d4e0519 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 @@ -18,12 +18,20 @@ KREA2_TRANSFORMER_CONFIG, _convert_krea2_native_to_diffusers, _dequantize_scaled_fp8, + _drop_discarded_native_final_layers, + _drop_unconsumed_quantization_sidecars, _is_native_krea2_format, _normalize_qwen3vl_rope_config, _reject_incomplete_load, _remap_qwen3vl_singlefile_keys, + _resolve_quantized_module_paths, _strip_comfyui_prefix, ) +from invokeai.backend.quantization.int8_convrot import ( + CONVROT_GROUP_SIZE, + build_regular_hadamard, + extract_int8_convrot_markers, +) class TestNormalizeQwen3vlRopeConfig: @@ -376,3 +384,153 @@ 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]}" + + +def _quantized_int8_layer(path: str, weight: torch.Tensor) -> dict: + """A convrot-rotated int8 layer in the native key spelling, as Comfy-Org ships it.""" + import json + + out_f, in_f = weight.shape + h = build_regular_hadamard(CONVROT_GROUP_SIZE, dtype=weight.dtype) + rotated = (weight.view(out_f, in_f // CONVROT_GROUP_SIZE, CONVROT_GROUP_SIZE) @ h.T).view(out_f, in_f) + scale = rotated.abs().amax(dim=1, keepdim=True) / 127.0 + marker = json.dumps({"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": CONVROT_GROUP_SIZE}) + return { + f"{path}.weight": torch.clamp(torch.round(rotated / scale), -128, 127).to(torch.int8), + f"{path}.weight_scale": scale.to(torch.float32), + f"{path}.comfy_quant": torch.frombuffer(bytearray(marker.encode("utf-8")), dtype=torch.uint8), + } + + +class TestInt8LayersSurviveTheLoaderPipeline: + """Two steps between the file and the model would quietly ruin an int8 layer. + + Neither fails loudly if it is broken - both produce a state dict that loads cleanly and + generates noise - so each is pinned by a test that shows the damage. + """ + + def test_the_fp8_fold_leaves_an_int8_weight_and_its_scale_alone(self) -> None: + """The fp8 path keys off `.weight_scale` and never consults the marker. Folding the scale + into an int8 weight without un-rotating it is exactly the silent corruption the marker + exists to prevent, so it skips anything that is not floating point.""" + torch.manual_seed(0) + sd = _quantized_int8_layer("blocks.0.mlp.down", torch.randn(64, 2 * CONVROT_GROUP_SIZE)) + before = {k: v.clone() for k, v in sd.items()} + + out = _dequantize_scaled_fp8(sd, torch.float32) + + assert out["blocks.0.mlp.down.weight"].dtype is torch.int8 + assert torch.equal(out["blocks.0.mlp.down.weight"], before["blocks.0.mlp.down.weight"]) + assert "blocks.0.mlp.down.weight_scale" in out, "the scale the int8 layer still needs was consumed" + + def test_an_fp8_layer_is_still_folded(self) -> None: + sd = { + "blocks.0.attn.wq.weight": torch.ones(4, 4, dtype=torch.float8_e4m3fn), + "blocks.0.attn.wq.weight_scale": torch.tensor(2.0), + } + out = _dequantize_scaled_fp8(sd, torch.float32) + assert "blocks.0.attn.wq.weight_scale" not in out + assert torch.allclose(out["blocks.0.attn.wq.weight"], torch.full((4, 4), 2.0)) + + def test_the_key_conversion_renames_the_scale_but_orphans_the_marker(self) -> None: + """Why the module paths are resolved by following the weight's own rename rather than the + marker's key: the within-block renames are substring replacements of `.attn.wq.weight`, + which `.weight_scale` contains and `.comfy_quant` does not.""" + torch.manual_seed(1) + sd = _quantized_int8_layer("blocks.0.attn.wq", torch.randn(32, CONVROT_GROUP_SIZE)) + + converted = _convert_krea2_native_to_diffusers(sd) + + assert "transformer_blocks.0.attn.to_q.weight" in converted + assert "transformer_blocks.0.attn.to_q.weight_scale" in converted + assert "transformer_blocks.0.attn.wq.comfy_quant" in converted + assert "transformer_blocks.0.attn.to_q.comfy_quant" not in converted + + def test_the_marker_is_re_keyed_onto_the_module_the_weight_landed_on(self) -> None: + torch.manual_seed(2) + sd = _quantized_int8_layer("blocks.0.attn.wq", torch.randn(32, CONVROT_GROUP_SIZE)) + markers = extract_int8_convrot_markers(sd) + + key_map: dict[str, str] = {} + _convert_krea2_native_to_diffusers(sd, key_map=key_map) + + assert _resolve_quantized_module_paths(markers, key_map) == { + "transformer_blocks.0.attn.to_q": markers["blocks.0.attn.wq"] + } + + def test_a_diffusers_named_checkpoint_needs_no_re_keying(self) -> None: + """No conversion runs, so the empty map has to resolve to the paths as they are.""" + markers = {"transformer_blocks.0.attn.to_q": {"format": "int8_tensorwise"}} + assert _resolve_quantized_module_paths(markers, {}) == markers + + +class TestUnconsumedQuantizationSidecars: + def test_markers_and_activation_scales_are_dropped(self) -> None: + """`input_scale` is an activation scale for W8A8 inference; this code dequantizes the + weight and computes in bf16, so there is nothing to apply it to. One Qwen3-VL repack + ships 337 of them, under a spelling the previous filter did not match.""" + sd = { + "layer.weight": torch.ones(2, 2), + "layer.comfy_quant": torch.zeros(4, dtype=torch.uint8), + "layer.input_scale": torch.ones(1), + "other.scale_input": torch.ones(1), + } + assert set(_drop_unconsumed_quantization_sidecars(sd)) == {"layer.weight"} + + def test_a_clean_state_dict_is_unchanged(self) -> None: + sd = {"layer.weight": torch.ones(2, 2), "layer.bias": torch.zeros(2)} + assert set(_drop_unconsumed_quantization_sidecars(sd)) == set(sd) + + +class TestDiscardedFinalProjections: + def test_the_dropped_projections_take_their_quantization_metadata_with_them(self) -> None: + """`last.down`/`last.up` have no diffusers counterpart, and one real repack quantizes + `last.up` with a scale layout the decode refuses. Dropping them before the decode keeps a + tensor that is on its way to the bin from failing the load.""" + sd = { + "last.up.weight": torch.zeros(4, 4, dtype=torch.int8), + "last.up.weight_scale": torch.ones(2, 2), + "last.up.comfy_quant": torch.zeros(8, dtype=torch.uint8), + "last.down.weight": torch.zeros(4, 4), + "last.linear.weight": torch.ones(4, 4), + } + out = _drop_discarded_native_final_layers(sd) + assert set(out) == {"last.linear.weight"} + + def test_a_state_dict_without_them_is_returned_unchanged(self) -> None: + sd = {"last.linear.weight": torch.ones(2, 2)} + assert _drop_discarded_native_final_layers(sd) is sd + + +class TestEverySuffixSurvivesTheKeyConversion: + """A quantized layer travels as three keys, and the converter has to move all of them. + + The within-block renames are substring replacements of `.attn.wq.weight`, so `.weight_scale` + rides along by construction. The top-level renames are prefix slices, so they carry any + suffix. `last.linear` used to be neither -- an exact match per suffix, which silently left + `last.linear.weight_scale` under its old name. + """ + + @pytest.mark.parametrize( + ("native", "diffusers"), + [ + ("blocks.0.attn.wq", "transformer_blocks.0.attn.to_q"), + ("blocks.3.mlp.down", "transformer_blocks.3.ff.down"), + ("txtfusion.refiner_blocks.1.attn.wv", "text_fusion.refiner_blocks.1.attn.to_v"), + ("first", "img_in"), + ("tmlp.0", "time_embed.linear_1"), + ("tmlp.2", "time_embed.linear_2"), + ("tproj.1", "time_mod_proj"), + ("txtmlp.1", "txt_in.linear_1"), + ("txtmlp.3", "txt_in.linear_2"), + ("last.linear", "final_layer.linear"), + ], + ) + def test_the_weight_and_its_scale_land_on_the_same_module(self, native: str, diffusers: str) -> None: + sd = {f"{native}.weight": torch.zeros(4, 4, dtype=torch.int8), f"{native}.weight_scale": torch.ones(4, 1)} + key_map: dict[str, str] = {} + out = _convert_krea2_native_to_diffusers(sd, key_map=key_map) + + assert f"{diffusers}.weight" in out + assert f"{diffusers}.weight_scale" in out, "the scale was left behind under its old name" + assert key_map[f"{native}.weight"] == f"{diffusers}.weight" diff --git a/tests/backend/minimax_h3/test_int8_convrot.py b/tests/backend/quantization/test_int8_convrot.py similarity index 99% rename from tests/backend/minimax_h3/test_int8_convrot.py rename to tests/backend/quantization/test_int8_convrot.py index 0a2ef9db6fd..e4b0b890e64 100644 --- a/tests/backend/minimax_h3/test_int8_convrot.py +++ b/tests/backend/quantization/test_int8_convrot.py @@ -11,7 +11,7 @@ import pytest import torch -from invokeai.backend.minimax_h3.int8_convrot import ( +from invokeai.backend.quantization.int8_convrot import ( CONVROT_GROUP_SIZE, Int8ConvrotLinear, build_regular_hadamard, diff --git a/tests/backend/minimax_h3/test_int8_convrot_lora_sidecar.py b/tests/backend/quantization/test_int8_convrot_lora_sidecar.py similarity index 98% rename from tests/backend/minimax_h3/test_int8_convrot_lora_sidecar.py rename to tests/backend/quantization/test_int8_convrot_lora_sidecar.py index 47f75f08e66..efed3e89bad 100644 --- a/tests/backend/minimax_h3/test_int8_convrot_lora_sidecar.py +++ b/tests/backend/quantization/test_int8_convrot_lora_sidecar.py @@ -8,11 +8,6 @@ import torch -from invokeai.backend.minimax_h3.int8_convrot import ( - Int8ConvrotLinear, - build_regular_hadamard, - dequantize_convrot_weight, -) from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import ( AUTOCAST_MODULE_TYPE_MAPPING, apply_custom_layers_to_model, @@ -20,6 +15,11 @@ from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.layers.lora_layer import LoRALayer from invokeai.backend.patches.model_patch_raw import ModelPatchRaw +from invokeai.backend.quantization.int8_convrot import ( + Int8ConvrotLinear, + build_regular_hadamard, + dequantize_convrot_weight, +) IN_FEATURES = 16 OUT_FEATURES = 12 diff --git a/tests/backend/quantization/test_int8_convrot_state_dict.py b/tests/backend/quantization/test_int8_convrot_state_dict.py new file mode 100644 index 00000000000..fea15c6818b --- /dev/null +++ b/tests/backend/quantization/test_int8_convrot_state_dict.py @@ -0,0 +1,97 @@ +"""Tests for the state-dict side of the int8_convrot scheme: which layers a loader picks up, +and which scale layouts it will act on. + +The per-tensor mathematics is covered by ``test_int8_convrot.py``. What is pinned here is the +part a loader can get wrong without anything raising — every failure below would otherwise +surface as a model that loads cleanly and generates noise. +""" + +import json + +import pytest +import torch + +from invokeai.backend.quantization.int8_convrot import ( + CONVROT_GROUP_SIZE, + check_int8_scale_layout, + extract_int8_convrot_markers, +) + +MARKER = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": CONVROT_GROUP_SIZE} + + +def _marker_blob(marker: dict) -> torch.Tensor: + return torch.frombuffer(bytearray(json.dumps(marker).encode("utf-8")), dtype=torch.uint8) + + +class TestWhichLayersAreClaimed: + def test_an_int8_marker_is_claimed_and_removed(self) -> None: + sd = { + "blocks.0.attn.wq.weight": torch.zeros(4, CONVROT_GROUP_SIZE, dtype=torch.int8), + "blocks.0.attn.wq.weight_scale": torch.ones(4, 1), + "blocks.0.attn.wq.comfy_quant": _marker_blob(MARKER), + } + markers = extract_int8_convrot_markers(sd) + + assert markers == {"blocks.0.attn.wq": MARKER} + # The marker is consumed; the weight and its scale stay for the loader to install. + assert set(sd) == {"blocks.0.attn.wq.weight", "blocks.0.attn.wq.weight_scale"} + + def test_a_marker_for_another_format_is_left_untouched(self) -> None: + """ComfyUI's fp8_scaled repacks share this key layout and belong to the fp8 path.""" + sd = { + "layer.weight": torch.zeros(4, 4, dtype=torch.float8_e4m3fn), + "layer.weight_scale": torch.ones(1), + "layer.comfy_quant": _marker_blob({"format": "float8_e4m3fn"}), + } + assert extract_int8_convrot_markers(sd) == {} + assert "layer.comfy_quant" in sd + + def test_unmarked_weights_are_not_claimed(self) -> None: + """Mixed precision is the norm, not the exception: one Krea-2 build leaves 40 weights in + bf16 with no marker, and the two Qwen3-VL encoders leave over 200 each. A decision made + per file rather than per tensor would be wrong on all of them.""" + sd = { + "txtfusion.0.weight": torch.zeros(8, 8, dtype=torch.bfloat16), + "blocks.0.attn.wq.weight": torch.zeros(4, CONVROT_GROUP_SIZE, dtype=torch.int8), + "blocks.0.attn.wq.weight_scale": torch.ones(4, 1), + "blocks.0.attn.wq.comfy_quant": _marker_blob(MARKER), + } + assert set(extract_int8_convrot_markers(sd)) == {"blocks.0.attn.wq"} + + def test_a_state_dict_with_no_markers_is_unchanged(self) -> None: + sd = {"a.weight": torch.zeros(2, 2)} + assert extract_int8_convrot_markers(sd) == {} + assert set(sd) == {"a.weight"} + + def test_the_group_size_travels_with_each_marker(self) -> None: + """Every marker in the Krea-2 build says 256, but the flag is per tensor and another + producer may vary it, so it is read rather than assumed.""" + sd = {"layer.comfy_quant": _marker_blob({"format": "int8_tensorwise", "convrot_groupsize": 64})} + assert extract_int8_convrot_markers(sd)["layer"]["convrot_groupsize"] == 64 + + +class TestTheScaleLayoutsSeenInRealCheckpoints: + """Three appear across the checkpoints this was tested against, so none is hypothetical.""" + + def test_per_output_channel(self) -> None: + check_int8_scale_layout("layer", torch.zeros(64, 256, dtype=torch.int8), torch.ones(64, 1)) + + def test_per_output_channel_without_the_trailing_axis(self) -> None: + check_int8_scale_layout("layer", torch.zeros(64, 256, dtype=torch.int8), torch.ones(64)) + + def test_a_scalar_scale(self) -> None: + """One Qwen3-VL encoder stores every scale as a bare scalar, and Krea-2's + `txtfusion.projector` does the same.""" + check_int8_scale_layout("layer", torch.zeros(1, 12, dtype=torch.int8), torch.tensor(0.01)) + + def test_a_blockwise_grid_is_refused_by_name(self) -> None: + """Also observed: a 6144x6144 weight with a [48, 48] scale, i.e. a 128x128 block grid. + Broadcasting would either raise somewhere uninformative or, for an unlucky shape, + silently scale the wrong axis.""" + with pytest.raises(ValueError, match=r"Blockwise scale grids"): + check_int8_scale_layout("last.up", torch.zeros(512, 512, dtype=torch.int8), torch.ones(4, 4)) + + def test_the_message_names_the_layer_and_both_shapes(self) -> None: + with pytest.raises(ValueError, match=r"'last\.up' has a \(4, 4\) scale for a \(512, 512\) weight"): + check_int8_scale_layout("last.up", torch.zeros(512, 512, dtype=torch.int8), torch.ones(4, 4)) From 8c9de55aa9a2f2611ada512ef7b7c6bc4f42bf0e Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Tue, 1 Sep 2026 19:56:37 +0200 Subject: [PATCH 2/2] feat(quantization): read int8_convrot Z-Image checkpoints Z-Image is the model this format is actually useful for. Krea-2 is 12.6 GiB as int8 -- too big for the 12 GB cards that fall back to GGUF today -- while Z-Image is 5.8 GiB, and measured here it drops a 1024px generation from 11.74 GiB resident to 6.04 GiB at 6 % more time. Moves the architecture-agnostic half of the Krea-2 work into `backend/quantization/int8_convrot.py` (resolve_quantized_module_paths, swap_in_int8_linears, cast_unquantized, drop_unconsumed_quantization_sidecars) so Z-Image imports it rather than copying it. No behaviour change. Z-Image's own hazard is the fused QKV: `attention.qkv.weight` is split into to_q/to_k/to_v, and the split already handled the weight but dropped the scale and the marker under the fused name -- 408 keys that then reached a strict `load_state_dict`. The per-output-channel scale now splits with its weight and the marker is copied to all three. That decision is made from the key suffix, not the tensor shape: a 72-byte JSON marker is also divisible by three, and cutting it into thirds yields three fragments of broken JSON. Because this converter carries `.comfy_quant` onto the final module names, the markers are read after the conversion and need no re-keying -- the opposite of Krea-2, whose converter orphans them. Adds a guard for an int8 weight with no marker: it would be handed to a float Linear and fail only at forward time, if at all. Verified against the real 5.75 GiB checkpoint and the bf16 diffusers release of the same model: 272 quantized modules with no orphans, no missing and no extra keys; weights at corr 0.99978-1.00000 against unquantized ground truth (0.054-0.065 without the un-rotation) at 0.85-1 % relative error; and a generation that produces the same photograph. Co-Authored-By: Claude Opus 5 (1M context) --- .../model_manager/load/model_loaders/krea2.py | 103 ++------------- .../load/model_loaders/z_image.py | 58 ++++++-- invokeai/backend/quantization/int8_convrot.py | 80 +++++++++++ .../load/test_krea2_state_dict_utils.py | 12 +- .../load/test_z_image_loader_boundaries.py | 125 ++++++++++++++++++ .../load/test_z_image_state_dict_utils.py | 78 ++++++++++- 6 files changed, 350 insertions(+), 106 deletions(-) create mode 100644 tests/backend/model_manager/load/test_z_image_loader_boundaries.py diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index 51156a7d607..0dd981874ac 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/invokeai/backend/model_manager/load/model_loaders/krea2.py @@ -27,10 +27,11 @@ from invokeai.backend.model_manager.util.qwen3_vl import normalize_qwen3vl_rope_config from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.quantization.int8_convrot import ( - CONVROT_GROUP_SIZE, - Int8ConvrotLinear, - check_int8_scale_layout, + cast_unquantized, + drop_unconsumed_quantization_sidecars, extract_int8_convrot_markers, + resolve_quantized_module_paths, + swap_in_int8_linears, ) from invokeai.backend.util.devices import TorchDevice @@ -402,12 +403,12 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: # 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) - sd = _drop_unconsumed_quantization_sidecars(sd) + sd = drop_unconsumed_quantization_sidecars(sd) # Native/ComfyUI key naming → diffusers Krea2Transformer2DModel keys. key_map: dict[str, str] = {} if _is_native_krea2_format(sd): sd = _convert_krea2_native_to_diffusers(sd, key_map=key_map) - quantized = _resolve_quantized_module_paths(int8_markers, key_map) + quantized = resolve_quantized_module_paths(int8_markers, key_map) with accelerate.init_empty_weights(): model = Krea2Transformer2DModel(**KREA2_TRANSFORMER_CONFIG) @@ -416,8 +417,8 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: # ask the cache to free ~12 GB that this load never uses. new_sd_size = sum(ten.nelement() * max(ten.element_size(), model_dtype.itemsize) for ten in sd.values()) self._ram_cache.make_room(new_sd_size) - _cast_unquantized(sd, model_dtype, quantized) - _swap_in_int8_linears(model, sd, quantized) + cast_unquantized(sd, model_dtype, quantized) + swap_in_int8_linears(model, sd, quantized) model.load_state_dict(sd, assign=True, strict=False) _reject_incomplete_load(model, what="Krea-2 single-file checkpoint") @@ -553,86 +554,6 @@ def _remap_qwen3vl_singlefile_keys(sd: dict[str, Any], *, key_map: dict[str, str return out -def _drop_unconsumed_quantization_sidecars(sd: dict[str, Any]) -> dict[str, Any]: - """Remove quantization metadata no loader here consumes. - - - ``.comfy_quant`` markers whose format was handled elsewhere, or not at all. - - ``.input_scale`` / ``.scale_input``: activation scales for W8A8 inference. This code - dequantizes the weight and computes in bf16, so there is nothing to apply them to. (Both - spellings appear in the wild; one Qwen3-VL repack ships 337 of the former.) - - `load_state_dict(strict=False)` would ignore them, but they are still cast and still counted - against the RAM reservation - and a loader that later switches to strict would fail on them. - """ - return { - k: v - for k, v in sd.items() - if not (isinstance(k, str) and (k.endswith(".comfy_quant") or "input_scale" in k or "scale_input" in k)) - } - - -def _resolve_quantized_module_paths( - markers: dict[str, dict[str, Any]], key_map: dict[str, str] -) -> dict[str, dict[str, Any]]: - """Re-key markers from the checkpoint's names to the built model's names. - - The markers are read in the checkpoint's key space, because that is the only place where a - marker and its weight are reliably paired: the key conversions rename `.weight` (and, by - substring, `.weight_scale`) but leave `.comfy_quant` behind on the old name. Following the - weight's own rename is therefore the only mapping that cannot drift from the conversion. - """ - resolved: dict[str, dict[str, Any]] = {} - for path, marker in markers.items(): - weight_key = key_map.get(f"{path}.weight", f"{path}.weight") - resolved[weight_key[: -len(".weight")]] = marker - return resolved - - -def _swap_in_int8_linears(model: Any, sd: dict[str, Any], quantized: dict[str, dict[str, Any]]) -> None: - """Replace each quantized ``nn.Linear`` with an ``Int8ConvrotLinear`` sized from the state dict. - - The weights stay int8 and rotated as stored; the layer dequantizes and derotates per forward. - That keeps a 12 GB checkpoint at 12 GB resident instead of the ~24 GB a dense decode would - produce, on every platform and without the fp8-storage opt-in (which is off by default and - unavailable outside CUDA/XPU). - - Its persistent buffers are named ``weight``/``weight_scale`` -- the checkpoint's own spelling -- - so the ``load_state_dict`` that follows assigns the quantized tensors straight into them. - """ - for path, marker in quantized.items(): - weight, scale = sd.get(f"{path}.weight"), sd.get(f"{path}.weight_scale") - if weight is None or scale is None: - raise ValueError( - f"'{path}' is marked int8_tensorwise but is missing its " - f"{'weight' if weight is None else 'weight_scale'}." - ) - check_int8_scale_layout(path, weight, scale) - parent_path, _, attribute = path.rpartition(".") - setattr( - model.get_submodule(parent_path) if parent_path else model, - attribute, - Int8ConvrotLinear( - weight=weight, - weight_scale=scale, - convrot=bool(marker.get("convrot", False)), - bias=sd.get(f"{path}.bias"), - group_size=int(marker.get("convrot_groupsize", CONVROT_GROUP_SIZE)), - ), - ) - - -def _cast_unquantized(sd: dict[str, Any], dtype: "torch.dtype", quantized: dict[str, dict[str, Any]]) -> None: - """Cast the dense tensors to the compute dtype, leaving the quantized payloads alone. - - An int8 weight cast to bf16 is no longer int8, and its float32 scale is what - ``Int8ConvrotLinear`` multiplies by -- both have to reach ``load_state_dict`` as stored. - """ - pinned = {key for path in quantized for key in (f"{path}.weight", f"{path}.weight_scale")} - for key in sd: - if key not in pinned: - sd[key] = sd[key].to(dtype) - - def _reject_incomplete_load(model: Any, *, what: str) -> None: """Raise if a ``load_state_dict(strict=False)`` left required tensors on the meta device. @@ -729,10 +650,10 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod ) 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) - sd = _drop_unconsumed_quantization_sidecars(sd) + sd = drop_unconsumed_quantization_sidecars(sd) key_map: dict[str, str] = {} sd = _remap_qwen3vl_singlefile_keys(sd, key_map=key_map) - quantized = _resolve_quantized_module_paths(int8_markers, key_map) + quantized = resolve_quantized_module_paths(int8_markers, key_map) te_config = self._load_hf_config() with accelerate.init_empty_weights(): @@ -740,8 +661,8 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod new_sd_size = sum(ten.nelement() * max(ten.element_size(), model_dtype.itemsize) for ten in sd.values()) self._ram_cache.make_room(new_sd_size) - _cast_unquantized(sd, model_dtype, quantized) - _swap_in_int8_linears(model, sd, quantized) + cast_unquantized(sd, model_dtype, quantized) + swap_in_int8_linears(model, sd, quantized) model.load_state_dict(sd, assign=True, strict=False) _reject_incomplete_load(model, what="Qwen3-VL encoder checkpoint") diff --git a/invokeai/backend/model_manager/load/model_loaders/z_image.py b/invokeai/backend/model_manager/load/model_loaders/z_image.py index 0ca5dfb010a..0164405e7ec 100644 --- a/invokeai/backend/model_manager/load/model_loaders/z_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/z_image.py @@ -35,11 +35,39 @@ SubModelType, ) from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader +from invokeai.backend.quantization.int8_convrot import ( + cast_unquantized, + drop_unconsumed_quantization_sidecars, + extract_int8_convrot_markers, + swap_in_int8_linears, +) 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 from invokeai.backend.qwen3.qwen3_tokenizer import load_bundled_qwen3_tokenizer from invokeai.backend.util.devices import TorchDevice +# Side-channel suffixes whose value is per output channel, and therefore splits with the weight. +# Everything else describes the layer as a whole -- a `comfy_quant` JSON blob most of all, whose +# byte length says nothing about how many outputs it covers. +_PER_ROW_SIDECHANNEL_SUFFIXES = ("weight_scale", "scale_weight") + + +def _split_qkv_sidechannel(value: Any, suffix: str, target: str) -> Any: + """One third of a fused-QKV side-channel tensor, or the whole thing when it is not per-row. + + A per-output-channel scale is ``[3 * dim, 1]`` and splits exactly like the weight it belongs + to. A per-tensor scale, and any marker, is copied to all three unchanged. The decision is made + from the suffix rather than the shape: a 72-byte JSON marker is also divisible by three. + """ + if suffix not in _PER_ROW_SIDECHANNEL_SUFFIXES: + return value + shape = getattr(value, "shape", None) + if shape is None or len(shape) == 0 or shape[0] <= 3 or shape[0] % 3 != 0: + return value + index = ("to_q", "to_k", "to_v").index(target) + dim = shape[0] // 3 + return value[index * dim : (index + 1) * dim] + 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,10 +125,12 @@ 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 + # Quantization side-channels have to follow their weight through the split, or they + # are left behind under a module name that no longer exists -- and a loader then sees + # a quantized weight with no scale, or a scale it cannot pair with anything. if suffix not in ("weight", "bias"): - new_sd[key] = value + for target in ("to_q", "to_k", "to_v"): + new_sd[f"{prefix}.attention.{target}.{suffix}"] = _split_qkv_sidechannel(value, suffix, target) continue # Split the fused QKV tensor into Q, K, V @@ -473,13 +503,25 @@ def _load_from_singlefile( 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()]) + # ComfyUI 'int8_tensorwise' checkpoints. The markers are read *after* the key conversion + # above, which carries them (and their scales) through the fused-QKV split onto the module + # names the model actually has -- so no re-keying is needed here. + int8_markers = extract_int8_convrot_markers(sd) + sd = drop_unconsumed_quantization_sidecars(sd) + orphans = sorted(k for k, v in sd.items() if v.dtype is torch.int8 and k[: -len(".weight")] not in int8_markers) + if orphans: + raise ValueError( + f"Z-Image checkpoint has {len(orphans)} int8 weight(s) with no `comfy_quant` marker, " + f"e.g. {orphans[:3]}. Loading them would produce a model that runs and generates noise." + ) + + # int8 payloads are one byte, not two: reserving the compute dtype's width for them would + # ask the cache to free memory this load never uses. + new_sd_size = sum(ten.nelement() * max(ten.element_size(), model_dtype.itemsize) for ten in sd.values()) self._ram_cache.make_room(new_sd_size) - # Convert to target dtype - for k in sd.keys(): - sd[k] = sd[k].to(model_dtype) + cast_unquantized(sd, model_dtype, int8_markers) + swap_in_int8_linears(model, sd, int8_markers) model.load_state_dict(sd, assign=True) diff --git a/invokeai/backend/quantization/int8_convrot.py b/invokeai/backend/quantization/int8_convrot.py index 032e4c3f571..f891ef4b1f0 100644 --- a/invokeai/backend/quantization/int8_convrot.py +++ b/invokeai/backend/quantization/int8_convrot.py @@ -191,3 +191,83 @@ def check_int8_scale_layout(path: str, weight: torch.Tensor, scale: torch.Tensor f"'{path}' has a {tuple(scale.shape)} scale for a {tuple(weight.shape)} weight, which is " "neither per-output-channel nor per-tensor. Blockwise scale grids are not implemented." ) + + +def drop_unconsumed_quantization_sidecars(sd: dict[str, Any]) -> dict[str, Any]: + """Remove quantization metadata no loader here consumes. + + - ``.comfy_quant`` markers whose format was handled elsewhere, or not at all. + - ``.input_scale`` / ``.scale_input``: activation scales for W8A8 inference. This code + dequantizes the weight and computes in bf16, so there is nothing to apply them to. (Both + spellings appear in the wild; one Qwen3-VL repack ships 337 of the former.) + + `load_state_dict(strict=False)` would ignore them, but they are still cast and still counted + against the RAM reservation - and a loader that later switches to strict would fail on them. + """ + return { + k: v + for k, v in sd.items() + if not (isinstance(k, str) and (k.endswith(".comfy_quant") or "input_scale" in k or "scale_input" in k)) + } + + +def resolve_quantized_module_paths( + markers: dict[str, dict[str, Any]], key_map: dict[str, str] +) -> dict[str, dict[str, Any]]: + """Re-key markers from the checkpoint's names to the built model's names. + + The markers are read in the checkpoint's key space, because that is the only place where a + marker and its weight are reliably paired: the key conversions rename `.weight` (and, by + substring, `.weight_scale`) but leave `.comfy_quant` behind on the old name. Following the + weight's own rename is therefore the only mapping that cannot drift from the conversion. + """ + resolved: dict[str, dict[str, Any]] = {} + for path, marker in markers.items(): + weight_key = key_map.get(f"{path}.weight", f"{path}.weight") + resolved[weight_key[: -len(".weight")]] = marker + return resolved + + +def swap_in_int8_linears(model: torch.nn.Module, sd: dict[str, Any], quantized: dict[str, dict[str, Any]]) -> None: + """Replace each quantized ``nn.Linear`` with an ``Int8ConvrotLinear`` sized from the state dict. + + The weights stay int8 and rotated as stored; the layer dequantizes and derotates per forward. + That keeps a 12 GB checkpoint at 12 GB resident instead of the ~24 GB a dense decode would + produce, on every platform and without the fp8-storage opt-in (which is off by default and + unavailable outside CUDA/XPU). + + Its persistent buffers are named ``weight``/``weight_scale`` -- the checkpoint's own spelling -- + so the ``load_state_dict`` that follows assigns the quantized tensors straight into them. + """ + for path, marker in quantized.items(): + weight, scale = sd.get(f"{path}.weight"), sd.get(f"{path}.weight_scale") + if weight is None or scale is None: + raise ValueError( + f"'{path}' is marked int8_tensorwise but is missing its " + f"{'weight' if weight is None else 'weight_scale'}." + ) + check_int8_scale_layout(path, weight, scale) + parent_path, _, attribute = path.rpartition(".") + setattr( + model.get_submodule(parent_path) if parent_path else model, + attribute, + Int8ConvrotLinear( + weight=weight, + weight_scale=scale, + convrot=bool(marker.get("convrot", False)), + bias=sd.get(f"{path}.bias"), + group_size=int(marker.get("convrot_groupsize", CONVROT_GROUP_SIZE)), + ), + ) + + +def cast_unquantized(sd: dict[str, Any], dtype: torch.dtype, quantized: dict[str, dict[str, Any]]) -> None: + """Cast the dense tensors to the compute dtype, leaving the quantized payloads alone. + + An int8 weight cast to bf16 is no longer int8, and its float32 scale is what + ``Int8ConvrotLinear`` multiplies by -- both have to reach ``load_state_dict`` as stored. + """ + pinned = {key for path in quantized for key in (f"{path}.weight", f"{path}.weight_scale")} + for key in sd: + if key not in pinned: + sd[key] = sd[key].to(dtype) 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 8598d4e0519..31cc6ff7bc6 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 @@ -19,18 +19,18 @@ _convert_krea2_native_to_diffusers, _dequantize_scaled_fp8, _drop_discarded_native_final_layers, - _drop_unconsumed_quantization_sidecars, _is_native_krea2_format, _normalize_qwen3vl_rope_config, _reject_incomplete_load, _remap_qwen3vl_singlefile_keys, - _resolve_quantized_module_paths, _strip_comfyui_prefix, ) from invokeai.backend.quantization.int8_convrot import ( CONVROT_GROUP_SIZE, build_regular_hadamard, + drop_unconsumed_quantization_sidecars, extract_int8_convrot_markers, + resolve_quantized_module_paths, ) @@ -454,14 +454,14 @@ def test_the_marker_is_re_keyed_onto_the_module_the_weight_landed_on(self) -> No key_map: dict[str, str] = {} _convert_krea2_native_to_diffusers(sd, key_map=key_map) - assert _resolve_quantized_module_paths(markers, key_map) == { + assert resolve_quantized_module_paths(markers, key_map) == { "transformer_blocks.0.attn.to_q": markers["blocks.0.attn.wq"] } def test_a_diffusers_named_checkpoint_needs_no_re_keying(self) -> None: """No conversion runs, so the empty map has to resolve to the paths as they are.""" markers = {"transformer_blocks.0.attn.to_q": {"format": "int8_tensorwise"}} - assert _resolve_quantized_module_paths(markers, {}) == markers + assert resolve_quantized_module_paths(markers, {}) == markers class TestUnconsumedQuantizationSidecars: @@ -475,11 +475,11 @@ def test_markers_and_activation_scales_are_dropped(self) -> None: "layer.input_scale": torch.ones(1), "other.scale_input": torch.ones(1), } - assert set(_drop_unconsumed_quantization_sidecars(sd)) == {"layer.weight"} + assert set(drop_unconsumed_quantization_sidecars(sd)) == {"layer.weight"} def test_a_clean_state_dict_is_unchanged(self) -> None: sd = {"layer.weight": torch.ones(2, 2), "layer.bias": torch.zeros(2)} - assert set(_drop_unconsumed_quantization_sidecars(sd)) == set(sd) + assert set(drop_unconsumed_quantization_sidecars(sd)) == set(sd) class TestDiscardedFinalProjections: diff --git a/tests/backend/model_manager/load/test_z_image_loader_boundaries.py b/tests/backend/model_manager/load/test_z_image_loader_boundaries.py new file mode 100644 index 00000000000..d111199bd11 --- /dev/null +++ b/tests/backend/model_manager/load/test_z_image_loader_boundaries.py @@ -0,0 +1,125 @@ +"""Loader-level tests for the Z-Image single-file path. + +The state-dict helpers are covered elsewhere. What is pinned here is that the loader *calls* +them: deleting the swap would leave every unit test green while the loader produced a model that +loads cleanly and generates noise. These drive `_load_from_singlefile` itself and check what +reaches the module. +""" + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from invokeai.backend.model_manager.configs.main import Main_Checkpoint_ZImage_Config +from invokeai.backend.model_manager.load.model_loaders.z_image import ZImageCheckpointModel +from invokeai.backend.quantization.int8_convrot import CONVROT_GROUP_SIZE, Int8ConvrotLinear, build_regular_hadamard + +MARKER = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": CONVROT_GROUP_SIZE} + + +def _marker_blob(marker: dict) -> torch.Tensor: + return torch.frombuffer(bytearray(json.dumps(marker).encode("utf-8")), dtype=torch.uint8) + + +def _quantize_convrot(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Mirror of comfy-quants: rotate along the input dim, then per-output-channel int8.""" + out_f, in_f = weight.shape + h = build_regular_hadamard(CONVROT_GROUP_SIZE, dtype=weight.dtype) + rotated = (weight.view(out_f, in_f // CONVROT_GROUP_SIZE, CONVROT_GROUP_SIZE) @ h.T).view(out_f, in_f) + scale = rotated.abs().amax(dim=1, keepdim=True) / 127.0 + return torch.clamp(torch.round(rotated / scale), -128, 127).to(torch.int8), scale.to(torch.float32) + + +class _TinyBlock(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.proj = torch.nn.Linear(CONVROT_GROUP_SIZE, 4, bias=False) + + +class _TinyZImage(torch.nn.Module): + """Stands in for ZImageTransformer2DModel. `layers.` is one of the loader's valid prefixes, + so the state dict survives its filter.""" + + def __init__(self, **_kwargs) -> None: + super().__init__() + self.layers = torch.nn.ModuleList([_TinyBlock()]) + + +def _driver(monkeypatch, tmp_path, state_dict: dict) -> tuple[ZImageCheckpointModel, Main_Checkpoint_ZImage_Config]: + import diffusers + from safetensors import torch as safetensors_torch + + checkpoint = tmp_path / "z_image_int8_convrot.safetensors" + checkpoint.touch() + config = Main_Checkpoint_ZImage_Config.model_construct(path=str(checkpoint), name="z-image") + + loader = object.__new__(ZImageCheckpointModel) + loader._ram_cache = SimpleNamespace(make_room=MagicMock()) + loader._logger = MagicMock() + loader._torch_device = torch.device("cpu") + loader._torch_dtype = torch.float32 + loader._apply_fp8_layerwise_casting = lambda model, _config, _submodel: model + + monkeypatch.setattr(diffusers, "ZImageTransformer2DModel", _TinyZImage, raising=False) + monkeypatch.setattr(safetensors_torch, "load_file", lambda _path: state_dict) + monkeypatch.setattr( + "invokeai.backend.model_manager.load.model_loaders.z_image.TorchDevice.choose_torch_device", + lambda: torch.device("cpu"), + ) + monkeypatch.setattr( + "invokeai.backend.model_manager.load.model_loaders.z_image.TorchDevice.choose_bfloat16_safe_dtype", + lambda _device: torch.float32, + ) + return loader, config + + +def test_an_int8_checkpoint_loads_int8_resident_and_un_rotated(monkeypatch, tmp_path) -> None: + torch.manual_seed(0) + original = torch.randn(4, CONVROT_GROUP_SIZE) + quantized, scale = _quantize_convrot(original) + state_dict = { + "layers.0.proj.weight": quantized, + "layers.0.proj.weight_scale": scale, + "layers.0.proj.comfy_quant": _marker_blob(MARKER), + } + loader, config = _driver(monkeypatch, tmp_path, state_dict) + + model = loader._load_from_singlefile(config) + + # Resident, not decoded: that is what keeps a 5.8 GB checkpoint at 5.8 GB. + assert isinstance(model.layers[0].proj, Int8ConvrotLinear) + assert model.layers[0].proj.weight.dtype is torch.int8 + + dequantized = model.layers[0].proj._dequantized_weight(torch.device("cpu"), torch.float32).flatten() + assert torch.corrcoef(torch.stack([dequantized, original.flatten()]))[0, 1] > 0.999 + # And specifically not the scaled-but-still-rotated weight, which is what a loader that only + # applied the scale would produce -- silently. + rotated = (quantized.float() * scale).flatten() + assert torch.corrcoef(torch.stack([rotated, original.flatten()]))[0, 1].abs() < 0.2 + + +def test_an_int8_weight_without_a_marker_is_refused(monkeypatch, tmp_path) -> None: + """A quantized weight the loader does not recognise would be handed to a float Linear and only + fail at forward time, if at all. Refuse at load, and say which layers.""" + torch.manual_seed(1) + quantized, scale = _quantize_convrot(torch.randn(4, CONVROT_GROUP_SIZE)) + state_dict = {"layers.0.proj.weight": quantized, "layers.0.proj.weight_scale": scale} + loader, config = _driver(monkeypatch, tmp_path, state_dict) + + with pytest.raises(ValueError, match=r"int8 weight\(s\) with no `comfy_quant` marker"): + loader._load_from_singlefile(config) + + +def test_an_unquantized_checkpoint_is_unaffected(monkeypatch, tmp_path) -> None: + torch.manual_seed(2) + weight = torch.randn(4, CONVROT_GROUP_SIZE) + loader, config = _driver(monkeypatch, tmp_path, {"layers.0.proj.weight": weight}) + + model = loader._load_from_singlefile(config) + + assert isinstance(model.layers[0].proj, torch.nn.Linear) + assert not isinstance(model.layers[0].proj, Int8ConvrotLinear) + assert torch.equal(model.layers[0].proj.weight, weight) 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..242ba397b9d 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,15 @@ """Unit tests for the Z-Image GGUF/ComfyUI -> diffusers state-dict converter.""" +import json + +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, + _split_qkv_sidechannel, +) +from invokeai.backend.quantization.int8_convrot import extract_int8_convrot_markers 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 +71,72 @@ 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]) + + +def _marker_blob(marker: dict) -> torch.Tensor: + return torch.frombuffer(bytearray(json.dumps(marker).encode("utf-8")), dtype=torch.uint8) + + +MARKER = {"format": "int8_tensorwise", "convrot": True, "convrot_groupsize": 256} + + +class TestTheFusedQkvSplitCarriesQuantizationMetadata: + """A quantized fused QKV travels as three keys, and all three have to reach the same three + modules. The weight already split; the scale and the marker did not, and were left behind + under a module name the model does not have.""" + + def test_a_per_output_channel_scale_splits_with_its_weight(self) -> None: + scale = torch.arange(3 * 4, dtype=torch.float32).reshape(3 * 4, 1) + pieces = [_split_qkv_sidechannel(scale, "weight_scale", t) for t in ("to_q", "to_k", "to_v")] + assert [tuple(p.shape) for p in pieces] == [(4, 1)] * 3 + assert torch.equal(torch.cat(pieces), scale) + + def test_a_marker_is_copied_whole_to_all_three(self) -> None: + """The regression that motivated the suffix check: a 72-byte JSON blob is divisible by + three, so a rule based on the tensor's shape cuts it into three fragments of broken JSON.""" + blob = _marker_blob(MARKER) + assert len(blob) % 3 == 0, "the fixture has to reproduce the divisible-by-three trap" + for target in ("to_q", "to_k", "to_v"): + piece = _split_qkv_sidechannel(blob, "comfy_quant", target) + assert torch.equal(piece, blob) + assert json.loads(bytes(piece.numpy().tobytes()).decode()) == MARKER + + def test_a_per_tensor_scale_is_copied_rather_than_split(self) -> None: + for scale in (torch.tensor(0.5), torch.tensor([0.5])): + assert torch.equal(_split_qkv_sidechannel(scale, "weight_scale", "to_k"), scale) + + def test_weight_scale_and_marker_land_on_the_same_three_modules(self) -> None: + prefix = "layers.0.attention" + sd = { + f"{prefix}.qkv.weight": torch.arange(3 * 4 * 8, dtype=torch.int8).reshape(3 * 4, 8), + f"{prefix}.qkv.weight_scale": torch.arange(3 * 4, dtype=torch.float32).reshape(3 * 4, 1), + f"{prefix}.qkv.comfy_quant": _marker_blob(MARKER), + # `x_embedder.` is what makes the loader run this conversion at all. + "x_embedder.weight": torch.zeros(2, 2), + } + out = _convert_z_image_gguf_to_diffusers(sd) + + for target in ("to_q", "to_k", "to_v"): + assert f"{prefix}.{target}.weight" in out + assert f"{prefix}.{target}.weight_scale" in out + assert f"{prefix}.{target}.comfy_quant" in out + assert not any(".qkv." in k for k in out), "the fused keys must not survive" + + def test_the_markers_are_readable_after_the_conversion(self) -> None: + """Why Z-Image reads them after converting rather than before: unlike Krea-2's converter, + this one carries `.comfy_quant` onto the final module names, so no re-keying is needed.""" + prefix = "layers.0.attention" + sd = { + f"{prefix}.qkv.weight": torch.zeros(3 * 4, 256, dtype=torch.int8), + f"{prefix}.qkv.weight_scale": torch.ones(3 * 4, 1), + f"{prefix}.qkv.comfy_quant": _marker_blob(MARKER), + "x_embedder.weight": torch.zeros(2, 2), + } + markers = extract_int8_convrot_markers(_convert_z_image_gguf_to_diffusers(sd)) + assert set(markers) == {f"{prefix}.to_q", f"{prefix}.to_k", f"{prefix}.to_v"} + assert all(m == MARKER for m in markers.values()) + + def test_a_qkv_weight_that_does_not_divide_by_three_is_refused(self) -> None: + sd = {"layers.0.attention.qkv.weight": torch.zeros(7, 8), "x_embedder.weight": torch.zeros(2, 2)} + with pytest.raises(ValueError, match="not divisible by 3"): + _convert_z_image_gguf_to_diffusers(sd)