diff --git a/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx b/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx index 5656dfa8b69..c75db0e05e5 100644 --- a/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/02.model-families.mdx @@ -1,7 +1,7 @@ --- title: Model Formats and Families description: An overview of the models InvokeAI supports — checkpoint, Diffusers and quantized formats, the components that make up a model, and the model families with their licenses. -lastUpdated: 2026-09-30 +lastUpdated: 2026-10-11 sidebar: order: 2 --- @@ -48,7 +48,7 @@ InvokeAI recognizes the format from the file's contents, so you install them lik - **NF4** (bitsandbytes): 4-bit, CUDA only. - **SDNQ**: see [SDNQ Quantization](/configuration/optimization/sdnq-quantization/). - **int8** (`int8_convrot`, from ComfyUI repackages): the weights stay int8 in memory on every device, including Apple Silicon, and need no setting to turn on. Supported for FLUX.1, FLUX.2 Klein, Z-Image, Krea-2, Ideogram 4, MiniMax H3, LTX-2 and the PiD decoders, and for the Qwen3 8B, Qwen3-VL and LTX-2 Gemma encoders. -- **nvfp4** (4-bit, in ComfyUI's checkpoint format, such as Comfy-Org's builds; other nvfp4 exports such as AWQ are refused): the weights stay packed in memory and are unpacked per layer, so no special GPU is needed. Supported for the Z-Image, Krea-2, Qwen-Image and LTX-2 transformers and for the Qwen3, Qwen2.5-VL and FLUX.2 Mistral encoders; the Ideogram 4 nvfp4 files are refused at install. +- **nvfp4** (4-bit, in ComfyUI's checkpoint format, such as Comfy-Org's builds; other nvfp4 exports such as AWQ are refused): the weights stay packed in memory and are unpacked per layer, so no special GPU is needed. Supported for the Z-Image, Krea-2, Qwen-Image and LTX-2 transformers and for the Qwen3, Qwen3-VL 4B/8B, Qwen2.5-VL and FLUX.2 Mistral encoders; the Ideogram 4 nvfp4 transformer files are refused at install. - **fp8** (plain or `fp8_scaled`): InvokeAI keeps the weights in fp8 when [FP8 Storage](/configuration/optimization/fp8-storage/) is on, which installation does for you, or when `fp8_compute` is enabled. - **MXFP8** (fp8 with one scale per block of 32 weights): loaded by unpacking to BF16. Turn FP8 Storage off for these; see the [FP8 Storage caution](/configuration/optimization/fp8-storage/). diff --git a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx index f80f88c5c96..6854445d619 100644 --- a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx @@ -1,7 +1,7 @@ --- title: Model Format Support description: Which file formats and quantizations InvokeAI loads for every model family, text encoder, VAE and adapter. -lastUpdated: 2026-10-10 +lastUpdated: 2026-10-11 sidebar: order: 3 --- @@ -57,7 +57,7 @@ Ministral 3B encoder. Others are recognized by their layer names and fail, or mi | T5-XXL | FLUX.1, SD 3 | ✓ | ✗ | ✗ | ✗ | ✗ | ✓ | ✓ (and bitsandbytes int8) | | CLIP-L / CLIP-G | SD, SDXL, FLUX.1, SD 3 | ✓ | ✗ | ✗ | ✗ | ✗ | ✗ | CLIP-L inside a FLUX.1 SDNQ pipeline | | Qwen3 0.6B / 4B / 8B | Z-Image, FLUX.2 Klein, Anima | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | -| Qwen3-VL 4B / 8B | Krea-2, Ideogram 4 | ✓ | ✓ | ✓ | ✓ | ✗ | ✓ | ✗ | +| Qwen3-VL 4B / 8B | Krea-2, Ideogram 4 | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ | ✗ | | Qwen3-VL 32B (truncated) | MiniMax H3 | in the Diffusers folder | ✓ | L | ✓ | L | ✗ | ✗ | | Qwen3.5 4B | Anima 3.8B | ✗ | ✓ | L | L | L | ✗ | ✗ | | Qwen2.5-VL | Qwen-Image | ✓ | ✓ | ✓ | L | ✓ | ✗ | ✗ | @@ -68,7 +68,11 @@ Ministral 3B encoder. Others are recognized by their layer names and fail, or mi | UMT5 | Wan 2.2 | ✓ | ✗ | ✗ | ✗ | ✗ | ✗ | ✗ | The Qwen3-VL GGUF encoder is the language model only; a separate `mmproj` vision file is refused, and split -multi-part GGUFs have to be joined before installing. The Gemma-4 encoder installs only as a folder: its +multi-part GGUFs have to be joined before installing. Qwen3-VL nvfp4 single files keep their layers packed: ComfyUI's +`qwen3vl_8b_nvfp4`, and repacks that name their nvfp4 layers the same way, in `.comfy_quant` markers or the +`_quantization_metadata` header. A single file with nvfp4 layers named neither way, with AWQ layers, or with packed +weights missing nvfp4's per-tensor `weight_scale_2` fails to install with a message; quantized encoder folders (AWQ, +ModelOpt, llm-compressor) are not covered by this. The Gemma-4 encoder installs only as a folder: its `config.json` and tokenizer files next to exactly one weight file, bf16 or int8. A standalone `.safetensors` file is not recognized. diff --git a/docs/src/content/docs/Users Guide/03.Models/Local Models/ideogram4.mdx b/docs/src/content/docs/Users Guide/03.Models/Local Models/ideogram4.mdx index 788326294e8..24b0b8486a2 100644 --- a/docs/src/content/docs/Users Guide/03.Models/Local Models/ideogram4.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/Local Models/ideogram4.mdx @@ -1,7 +1,7 @@ --- title: Ideogram 4 description: Generate images with the Ideogram 4 text-to-image models. -lastUpdated: 2026-10-08 +lastUpdated: 2026-10-11 sidebar: order: 6 --- @@ -33,7 +33,7 @@ Comfy-Org's [`Comfy-Org/Ideogram-4`](https://huggingface.co/Comfy-Org/Ideogram-4 - `text_encoders/qwen3vl_8b_fp8_scaled.safetensors` — the Qwen3-VL **8B** encoder (not the 4B one Krea-2 uses). A llama.cpp GGUF of the 8B encoder works as well; see [Qwen3-VL GGUF encoders](/users-guide/models/local-models/krea-2/#qwen3-vl-gguf-encoders). - `vae/flux2-vae.safetensors` — the 32-channel VAE, shared with FLUX.2, which is why it installs under the FLUX.2 base -Pick the unconditional branch, the encoder and the VAE in the **Components** section next to the model — and pick the unconditional branch of the *same* build as the model. The `nvfp4` repackage in that repo is **not** supported yet and is refused at install. +Pick the unconditional branch, the encoder and the VAE in the **Components** section next to the model — and pick the unconditional branch of the *same* build as the model. The `nvfp4` transformers in that repo are **not** supported yet and are refused at install; its nvfp4 encoder, `text_encoders/qwen3vl_8b_nvfp4.safetensors`, loads and keeps its layers packed. Which branch a file holds is read from the file itself, and from the filename for a build that does not record it, so keep these names as published. diff --git a/docs/src/content/docs/Users Guide/03.Models/Local Models/krea-2.mdx b/docs/src/content/docs/Users Guide/03.Models/Local Models/krea-2.mdx index d87a444e5cd..6ea24e0571c 100644 --- a/docs/src/content/docs/Users Guide/03.Models/Local Models/krea-2.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/Local Models/krea-2.mdx @@ -1,7 +1,7 @@ --- title: Krea-2 description: Generate images with the Krea-2 text-to-image models (Turbo and Raw), including the GGUF / single-file workflow, the conditioning enhancers and training-free style reference. -lastUpdated: 2026-09-30 +lastUpdated: 2026-10-11 sidebar: order: 7 --- @@ -70,7 +70,8 @@ the bundled components. :::note[Qwen3-VL encoder] Single-file Qwen3-VL encoders bundle only the weights. Their tokenizer and architecture config ship with InvokeAI, so these encoders load without network access. A single-file fp8 encoder is kept -resident in fp8 (roughly half the VRAM of the bf16 encoder). +resident in fp8 (roughly half the VRAM of the bf16 encoder), and an nvfp4 encoder keeps its layers packed +at 4 bits. ::: ### Qwen3-VL GGUF encoders diff --git a/invokeai/backend/model_manager/configs/qwen3_vl_encoder.py b/invokeai/backend/model_manager/configs/qwen3_vl_encoder.py index 8d44eb46eb8..891772d7da3 100644 --- a/invokeai/backend/model_manager/configs/qwen3_vl_encoder.py +++ b/invokeai/backend/model_manager/configs/qwen3_vl_encoder.py @@ -15,6 +15,10 @@ ) from invokeai.backend.model_manager.model_on_disk import ModelOnDisk from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, Qwen3VLVariantType +from invokeai.backend.model_manager.util.qwen3_vl import drop_qwen3vl_unloaded_keys +from invokeai.backend.quantization.fp8_scaled import parse_quantization_metadata +from invokeai.backend.quantization.int8_convrot import read_comfy_quant_markers +from invokeai.backend.quantization.nvfp4 import WEIGHT_SCALE_2_SUFFIX, find_nvfp4_layers # The two Qwen3-VL encoders v7 can build, by language-model width. Both have 36 layers, so the # width alone identifies the variant. Anything else -- including MiniMax H3's truncated @@ -94,6 +98,25 @@ def _variant_from_checkpoint_shape(state_dict: dict[str | int, Any]) -> Qwen3VLV return variant +def _raise_for_undecodable_nvfp4_layers(mod: ModelOnDisk, state_dict: dict[str | int, Any]) -> None: + """Refuse at install time an nvfp4 layer the loader would refuse at the first generation. + + The loader's own check, over the tensors the loader reads -- so a visual tower or LM head it drops cannot + get a file refused. It turns away AWQ builds, packed weights without nvfp4's per-tensor global scale, + malformed layers, and layers neither a ComfyUI ``.comfy_quant`` marker nor the safetensors header names as + nvfp4 (ModelOpt-style exports among them). Identification holds meta tensors, whose markers have no bytes to + parse, so those are read from the file -- and only from a file that has an nvfp4 layer to name. + """ + consumed = drop_qwen3vl_unloaded_keys(state_dict) + names = parse_quantization_metadata(mod.metadata()) + if any(isinstance(key, str) and key.endswith(WEIGHT_SCALE_2_SUFFIX) for key in consumed): + names = {**names, **read_comfy_quant_markers(mod.path)} + try: + find_nvfp4_layers(consumed, header_layers=names) + except ValueError as e: + raise InvalidMatchError(f"this Qwen3-VL encoder cannot be loaded: {e}") from e + + class Qwen3VLEncoder_Qwen3VLEncoder_Config(Config_Base): """Configuration for standalone Qwen3-VL text encoder models (diffusers-like directory format). @@ -210,6 +233,9 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - state_dict = mod.load_state_dict() if not _is_qwen3_vl_encoder_state_dict(state_dict): raise NotAMatchError("state dict does not look like a single-file Qwen3-VL encoder") + # Ahead of the variant, whose width is read off the embedding: an undecodable repack that packs the + # embedding too would fail there instead, and register as an unknown model with the reason only in the log. + _raise_for_undecodable_nvfp4_layers(mod, state_dict) variant = override_fields.pop("variant", None) or _variant_from_checkpoint_shape(state_dict) return cls(**override_fields, variant=variant) diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index 9f4a9a82bab..be23d63ee2f 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/invokeai/backend/model_manager/load/model_loaders/krea2.py @@ -34,8 +34,8 @@ ) from invokeai.backend.model_manager.util.llamacpp_keys import convert_llamacpp_decoder_keys from invokeai.backend.model_manager.util.qwen3_vl import ( + drop_qwen3vl_unloaded_keys, drop_qwen3vl_visual_tower, - drop_qwen3vl_visual_tower_keys, normalize_qwen3vl_rope_config, qwen3vl_target_key, ) @@ -43,6 +43,7 @@ INPUT_SCALE_SUFFIXES, attach_fp8_scales, cast_state_dict, + count_fp8_weights, dequantize_fp8_scaled, detach_layer_sidechannel, extract_comfy_quant_hints, @@ -866,9 +867,10 @@ def _load_te_config(self, config: _Qwen3VLSingleFileConfig) -> Any: class Qwen3VLEncoderCheckpointLoader(_Qwen3VLEncoderSingleFileLoader[Qwen3VLEncoder_Checkpoint_Config]): """Loads a single-file Qwen3-VL encoder checkpoint (e.g. ComfyUI ``qwen3vl_4b_bf16`` / ``_fp8_scaled``). - The checkpoint bundles the language model + visual tower; the tower is dropped on the way in (see - ``drop_qwen3vl_visual_tower_keys``). ComfyUI 'scaled fp8' weights are dequantized to the compute - dtype on load. + The checkpoint bundles the language model + visual tower; the tower is dropped on the way in, with the + LM head (see ``drop_qwen3vl_unloaded_keys``). ComfyUI 'scaled fp8' weights are kept for the fp8 matmul or fp8 + storage where the device has either, and folded into the compute dtype otherwise. ``int8_tensorwise`` + layers stay int8 and Comfy's ``nvfp4`` layers stay packed, both dequantized per forward. """ CONFIG_CLASS = Qwen3VLEncoder_Checkpoint_Config @@ -886,7 +888,19 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod # Ahead of every pass that reads tensors, so the tower is never dequantized, cast, copied or # reserved for. (The file *header* is still read below for quantization metadata, which can # therefore still name visual layers; those hints match no surviving key and are ignored.) - sd = drop_qwen3vl_visual_tower_keys(sd) + # The LM head goes with it: this model has none, and a packed one would name no module. + sd = drop_qwen3vl_unloaded_keys(sd) + metadata = read_safetensors_metadata(model_path, self._logger) + + # Comfy's nvfp4 build packs every language-model projection and names the layers in `.comfy_quant` + # markers (a header in some repacks). Take them out before either side-channel branch reads a scale: + # the int8 branch strips every remaining marker -- one of the two things that name an nvfp4 layer -- + # and the fp8 extraction pops every `weight_scale` and discards those whose weight is not float8, + # nvfp4's block scales included. `install_nvfp4_layers` puts them back, packed, under the paths the + # remap gives their weights. + nvfp4_payloads = pop_nvfp4_layers(sd, header_layers=parse_quantization_metadata(metadata)) + nvfp4_payloads = _remap_qwen3vl_singlefile_keys(nvfp4_payloads) + # Same one-or-the-other split as the transformer above, and here it also guards the fp8 # *detection*: an int8 layer ships a `.weight_scale` too, so probing for scales without # ruling out int8 first would keep the encoder "fp8-resident" over weights that were never @@ -913,7 +927,6 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod # only what this path already left dense -- work with no saving on the layers that matter. use_fp8_storage = False else: - 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)} @@ -942,19 +955,22 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod # for the byte count the file already had. use_fp8_storage = source_is_fp8 and _device_supports_fp8_storage(self._torch_device, self._logger) + # Every layer that arrived scaled, before the split folds the ones the matmul cannot take. The storage + # pass below must leave those alone: cast again, they would be re-quantized *without* their scale. + scaled_paths = set(fp8_layers) + te_config = self._load_te_config(config) with accelerate.init_empty_weights(): model = Qwen3VLModel._from_config(te_config) # Its weights were dropped from the state dict above, so the module has to go too -- left in # place it would stay on the meta device and `_reject_incomplete_load` would rightly refuse. drop_qwen3vl_visual_tower(model) + # `Qwen3VLModel` declares no precision-sensitive modules, but read them rather than assume: in the + # int8 split they are also what keeps a marker on a non-Linear (a "quantize everything" repack's 1-D + # norms) out of `swap_in_int8_linears`, and the nvfp4 install decodes rather than packs under them. + skip_patterns = _model_declared_skip_patterns(model) if int8_markers: - # `Qwen3VLModel` declares no precision-sensitive modules, but read them rather than - # assume: the split is also what keeps a marker on a non-Linear (a "quantize - # everything" repack's 1-D norms) out of `swap_in_int8_linears`. - skip_patterns = _model_declared_skip_patterns(model) - quantized = install_int8_convrot_layers( model, sd, @@ -963,6 +979,9 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod architecture="Qwen3-VL encoder", reserve=self._ram_cache.make_room, skip_patterns=skip_patterns, + # The nvfp4 layers taken out above are held beside the int8 ones, and this is the + # branch's only reservation. + extra_reserved_bytes=predict_nvfp4_install_size(model, nvfp4_payloads, model_dtype, skip_patterns), ) self._logger.info( f"Qwen3-VL encoder: kept {len(quantized)} of {len(int8_markers)} layer(s) in int8 " @@ -999,21 +1018,27 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod keep_fp8=keep_fp8, model=model, fp8_layers=fp8_layers, - nvfp4_payloads={}, + nvfp4_payloads=nvfp4_payloads, + skip_patterns=skip_patterns, ) if fp8_layers and not keep_fp8: # Neither consumer wants them packed: fold the scales into the weights. Both consumers # were resolved above, next to the matmul probe, because this fold is irreversible. dequantize_fp8_scaled(sd, fp8_layers, model_dtype) fp8_layers = {} - fp8_layers = split_fp8_scaled_layers(sd, fp8_layers, model_dtype, model=model) - # No `skip_patterns` here on purpose: this model declares none, and the storage pass below - # applies `_FP8_DEFAULT_SKIP_PATTERNS` itself. Those two lists used to have to not intersect - # on a layer class with no fp8-capable wrapper -- a `pos_embed` or a `patch_embed.proj` - # would arrive fp8 from the state dict, be skipped by the cast pass, and then raise on the - # first forward. `_apply_fp8_to_nn_module` now restores the compute dtype on the modules it - # skips, so the two lists are independent again. - cast_state_dict(sd, model_dtype, keep_fp8=keep_fp8, model=model) + # The model's own patterns (none, today) go to the reservation, the split, the cast and the nvfp4 + # install alike, so the prediction answers for what each of them then does. They are independent + # of `_FP8_DEFAULT_SKIP_PATTERNS`, which the storage pass below applies itself: a `pos_embed` or + # `patch_embed.proj` arriving fp8 and skipped there used to raise on the first forward, and + # `_apply_fp8_to_nn_module` now restores the compute dtype on the modules it skips. + fp8_layers = split_fp8_scaled_layers(sd, fp8_layers, model_dtype, model=model, skip_patterns=skip_patterns) + cast_state_dict(sd, model_dtype, keep_fp8=keep_fp8, model=model, skip_patterns=skip_patterns) + + if nvfp4_payloads: + packed = install_nvfp4_layers(model, sd, nvfp4_payloads, model_dtype, skip_patterns) + self._logger.info( + f"Qwen3-VL encoder '{config.name}': kept {packed} of {len(nvfp4_payloads)} nvfp4 layer(s) packed." + ) load_state_dict_ignoring_extras( model, sd, source="Qwen3-VL encoder checkpoint", assign=True, allow_missing=True @@ -1074,7 +1099,22 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod 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) + # + # A scaled layer reaches this branch only when the split folded it -- every one it had, or the + # attach branch above would have run. The Linears among them go to fp8 storage like any other, + # which is what this pass is for: an MXFP8 build's block scales are folded, and its language model + # is held in fp8 rather than at twice the size. A folded embedding is not: Comfy's nvfp4 8B build + # scales its token table, and cast again it would lose that scale for the accuracy cost measured + # above for exactly this table -- the reason the attach branch leaves every embedding alone. + self._apply_fp8_to_nn_module( + model, + storage_dtype=torch.float8_e4m3fn, + compute_dtype=model_dtype, + skip=lambda name, module: name in scaled_paths and not isinstance(module, torch.nn.Linear), + ) + # Counted rather than assumed: on that nvfp4 build every Linear is packed and the embedding skipped, so + # the pass leaves nothing in fp8, and saying it did would misdescribe the encoder. + if use_fp8_storage and count_fp8_weights(model): self._logger.info( f"FP8 layerwise casting enabled for Qwen3-VL encoder '{config.name}' " f"(storage=float8_e4m3fn, compute={model_dtype})." diff --git a/invokeai/backend/model_manager/util/qwen3_vl.py b/invokeai/backend/model_manager/util/qwen3_vl.py index c8dd8d652a5..709e5d9af7a 100644 --- a/invokeai/backend/model_manager/util/qwen3_vl.py +++ b/invokeai/backend/model_manager/util/qwen3_vl.py @@ -103,3 +103,18 @@ def drop_qwen3vl_visual_tower_keys(sd: dict[str, Any]) -> dict[str, Any]: and ``comfy_quant``). Enumerating suffixes instead would silently start leaving orphans behind. """ return {k: v for k, v in sd.items() if not (isinstance(k, str) and qwen3vl_target_key(k).startswith("visual."))} + + +def drop_qwen3vl_unloaded_keys(sd: dict[str, Any]) -> dict[str, Any]: + """Keep only what the single-file encoder loader reads: drop the visual tower and the LM head. + + ``Qwen3VLModel`` has no LM head, so a non-strict load would discard it anyway -- but only after the + side-channel handling had folded and reserved for it: on the 8B fp8 builds a 151936x4096 bf16 copy. + Identification judges a file's quantization by the same view, so it refuses nothing the loader would + never read. + """ + return { + k: v + for k, v in drop_qwen3vl_visual_tower_keys(sd).items() + if not (isinstance(k, str) and k.startswith("lm_head.")) + } diff --git a/invokeai/backend/quantization/fp8_scaled.py b/invokeai/backend/quantization/fp8_scaled.py index 46e0d8d46e0..349feb3503d 100644 --- a/invokeai/backend/quantization/fp8_scaled.py +++ b/invokeai/backend/quantization/fp8_scaled.py @@ -651,11 +651,34 @@ def dequantize_fp8_scaled( weight = sd.get(key) if weight is None: continue - weight = weight.float() - sd[key] = (weight * expand_weight_scale(weight, layer.weight_scale, path)).to(dtype) + sd[key] = _fold_scale(weight, expand_weight_scale(weight, layer.weight_scale, path), dtype) return sd +# The most elements a fold widens to float32 at once. The multiply holds two float32 copies of what it works +# on, so folding a tensor whole peaks at 10 bytes per element in RAM: 5.8 GiB for the 151936x4096 token +# table of Comfy's nvfp4 Qwen3-VL 8B, against a 1.2 GiB result and a reservation sized for that result. +_FOLD_SLAB_ELEMENTS = 1 << 24 + + +def _fold_scale(weight: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """``(weight.float() * scale).to(dtype)``, element for element, a slab of rows at a time. + + ``scale`` is already expanded against ``weight``: a 0-d scalar, or a tensor whose leading axis runs over + the weight's rows. Each slab rounds exactly as the whole would, since every element is computed alone. + """ + if weight.numel() <= _FOLD_SLAB_ELEMENTS: + return (weight.float() * scale).to(dtype) + rows = weight.shape[0] + step = max(1, _FOLD_SLAB_ELEMENTS // (weight.numel() // rows)) + per_row = scale.dim() > 0 and scale.shape[0] == rows + out = torch.empty(weight.shape, dtype=dtype, device=weight.device) + for start in range(0, rows, step): + slab = slice(start, start + step) + out[slab] = weight[slab].float() * (scale[slab] if per_row else scale) + return out + + def attach_fp8_scales( model: torch.nn.Module, layers: Mapping[str, Fp8ScaledLayer], diff --git a/invokeai/backend/quantization/nvfp4.py b/invokeai/backend/quantization/nvfp4.py index 743bd76c944..2b522b058ae 100644 --- a/invokeai/backend/quantization/nvfp4.py +++ b/invokeai/backend/quantization/nvfp4.py @@ -193,7 +193,7 @@ def nbytes(self) -> int: return sum(t.nelement() * t.element_size() for t in (self.weight, self.weight_scale, self.weight_scale_2)) -def _find_nvfp4_layers(sd: Mapping[str, Any], header_layers: Mapping[str, Any] | None) -> list[str]: +def find_nvfp4_layers(sd: Mapping[str, Any], header_layers: Mapping[str, Any] | None) -> list[str]: """Every nvfp4 layer in ``sd``, after refusing what this module would read wrong. Mutates nothing. Detection is structural -- a ``weight_scale_2`` beside the weight -- and decided per layer, never per @@ -201,6 +201,10 @@ def _find_nvfp4_layers(sd: Mapping[str, Any], header_layers: Mapping[str, Any] | with a ``weight_scale`` but no ``weight_scale_2``; malformed layers; and layers that neither a ``.comfy_quant`` marker nor a ``_quantization_metadata`` entry (``header_layers``, in ``sd``'s key space) names as nvfp4. + + Reads dtypes and shapes only, so identification can run it over meta tensors and refuse at install time + what a load would refuse. A meta marker has no bytes to parse, so there the markers' entries have to come + in through ``header_layers`` (see ``read_comfy_quant_markers``). """ layers = sorted( k[: -len(WEIGHT_SCALE_2_SUFFIX)] for k in sd if isinstance(k, str) and k.endswith(WEIGHT_SCALE_2_SUFFIX) @@ -285,7 +289,7 @@ def reject_nvfp4_layers_a_plain_fold_cannot_decode(sd: Mapping[str, Any], what: Detection is the union of two structural tests: a ``weight_scale_2`` beside the weight, which is what the decode keys on, and a packed ``uint8`` weight carrying a block scale without one, which - is the half-state :func:`_find_nvfp4_layers` refuses by name. Keying on the first alone would + is the half-state :func:`find_nvfp4_layers` refuses by name. Keying on the first alone would miss the second, and the fold takes it just as readily. Narrower than the decode in one way worth knowing: at the Wan seam a bundled nvfp4 text encoder @@ -299,7 +303,7 @@ def reject_nvfp4_layers_a_plain_fold_cannot_decode(sd: Mapping[str, Any], what: and key.endswith(WEIGHT_SCALE_2_SUFFIX) and f"{key[: -len(WEIGHT_SCALE_2_SUFFIX)]}.weight" in sd } - # A packed weight with a block scale and no global one is the half-state `_find_nvfp4_layers` + # A packed weight with a block scale and no global one is the half-state `find_nvfp4_layers` # refuses by name. Keying on `weight_scale_2` alone would miss it, and the fold takes it just as # readily -- so the detection here is the union, not the narrower test. half = { @@ -326,7 +330,7 @@ def pop_nvfp4_layers(sd: dict[str, Any], header_layers: Mapping[str, Any] | None is touched. """ payloads: dict[str, NVFP4Payload] = {} - for path in _find_nvfp4_layers(sd, header_layers): + for path in find_nvfp4_layers(sd, header_layers): payloads[path] = NVFP4Payload( weight=sd.pop(f"{path}.weight"), weight_scale=sd.pop(f"{path}.weight_scale"), diff --git a/tests/backend/model_manager/configs/test_qwen3_vl_encoder_config.py b/tests/backend/model_manager/configs/test_qwen3_vl_encoder_config.py index 456f8273dee..24bbb27bdd5 100644 --- a/tests/backend/model_manager/configs/test_qwen3_vl_encoder_config.py +++ b/tests/backend/model_manager/configs/test_qwen3_vl_encoder_config.py @@ -10,8 +10,10 @@ from unittest.mock import MagicMock, patch import pytest +import torch +from safetensors.torch import save_file -from invokeai.backend.model_manager.configs.identification_utils import NotAMatchError +from invokeai.backend.model_manager.configs.identification_utils import InvalidMatchError, NotAMatchError from invokeai.backend.model_manager.configs.qwen3_vl_encoder import ( Qwen3VLEncoder_Checkpoint_Config, Qwen3VLEncoder_Qwen3VLEncoder_Config, @@ -19,6 +21,7 @@ ) from invokeai.backend.model_manager.model_on_disk import ModelOnDisk from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, Qwen3VLVariantType +from tests.fixtures.quantized_payloads import comfy_quant_marker, nvfp4_signed_tensors, quantize_scaled_fp8 _REQUIRED_FIELDS = { "hash": "blake3:fakehash", @@ -165,6 +168,100 @@ def test_rejects_a_truncated_language_stack(self, _rfo, _rif) -> None: Qwen3VLEncoder_Checkpoint_Config.from_model_on_disk(mod, {**_REQUIRED_FIELDS}) +_PACKED = "model.layers.0.self_attn.q_proj" + + +class TestQwen3VLEncoderQuantizedCheckpoints: + """What the loader would refuse of a quantized single file is refused at install time; the rest installs. + + Written as real files and read through `ModelOnDisk`, because identification sees tensors on the meta + device -- dtypes and shapes, no bytes -- so whatever names a layer nvfp4 has to be read from the file. + """ + + @classmethod + def _identify( + cls, tmp_path: Path, mutate=lambda _sd: None, *, naming: str = "marker" + ) -> Qwen3VLEncoder_Checkpoint_Config: + packed, _ = nvfp4_signed_tensors(_PACKED, torch.ones(128, 64, dtype=torch.bool)) + state_dict = { + "model.embed_tokens.weight": torch.zeros(2, 2560, dtype=torch.bfloat16), + "model.layers.35.input_layernorm.weight": torch.ones(4, dtype=torch.bfloat16), + "model.visual.blocks.0.attn.qkv.weight": torch.zeros(4, 4, dtype=torch.bfloat16), + **packed, + } + metadata = None + if naming == "marker": + state_dict[f"{_PACKED}.comfy_quant"] = comfy_quant_marker({"format": "nvfp4"}) + else: + metadata = {"_quantization_metadata": json.dumps({"layers": {_PACKED: {"format": "nvfp4"}}})} + mutate(state_dict) + path = tmp_path / "qwen3vl_4b.safetensors" + save_file(state_dict, str(path), metadata=metadata) + return Qwen3VLEncoder_Checkpoint_Config.from_model_on_disk(ModelOnDisk(path), {**_REQUIRED_FIELDS}) + + @pytest.mark.parametrize("naming", ["marker", "header"]) + def test_a_comfy_nvfp4_build_installs(self, tmp_path: Path, naming: str) -> None: + assert self._identify(tmp_path, naming=naming).variant is Qwen3VLVariantType.Qwen3VL_4B + + def _replace_packed_layer(self, state_dict: dict, **tensors: torch.Tensor) -> None: + for key in [key for key in state_dict if key.startswith(f"{_PACKED}.")]: + del state_dict[key] + state_dict.update({f"{_PACKED}.{suffix}": tensor for suffix, tensor in tensors.items()}) + + def test_scaled_fp8_and_int8_builds_still_install(self, tmp_path: Path) -> None: + # Both pair a weight with a `weight_scale` too, and neither is 4-bit. + fp8 = quantize_scaled_fp8(torch.randn(128, 64)) + assert ( + self._identify( + tmp_path, lambda sd: self._replace_packed_layer(sd, weight=fp8.codes, weight_scale=fp8.scale) + ).variant + is Qwen3VLVariantType.Qwen3VL_4B + ) + int8 = { + "weight": torch.randint(-127, 128, (128, 64), dtype=torch.int8), + "weight_scale": torch.ones(128, 1), + "comfy_quant": comfy_quant_marker({"format": "int8_tensorwise"}), + } + assert ( + self._identify(tmp_path, lambda sd: self._replace_packed_layer(sd, **int8)).variant + is Qwen3VLVariantType.Qwen3VL_4B + ) + + @pytest.mark.parametrize( + ("mutate", "message"), + [ + # ModelOpt-style: the nvfp4 tensors, but nothing naming the layout ComfyUI's. + pytest.param(lambda sd: sd.pop(f"{_PACKED}.comfy_quant"), "no ComfyUI `comfy_quant` marker", id="unnamed"), + # A custom RTN repack published for this encoder: packed codes and a block scale only. + pytest.param(lambda sd: sd.pop(f"{_PACKED}.weight_scale_2"), "no weight_scale_2", id="no-global-scale"), + pytest.param(lambda sd: sd.update({f"{_PACKED}.pre_quant_scale": torch.ones(64)}), "AWQ", id="awq"), + ], + ) + def test_a_layer_the_loader_would_refuse_is_refused_at_install(self, tmp_path: Path, mutate, message: str) -> None: + with pytest.raises(InvalidMatchError, match=message): + self._identify(tmp_path, mutate) + + def test_a_repack_packing_its_embedding_too_is_refused_rather_than_unknown(self, tmp_path: Path) -> None: + # The variant is read off the embedding's width, which a packed table halves. Checked after that, such a + # file would fail there instead and register as an unknown model, its reason only in the log. + def packed_embedding_without_global_scale(state_dict: dict) -> None: + state_dict["model.embed_tokens.weight"] = torch.zeros(2, 1280, dtype=torch.uint8) + state_dict["model.embed_tokens.weight_scale"] = torch.zeros(2, 160).to(torch.float8_e4m3fn) + + with pytest.raises(InvalidMatchError, match="no weight_scale_2"): + self._identify(tmp_path, packed_embedding_without_global_scale) + + def test_a_layer_the_loader_never_reads_cannot_get_a_file_refused(self, tmp_path: Path) -> None: + # The loader drops the visual tower and the LM head before it reads a scale, so their layout is not this + # file's to answer for. + def unread_half_state_layers(state_dict: dict) -> None: + for path in ("model.visual.blocks.0.mlp.linear_fc1", "lm_head"): + state_dict[f"{path}.weight"] = torch.zeros(128, 32, dtype=torch.uint8) + state_dict[f"{path}.weight_scale"] = torch.zeros(128, 4).to(torch.float8_e4m3fn) + + assert self._identify(tmp_path, unread_half_state_layers).variant is Qwen3VLVariantType.Qwen3VL_4B + + class TestQwen3VLEncoderDirectoryConfig: @staticmethod def _write_config(path: Path, *, hidden_size: int = 2560, num_hidden_layers: int = 36) -> None: 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 29082594df3..1659b5cbeb5 100644 --- a/tests/backend/model_manager/load/test_krea2_loader_boundaries.py +++ b/tests/backend/model_manager/load/test_krea2_loader_boundaries.py @@ -926,7 +926,7 @@ def test_an_nvfp4_layer_missing_its_global_scale_is_refused_before_the_cache_is_ A packed uint8 weight with a block-scale grid and no `weight_scale_2` is the state a guard keyed on `weight_scale_2` -- the key the decode itself keys on -- lets straight through. - `_find_nvfp4_layers` refuses it and `test_nvfp4.py` pins that; what only a seam can answer is + `find_nvfp4_layers` refuses it and `test_nvfp4.py` pins that; what only a seam can answer is whether this loader reaches the detector before it asks the cache for room. Krea-2 pops on the *native* keys, before the conversion that remaps the payload paths, so the refusal has to survive being upstream of the rename as well as upstream of the reservation. diff --git a/tests/backend/model_manager/load/test_ltx2_loader.py b/tests/backend/model_manager/load/test_ltx2_loader.py index 5604ade400d..93c59d32432 100644 --- a/tests/backend/model_manager/load/test_ltx2_loader.py +++ b/tests/backend/model_manager/load/test_ltx2_loader.py @@ -582,7 +582,7 @@ def test_an_nvfp4_layer_missing_its_global_scale_is_refused_before_the_cache_is_ """The degraded half-state, at this seam rather than at the detector. A packed uint8 weight with a block-scale grid and no `weight_scale_2` is what a guard keyed on - `weight_scale_2` -- the key the decode keys on -- lets straight through. `_find_nvfp4_layers` + `weight_scale_2` -- the key the decode keys on -- lets straight through. `find_nvfp4_layers` refuses it and `test_nvfp4.py` pins that; what only a seam can answer is whether this loader reaches the detector before it asks the cache for room. Here a key conversion stands between the pop and the reservation, and nothing but their order keeps them that way round. diff --git a/tests/backend/model_manager/load/test_mistral_encoder_checkpoint_loader.py b/tests/backend/model_manager/load/test_mistral_encoder_checkpoint_loader.py index 7cd8662fbdf..fe8dc505a72 100644 --- a/tests/backend/model_manager/load/test_mistral_encoder_checkpoint_loader.py +++ b/tests/backend/model_manager/load/test_mistral_encoder_checkpoint_loader.py @@ -304,7 +304,7 @@ def test_an_nvfp4_layer_missing_its_global_scale_is_refused_before_the_cache_is_ """The degraded half-state, at this seam rather than at the detector. A packed uint8 weight with a block-scale grid and no `weight_scale_2` is what a guard keyed on - `weight_scale_2` -- the key the decode itself keys on -- lets straight through. `_find_nvfp4_layers` + `weight_scale_2` -- the key the decode itself keys on -- lets straight through. `find_nvfp4_layers` refuses it and `test_nvfp4.py` pins that; what only a seam can answer is whether this loader still reaches the detector before it asks the cache for room, and on both fp8 branches. Nothing but the order of two statements keeps `pop_nvfp4_layers` above the reservation here. diff --git a/tests/backend/model_manager/load/test_qwen3_encoder_checkpoint_loader.py b/tests/backend/model_manager/load/test_qwen3_encoder_checkpoint_loader.py index 9982f04079a..64d30d43e27 100644 --- a/tests/backend/model_manager/load/test_qwen3_encoder_checkpoint_loader.py +++ b/tests/backend/model_manager/load/test_qwen3_encoder_checkpoint_loader.py @@ -259,7 +259,7 @@ def test_an_nvfp4_layer_missing_its_global_scale_is_refused_before_the_cache_is_ """The degraded nvfp4 half-state, at this seam rather than at the detector. A packed uint8 weight with a block-scale grid and no `weight_scale_2` is the state a guard keyed - on `weight_scale_2` -- the key the decode keys on -- lets straight through. `_find_nvfp4_layers` + on `weight_scale_2` -- the key the decode keys on -- lets straight through. `find_nvfp4_layers` refuses it and `test_nvfp4.py` pins that; what only a seam can answer is whether this loader reaches the detector before it asks the cache for room. Asserted through the cache double this file already uses, in the same idiom as the cells above, because this loader reads its header diff --git a/tests/backend/model_manager/load/test_qwen3_vl_encoder_nvfp4_loader.py b/tests/backend/model_manager/load/test_qwen3_vl_encoder_nvfp4_loader.py new file mode 100644 index 00000000000..0a825ec1133 --- /dev/null +++ b/tests/backend/model_manager/load/test_qwen3_vl_encoder_nvfp4_loader.py @@ -0,0 +1,309 @@ +"""The single-file Qwen3-VL encoder loader keeps Comfy's nvfp4 builds packed. + +`qwen3vl_8b_nvfp4.safetensors` (Comfy-Org, Ideogram 4) packs all seven projections of every layer, names them only in +`.comfy_quant` markers, and keeps the token embedding and the LM head in scaled fp8. Community 4B repacks follow +either that marker convention or name their layers in the safetensors header, one of the latter packing the visual +tower too. Both reach the loader shared by Krea-2 and Ideogram 4. What these tests pin is the order: the packed +layers leave the state dict before either side-channel branch reads a scale, come back packed under the transformers +paths of a real `Qwen3VLModel`, are counted in the one reservation as they are held, and a folded scaled embedding is +not re-quantized without its scale -- while the Linears an MXFP8 build folds still go to fp8 storage as before. The +compute dtype is bf16, as in production: the packed global scale must not be cast with it. +""" + +import json + +import pytest +import safetensors.torch +import torch +from transformers import Qwen3VLConfig, Qwen3VLModel + +from invokeai.backend.model_manager.configs.qwen3_vl_encoder import Qwen3VLEncoder_Checkpoint_Config +from invokeai.backend.model_manager.load.model_loaders import krea2 +from invokeai.backend.model_manager.load.model_loaders.krea2 import Qwen3VLEncoderCheckpointLoader +from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear +from invokeai.backend.quantization.nvfp4 import NVFP4Linear +from tests.fixtures.loader_seams import Seam, prepare +from tests.fixtures.quantized_payloads import ( + MX_BLOCK_SIZE, + comfy_quant_marker, + mxfp8_marker, + mxfp8_tensors, + nvfp4_signed_tensors, + quantize_scaled_fp8, +) + +COMPUTE_DTYPE = torch.bfloat16 +HIDDEN = 128 +INTERMEDIATE = 256 +VOCAB = 32 +HEAD_DIM = 32 + +# Every projection, as the official build packs them: (layer path in the file, [out, in]). The widths are whole +# 128-row scale tiles, the smallest an nvfp4 layer can be. +PROJECTIONS = { + "model.layers.0.self_attn.q_proj": (HIDDEN, HIDDEN), + "model.layers.0.self_attn.k_proj": (HIDDEN, HIDDEN), + "model.layers.0.self_attn.v_proj": (HIDDEN, HIDDEN), + "model.layers.0.self_attn.o_proj": (HIDDEN, HIDDEN), + "model.layers.0.mlp.gate_proj": (INTERMEDIATE, HIDDEN), + "model.layers.0.mlp.up_proj": (INTERMEDIATE, HIDDEN), + "model.layers.0.mlp.down_proj": (HIDDEN, INTERMEDIATE), +} +# A packed visual layer, as the header-named 4B repack ships one. The tower is dropped before anything reads it. +VISUAL = "model.visual.blocks.0.mlp.linear_fc1" +# The bf16 the folded embedding and the norms (two per layer, q/k norms, the final one) are held in. +DENSE_ELEMENTS = VOCAB * HIDDEN + 2 * HIDDEN + 2 * HEAD_DIM + HIDDEN + +SEAM = Seam( + loader=Qwen3VLEncoderCheckpointLoader, + module=krea2, + entry="_load_text_encoder", + # The loader imports `load_file` inside the method, so the name it resolves is the package's. + load_file_host=safetensors.torch, + compute_dtype=COMPUTE_DTYPE, + patches_device=True, + sets_torch_dtype=False, + # This entry calls `_apply_fp8_to_nn_module` itself, and the storage pass is part of what is tested. + casts_fp8_storage=False, +) + + +def _te_config() -> Qwen3VLConfig: + """Qwen3-VL at toy width: one text layer whose projections all fit nvfp4's tiles, and a one-block tower.""" + return Qwen3VLConfig( + text_config={ + "hidden_size": HIDDEN, + "intermediate_size": INTERMEDIATE, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": HEAD_DIM, + "vocab_size": VOCAB, + "rope_scaling": {"rope_type": "default", "mrope_section": [8, 4, 4], "mrope_interleaved": True}, + }, + vision_config={ + "depth": 1, + "hidden_size": 64, + "intermediate_size": 128, + "num_heads": 2, + "out_hidden_size": HIDDEN, + "deepstack_visual_indexes": [], + }, + ) + + +def _transformers_path(path: str) -> str: + """Where `Qwen3VLModel` holds a language-model layer the file spells `model.X`.""" + return path.replace("model.", "language_model.", 1) + + +def _packed_bytes(rows: int, columns: int) -> int: + return rows * columns // 2 + rows * (columns // 16) + 4 + + +def _dense_checkpoint() -> dict[str, torch.Tensor]: + """The toy model's own state dict in bf16, spelled the way a Comfy export spells it. + + Starting from the real model is what makes every key the loader has to fill be there, and the norms carry + real values. + """ + torch.manual_seed(0) + reference = Qwen3VLModel(_te_config()).state_dict() + return { + f"model.{key.removeprefix('language_model.')}": value.to(torch.bfloat16) for key, value in reference.items() + } + + +def _checkpoint(naming: str) -> tuple[dict[str, torch.Tensor], dict[str, str] | None, dict[str, torch.Tensor]]: + """The official build's layout at toy width: (state dict, header, expected weight per transformers path).""" + state_dict = _dense_checkpoint() + header: dict[str, dict[str, str]] = {} + expected: dict[str, torch.Tensor] = {} + for path, (rows, columns) in PROJECTIONS.items(): + del state_dict[f"{path}.weight"] + tensors, weight = nvfp4_signed_tensors(path, torch.randint(0, 2, (rows, columns), dtype=torch.bool)) + state_dict.update(tensors) + expected[_transformers_path(path)] = weight + if naming == "marker": + state_dict[f"{path}.comfy_quant"] = comfy_quant_marker({"format": "nvfp4"}) + else: + header[path] = {"format": "nvfp4"} + if naming == "header": + del state_dict[f"{VISUAL}.weight"] + tensors, _ = nvfp4_signed_tensors(VISUAL, torch.ones(128, 64, dtype=torch.bool)) + state_dict.update(tensors) + header[VISUAL] = {"format": "nvfp4"} + + # The official build's scaled fp8 embedding and LM head. + for path, payload in ( + ("model.embed_tokens", quantize_scaled_fp8(torch.randn(VOCAB, HIDDEN))), + ("lm_head", quantize_scaled_fp8(torch.randn(VOCAB, HIDDEN))), + ): + state_dict[f"{path}.weight"] = payload.codes + state_dict[f"{path}.weight_scale"] = payload.scale + state_dict[f"{path}.comfy_quant"] = comfy_quant_marker( + {"format": "float8_e4m3fn", "full_precision_matrix_mult": False} + ) + expected[f"{path}.dequantized"] = payload.dequantized + metadata = {"_quantization_metadata": json.dumps({"layers": header})} if header else None + return state_dict, metadata, expected + + +def _prepare(monkeypatch: pytest.MonkeyPatch, state_dict, metadata, *, matmul: bool = False, storage: bool = False): + run = prepare( + SEAM, + monkeypatch, + state_dict=state_dict, + metadata=metadata, + observe=("dequantize_fp8_scaled", "split_fp8_scaled_layers"), + ) + run.loader._load_te_config = lambda _config: _te_config() + monkeypatch.setattr(krea2, "should_keep_fp8_weights", lambda _device: matmul) + monkeypatch.setattr(krea2, "_device_supports_fp8_storage", lambda *_args: storage) + return run + + +def _config() -> Qwen3VLEncoder_Checkpoint_Config: + return Qwen3VLEncoder_Checkpoint_Config.model_construct(path="qwen3vl_8b_nvfp4.safetensors", name="encoder") + + +def _logged(run) -> list[str]: + return [str(call.args[0]) for call in run.loader._logger.info.call_args_list] + + +def _assert_packed(model: torch.nn.Module, path: str, weight: torch.Tensor) -> None: + module = model.get_submodule(_transformers_path(path)) + assert isinstance(module, NVFP4Linear), path + assert module.weight.dtype is torch.uint8, path + assert module.weight_scale_2.dtype is torch.float32, path + # A float32 activation, not the compute dtype: CPU bf16 GEMM faults on part of the windows runner fleet. + # Every decoded weight is +-0.5, exact in either width. + x = torch.randn(3, module.in_features, dtype=torch.float32) + torch.testing.assert_close(module(x), torch.nn.functional.linear(x, weight)) + + +@pytest.mark.parametrize("mode", ["fold", "fp8_compute", "fp8_storage"]) +@pytest.mark.parametrize("naming", ["marker", "header"]) +def test_named_nvfp4_layers_stay_packed_beside_a_folded_embedding( + monkeypatch: pytest.MonkeyPatch, naming: str, mode: str +) -> None: + state_dict, metadata, expected = _checkpoint(naming) + run = _prepare(monkeypatch, state_dict, metadata, matmul=mode == "fp8_compute", storage=mode == "fp8_storage") + + model = run.load(_config()) + + for path in PROJECTIONS: + _assert_packed(model, path, expected[_transformers_path(path)]) + + # Comfy's embedding is scaled fp8, and no matmul takes an embedding: it is folded with its scale. Under fp8 + # storage it must stay folded -- cast again it would be re-quantized without the scale. With every Linear + # packed nothing else is left to cast, and the log must not claim otherwise. + embedding = model.language_model.embed_tokens.weight + assert embedding.dtype is COMPUTE_DTYPE + assert torch.equal(embedding, expected["model.embed_tokens.dequantized"].to(COMPUTE_DTYPE)) + assert not any("FP8 layerwise casting enabled" in message for message in _logged(run)) + + # One reservation, before the fold and the split widen anything. Counted by hand: the packed layers as held, + # the folded embedding and the norms at bf16, and the embedding's scalar scale. The LM head is in none of it: + # this model has none, and folding it first would only be paid for and thrown away. + packed = sum(_packed_bytes(rows, columns) for rows, columns in PROJECTIONS.values()) + assert run.reserved == [packed + DENSE_ELEMENTS * COMPUTE_DTYPE.itemsize + 4] + assert run.order and all(reserved == 1 for _step, reserved in run.order), run.order + + +@pytest.mark.parametrize("naming", ["marker", "header"]) +def test_an_int8_neighbour_leaves_the_nvfp4_layers_packed_and_reserved( + monkeypatch: pytest.MonkeyPatch, naming: str +) -> None: + """The int8 branch strips every marker it leaves, nvfp4's included, so the packed layers must be out first; and + its reservation is the only one, so it has to count them.""" + state_dict, metadata, expected = _checkpoint(naming) + int8 = "model.layers.0.mlp.down_proj" + # An int8 file carries no scaled fp8 beside it: that branch refuses foreign float8 scales. + for path in ("model.embed_tokens", "lm_head"): + state_dict[f"{path}.weight"] = expected[f"{path}.dequantized"].to(torch.bfloat16) + del state_dict[f"{path}.weight_scale"], state_dict[f"{path}.comfy_quant"] + for suffix in (".weight", ".weight_scale", ".weight_scale_2", ".comfy_quant"): + state_dict.pop(f"{int8}{suffix}", None) + if metadata is not None: + header = json.loads(metadata["_quantization_metadata"]) + del header["layers"][int8] + metadata = {"_quantization_metadata": json.dumps(header)} + codes = torch.randint(-127, 128, (HIDDEN, INTERMEDIATE), dtype=torch.int8) + scale = torch.rand(HIDDEN, 1) + 0.5 + state_dict[f"{int8}.weight"] = codes + state_dict[f"{int8}.weight_scale"] = scale + state_dict[f"{int8}.comfy_quant"] = comfy_quant_marker({"format": "int8_tensorwise", "convrot": False}) + run = _prepare(monkeypatch, state_dict, metadata) + + model = run.load(_config()) + + for path in PROJECTIONS: + if path != int8: + _assert_packed(model, path, expected[_transformers_path(path)]) + int8_module = model.get_submodule(_transformers_path(int8)) + assert isinstance(int8_module, Int8ConvrotLinear) + torch.testing.assert_close( + int8_module._dequantized_weight(torch.device("cpu"), torch.float32), codes.float() * scale + ) + + packed = sum(_packed_bytes(rows, columns) for path, (rows, columns) in PROJECTIONS.items() if path != int8) + held_int8 = HIDDEN * INTERMEDIATE + HIDDEN * scale.element_size() + assert run.reserved == [packed + held_int8 + DENSE_ELEMENTS * COMPUTE_DTYPE.itemsize] + + +def test_an_mxfp8_build_keeps_its_language_model_in_fp8_storage(monkeypatch: pytest.MonkeyPatch) -> None: + """Every MXFP8 Linear is folded -- `_scaled_mm` cannot apply a block-wise scale -- so its scaled layers reach the + storage-only pass, which must still hold them in fp8: only a folded *embedding* stays wide there. Kept at the + compute dtype instead, the language model would be resident at twice the size it had before.""" + state_dict = _dense_checkpoint() + expected: dict[str, torch.Tensor] = {} + torch.manual_seed(1) + for path, (rows, columns) in PROJECTIONS.items(): + # Exponent bytes around the bias: block scales from 0.5 to 4, exact in e4m3 as in bf16. + tensors, block_scales = mxfp8_tensors(path, torch.randint(126, 130, (rows, columns // MX_BLOCK_SIZE))) + state_dict.update(tensors) + state_dict[f"{path}.comfy_quant"] = comfy_quant_marker(mxfp8_marker()) + expected[_transformers_path(path)] = block_scales.repeat_interleave(MX_BLOCK_SIZE, dim=1) + run = _prepare(monkeypatch, state_dict, None, storage=True) + + model = run.load(_config()) + + for path, weight in expected.items(): + stored = model.get_submodule(path).weight + assert stored.dtype is torch.float8_e4m3fn, path + assert torch.equal(stored.float(), weight), path + assert any("FP8 layerwise casting enabled" in message for message in _logged(run)) + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + pytest.param( + lambda sd: sd.pop("model.layers.0.self_attn.q_proj.comfy_quant"), + "no ComfyUI `comfy_quant` marker", + id="unnamed", + ), + pytest.param( + lambda sd: sd.pop("model.layers.0.self_attn.q_proj.weight_scale_2"), + "no weight_scale_2", + id="no-global-scale", + ), + pytest.param( + lambda sd: sd.update({"model.layers.0.self_attn.q_proj.pre_quant_scale": torch.ones(HIDDEN)}), + "AWQ", + id="awq", + ), + ], +) +def test_a_layer_the_decode_would_misread_is_refused_before_the_loader_reserves_room( + monkeypatch: pytest.MonkeyPatch, mutate, message: str +) -> None: + state_dict, metadata, _ = _checkpoint("marker") + mutate(state_dict) + run = _prepare(monkeypatch, state_dict, metadata) + + with pytest.raises(ValueError, match=message): + run.load(_config()) + + assert run.reserved == [] diff --git a/tests/backend/model_manager/load/test_qwen_image_nvfp4_loader.py b/tests/backend/model_manager/load/test_qwen_image_nvfp4_loader.py index fa461b027fa..4dde6e2f88a 100644 --- a/tests/backend/model_manager/load/test_qwen_image_nvfp4_loader.py +++ b/tests/backend/model_manager/load/test_qwen_image_nvfp4_loader.py @@ -325,7 +325,7 @@ def test_an_nvfp4_layer_missing_its_global_scale_is_refused_before_the_cache_is_ """The degraded half-state, at the seam rather than at the detector. A packed uint8 weight with a block-scale grid and no `weight_scale_2` is the shape a guard keyed - on `weight_scale_2` -- what the decode keys on -- lets straight through. `_find_nvfp4_layers` + on `weight_scale_2` -- what the decode keys on -- lets straight through. `find_nvfp4_layers` refuses it, and `test_nvfp4.py` pins that; what only a seam can answer is whether this loader still reaches the detector *before* it asks the cache for room, which nothing but the order of those two statements secures. 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 index 20a997ac9d2..e82b77d463f 100644 --- a/tests/backend/model_manager/load/test_scaled_fp8_fold_axis.py +++ b/tests/backend/model_manager/load/test_scaled_fp8_fold_axis.py @@ -201,7 +201,7 @@ def test_an_nvfp4_layer_is_refused_rather_than_folded_over_its_packed_codes(self def test_an_nvfp4_layer_missing_its_global_scale_is_refused_too(self) -> None: """The half-state: packed codes and a block-scale grid, no `weight_scale_2`. It is the shape - `_find_nvfp4_layers` refuses by name, so a build that loses the global scale somewhere + `find_nvfp4_layers` refuses by name, so a build that loses the global scale somewhere upstream arrives here looking like ordinary scaled fp8 -- and the fold takes it exactly as readily as the whole layer, for the same reason and to the same end. Keying the guard on `weight_scale_2` alone, which is what the decode keys on, would let this one through.""" 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 index 7f83938ad2c..2de7454d840 100644 --- a/tests/backend/model_manager/load/test_z_image_loader_boundaries.py +++ b/tests/backend/model_manager/load/test_z_image_loader_boundaries.py @@ -422,7 +422,7 @@ def test_an_nvfp4_layer_missing_its_global_scale_is_refused_before_the_cache_is_ """The degraded half-state, at this seam rather than at the detector. A packed uint8 weight with a block-scale grid and no `weight_scale_2` is the shape a guard keyed - on `weight_scale_2` -- the key the decode keys on -- lets straight through. `_find_nvfp4_layers` + on `weight_scale_2` -- the key the decode keys on -- lets straight through. `find_nvfp4_layers` refuses it and `test_nvfp4.py` pins that; what only a seam can answer is whether this loader reaches the detector before it asks the cache for room, which nothing but the order of those two statements secures. diff --git a/tests/backend/quantization/test_fp8_scaled.py b/tests/backend/quantization/test_fp8_scaled.py index 606d81e18c8..cb07a487c82 100644 --- a/tests/backend/quantization/test_fp8_scaled.py +++ b/tests/backend/quantization/test_fp8_scaled.py @@ -6,6 +6,7 @@ import pytest import torch +from invokeai.backend.quantization import fp8_scaled from invokeai.backend.quantization.block_scale_tiles import unblock_scale_grid from invokeai.backend.quantization.fp8_scaled import ( FP8_DTYPE, @@ -862,6 +863,35 @@ def test_it_is_never_left_quantized(self) -> None: assert torch.allclose(sd["lin.weight"].float(), torch.full((4, 4), 0.5)) +class TestFoldInSlabs: + """A tensor larger than a slab is folded a slab of rows at a time, and must come out as the whole fold would. + + The expectations are the definition -- float32 product, one rounding to the target dtype -- written out per + scale layout rather than taken from the expansion the fold uses. The slab is shrunk so a 37-row weight splits + into slabs of two rows and an odd one at the end, the boundaries a row-indexing mistake would show at. + """ + + @pytest.mark.parametrize("layout", ["per_tensor", "per_row", "block_wise"]) + def test_slabs_fold_exactly_as_the_whole(self, monkeypatch: pytest.MonkeyPatch, layout: str) -> None: + monkeypatch.setattr(fp8_scaled, "_FOLD_SLAB_ELEMENTS", 100) + torch.manual_seed(0) + codes = torch.randn(37, 48).to(FP8_DTYPE) + if layout == "per_tensor": + scale = torch.tensor(0.37) + expanded = scale + elif layout == "per_row": + scale = torch.rand(37) + 0.5 + expanded = scale[:, None] + else: # one scale per 16 columns + scale = torch.rand(37, 3) + 0.5 + expanded = scale.repeat_interleave(16, dim=1) + sd = {"lin.weight": codes} + + dequantize_fp8_scaled(sd, {"lin": Fp8ScaledLayer(weight_scale=scale)}, torch.bfloat16) + + assert torch.equal(sd["lin.weight"], (codes.float() * expanded).to(torch.bfloat16)) + + class TestBlockWiseScale: """A 2-D ``weight_scale`` has one entry per *block* of weight elements, not per row."""