From 8bf49dd57d2777e118f7c20e9f087458021d17ff Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Sun, 30 Aug 2026 21:43:58 +0200 Subject: [PATCH 1/2] fix(mm): identify sharded Qwen3 encoder folders instead of falling back to unknown The starter-model download `black-forest-labs/FLUX.2-klein-{4B,9B}::text_encoder+tokenizer` lands the Qwen3 encoder as several `model-0000N-of-0000M.safetensors` shards. The SDNQ rejection guard in `Qwen3Encoder_Qwen3Encoder_Config` called `mod.load_state_dict()`, which raises `ValueError("Multiple weight files found for this model")` - not a `NotAMatchError` - when a folder holds more than one weight file. That aborted this config's probe entirely, so no candidate matched and the encoder was stored as `unknown`. Make the SDNQ key check shard-safe: read tensor *names* from the safetensors headers per shard instead of loading a state dict. That is cheap (no tensor data is materialized), works for any number of shards, and keeps detecting SDNQ weight+scale pairs even when they are split across shards. The mirrored fallback in `Qwen3Encoder_SDNQ_Folder_Config` had the same crash for sharded SDNQ folders and now uses the same helper, with its file scope unchanged. Verified against a real `FLUX.2-klein-9B::text_encoder+tokenizer` install, which now identifies as `Qwen3Encoder_Qwen3Encoder_Config` / variant `qwen3_8b`. Co-Authored-By: Claude Opus 5 (1M context) --- .../model_manager/configs/qwen3_encoder.py | 49 ++++++++++++++---- .../configs/test_qwen3_encoder_config.py | 50 +++++++++++++++++++ 2 files changed, 88 insertions(+), 11 deletions(-) diff --git a/invokeai/backend/model_manager/configs/qwen3_encoder.py b/invokeai/backend/model_manager/configs/qwen3_encoder.py index 1cd98729b85..92b6352fdf9 100644 --- a/invokeai/backend/model_manager/configs/qwen3_encoder.py +++ b/invokeai/backend/model_manager/configs/qwen3_encoder.py @@ -1,7 +1,10 @@ import json +from collections.abc import Iterable +from pathlib import Path from typing import Any, Literal, Optional, Self from pydantic import Field +from safetensors import safe_open from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Config_Base from invokeai.backend.model_manager.configs.identification_utils import ( @@ -52,17 +55,44 @@ def _has_sdnq_tensors(state_dict: dict[str | int, Any]) -> bool: return any(isinstance(v, SDNQTensor) for v in state_dict.values()) -def _has_sdnq_keys(state_dict: dict[str | int, Any]) -> bool: - """Check if state dict has SDNQ-style keys (weight + scale pairs).""" - keys = {k for k in state_dict.keys() if isinstance(k, str)} - for key in keys: +def _keys_look_sdnq(keys: Iterable[str]) -> bool: + """Check if a set of tensor names has SDNQ-style keys (weight + scale pairs).""" + key_set = {k for k in keys if isinstance(k, str)} + for key in key_set: if key.endswith(".weight"): base = key[:-7] - if f"{base}.scale" in keys: + if f"{base}.scale" in key_set: return True return False +def _has_sdnq_keys(state_dict: dict[str | int, Any]) -> bool: + """Check if state dict has SDNQ-style keys (weight + scale pairs).""" + return _keys_look_sdnq(k for k in state_dict.keys() if isinstance(k, str)) + + +def _files_look_sdnq_quantized(files: Iterable[Path]) -> bool: + """Best-effort SDNQ key check over safetensors files, safe for sharded checkpoints. + + ``ModelOnDisk.load_state_dict()`` refuses to pick a file when a folder holds more than one weight + file, so calling it on a *sharded* encoder raises ValueError instead of a NotAMatchError - which + aborts identification for that config entirely. FLUX.2 Klein's ``text_encoder+tokenizer`` download + ships the Qwen3 encoder as 2-4 safetensors shards, so the SDNQ fallback must never go through + ``load_state_dict()``. We only need tensor *names* here, so read them from the safetensors headers: + cheap, per-shard, and no tensor data is materialized. + """ + keys: set[str] = set() + for file in files: + if file.suffix != ".safetensors": + continue + try: + with safe_open(file, framework="pt", device="cpu") as f: + keys.update(f.keys()) + except Exception: + continue + return _keys_look_sdnq(keys) + + def _has_t5_encoder_keys(state_dict: dict[str | int, Any]) -> bool: """Check if state dict looks like a llama.cpp T5 encoder. @@ -357,7 +387,7 @@ def _reject_if_sdnq_quantized(cls, mod: ModelOnDisk) -> None: if quant_config.get("quant_method") == "sdnq": raise NotAMatchError("folder is SDNQ-quantized; use Qwen3Encoder_SDNQ_Folder_Config") - if _has_sdnq_keys(mod.load_state_dict()): + if _files_look_sdnq_quantized(mod.weight_files()): raise NotAMatchError("state dict looks SDNQ-quantized; use Qwen3Encoder_SDNQ_Folder_Config") @classmethod @@ -539,11 +569,8 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - # Fallback: check if safetensors files have SDNQ-style keys if not matched: - safetensors_files = list(mod.path.glob("*.safetensors")) - if safetensors_files: - state_dict = mod.load_state_dict() - if _has_sdnq_keys(state_dict): - matched = True + if _files_look_sdnq_quantized(mod.path.glob("*.safetensors")): + matched = True if not matched: raise NotAMatchError("directory does not look like an SDNQ-quantized Qwen3 encoder") diff --git a/tests/backend/model_manager/configs/test_qwen3_encoder_config.py b/tests/backend/model_manager/configs/test_qwen3_encoder_config.py index 36f0cb404ca..f406a7d0227 100644 --- a/tests/backend/model_manager/configs/test_qwen3_encoder_config.py +++ b/tests/backend/model_manager/configs/test_qwen3_encoder_config.py @@ -15,6 +15,8 @@ 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.qwen3_encoder import ( @@ -24,6 +26,7 @@ _has_gemma2_keys, _has_qwen_vl_visual_tower, ) +from invokeai.backend.model_manager.model_on_disk import ModelOnDisk _OVERRIDE_FIELDS: dict[str, object] = { "hash": "blake3:fakehash", @@ -145,3 +148,50 @@ def test_gguf_config_rejects_gemma_state_dict() -> None: } with pytest.raises(NotAMatchError, match="Gemma-2"): Qwen3Encoder_GGUF_Config._validate_looks_like_qwen3_model(mod) + + +class TestShardedQwen3EncoderFolder: + """A sharded text_encoder folder must still identify as a Qwen3 encoder. + + The starter-model download + `black-forest-labs/FLUX.2-klein-{4B,9B}::text_encoder+tokenizer` lands the encoder as several + `model-0000N-of-0000M.safetensors` shards. The SDNQ rejection guard used to call + `mod.load_state_dict()`, which raises ValueError ("Multiple weight files found") rather than a + NotAMatchError when a folder holds more than one weight file. That aborted this config's probe, + so the model was stored as `unknown`. + """ + + @staticmethod + def _make_sharded_encoder(root: Path, *, sdnq: bool = False) -> Path: + text_encoder = root / "text_encoder" + text_encoder.mkdir(parents=True) + _write_config(text_encoder / "config.json", hidden_size=4096, architecture="Qwen3ForCausalLM") + + shards: list[dict[str, torch.Tensor]] = [ + {"model.embed_tokens.weight": torch.zeros(8, 4096, dtype=torch.uint8 if sdnq else torch.float32)}, + {"model.layers.0.self_attn.q_proj.weight": torch.zeros(8, 8, dtype=torch.uint8 if sdnq else torch.float32)}, + ] + if sdnq: + shards[1]["model.layers.0.self_attn.q_proj.scale"] = torch.zeros(8, 1, dtype=torch.float32) + + for i, shard in enumerate(shards, start=1): + save_file(shard, str(text_encoder / f"model-0000{i}-of-00002.safetensors")) + + # Tokenizer files live in their own subfolder for the `text_encoder+tokenizer` download layout. + tokenizer = root / "tokenizer" + tokenizer.mkdir() + (tokenizer / "tokenizer_config.json").write_text("{}") + return root + + def test_sharded_encoder_matches(self, tmp_path: Path) -> None: + root = self._make_sharded_encoder(tmp_path / "klein-9b-encoder") + config = Qwen3Encoder_Qwen3Encoder_Config.from_model_on_disk(ModelOnDisk(root), dict(_OVERRIDE_FIELDS)) + + assert config.type.value == "qwen3_encoder" + assert config.variant.value == "qwen3_8b" + + def test_sharded_sdnq_encoder_is_still_rejected(self, tmp_path: Path) -> None: + """The shard-safe check must keep detecting SDNQ weight+scale pairs across shards.""" + root = self._make_sharded_encoder(tmp_path / "klein-9b-encoder-sdnq", sdnq=True) + with pytest.raises(NotAMatchError, match="SDNQ"): + Qwen3Encoder_Qwen3Encoder_Config.from_model_on_disk(ModelOnDisk(root), dict(_OVERRIDE_FIELDS)) From fc753bd2a30e5a136492d5273f920aaea754f20a Mon Sep 17 00:00:00 2001 From: Alexander Eichhorn Date: Mon, 7 Sep 2026 23:01:32 +0200 Subject: [PATCH 2/2] fix(mm): stop the SDNQ guards from aborting folder-encoder identification The SDNQ support added an "is this folder SDNQ-quantized?" guard to Qwen3Encoder_Qwen3Encoder_Config - a folder config that until then never touched weights - and implemented it with `mod.load_state_dict()`. That call refuses to pick a file when a directory holds more than one weight file and raises ValueError, which is not a NotAMatchError, so the factory records it as an error, no candidate matches and the model falls back to Unknown_Config. Every folder-layout Qwen3 encoder we ship as a starter model is sharded (FLUX.2 Klein 9B: 4, Klein 4B: 2, Z-Image text_encoder: 3), so all three install as "Unable to identify model". The SDNQ checks in identification now read tensor names from the safetensors headers via safetensors_tensor_names() / safetensors_have_sdnq_keys() in the sdnq detection module - the module the SDNQ PR created so this question has one implementation. Reading headers is per file, so it works for any number of shards, and it resolves weight/scale pairs across the union of all of them: sharding splits a checkpoint by tensor order and routinely separates a weight from its scale. Three further gaps closed while here: - Qwen3Encoder_Qwen3Encoder_Config and Qwen3Encoder_SDNQ_Folder_Config must partition folders but asked different questions: the unquantized one checked the marker in text_encoder/ too and keys across every shard, the SDNQ one only the root and only safetensors directly in it. A markerless SDNQ encoder in the nested layout was rejected by both. Both now call one shared predicate. - A corrupt quantization_config.json aborted the SDNQ probe with a JSONDecodeError instead of falling through to the key check. - The Qwen3-only q_norm/k_norm fallback read the state dict inside a bare except, so a sharded folder yielded no signal and was rejected. Single-file and GGUF configs keep using the state dict, which is correct for them. Verified against a real FLUX.2-klein-9B::text_encoder+tokenizer install: it now identifies as Qwen3Encoder_Qwen3Encoder_Config / variant qwen3_8b, and still does with ModelOnDisk.load_state_dict patched to raise. Closes #9567 --- .../model_manager/configs/qwen3_encoder.py | 136 ++++++++---------- .../backend/quantization/sdnq/detection.py | 75 +++++++--- .../configs/test_qwen3_encoder_config.py | 97 +++++++++++++ 3 files changed, 206 insertions(+), 102 deletions(-) diff --git a/invokeai/backend/model_manager/configs/qwen3_encoder.py b/invokeai/backend/model_manager/configs/qwen3_encoder.py index 92b6352fdf9..ed0f116ef05 100644 --- a/invokeai/backend/model_manager/configs/qwen3_encoder.py +++ b/invokeai/backend/model_manager/configs/qwen3_encoder.py @@ -1,10 +1,8 @@ import json from collections.abc import Iterable -from pathlib import Path from typing import Any, Literal, Optional, Self from pydantic import Field -from safetensors import safe_open from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Config_Base from invokeai.backend.model_manager.configs.identification_utils import ( @@ -17,6 +15,11 @@ from invokeai.backend.model_manager.model_on_disk import ModelOnDisk from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, Qwen3VariantType from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor +from invokeai.backend.quantization.sdnq.detection import ( + folder_has_sdnq_marker, + safetensors_have_sdnq_keys, + safetensors_tensor_names, +) from invokeai.backend.quantization.sdnq.sdnq_tensor import SDNQTensor @@ -55,42 +58,37 @@ def _has_sdnq_tensors(state_dict: dict[str | int, Any]) -> bool: return any(isinstance(v, SDNQTensor) for v in state_dict.values()) -def _keys_look_sdnq(keys: Iterable[str]) -> bool: - """Check if a set of tensor names has SDNQ-style keys (weight + scale pairs).""" - key_set = {k for k in keys if isinstance(k, str)} - for key in key_set: - if key.endswith(".weight"): - base = key[:-7] - if f"{base}.scale" in key_set: - return True - return False - - def _has_sdnq_keys(state_dict: dict[str | int, Any]) -> bool: """Check if state dict has SDNQ-style keys (weight + scale pairs).""" - return _keys_look_sdnq(k for k in state_dict.keys() if isinstance(k, str)) + keys = {k for k in state_dict.keys() if isinstance(k, str)} + return any(key.endswith(".weight") and f"{key[: -len('.weight')]}.scale" in keys for key in keys) -def _files_look_sdnq_quantized(files: Iterable[Path]) -> bool: - """Best-effort SDNQ key check over safetensors files, safe for sharded checkpoints. +def _folder_tensor_names(mod: ModelOnDisk) -> set[str]: + """Tensor names declared by every safetensors shard under `mod.path`. - ``ModelOnDisk.load_state_dict()`` refuses to pick a file when a folder holds more than one weight - file, so calling it on a *sharded* encoder raises ValueError instead of a NotAMatchError - which - aborts identification for that config entirely. FLUX.2 Klein's ``text_encoder+tokenizer`` download - ships the Qwen3 encoder as 2-4 safetensors shards, so the SDNQ fallback must never go through - ``load_state_dict()``. We only need tensor *names* here, so read them from the safetensors headers: - cheap, per-shard, and no tensor data is materialized. + Folder probes must not go through ``ModelOnDisk.load_state_dict()``. It refuses to pick a file + when a folder holds more than one weight file and raises ``ValueError`` — not a + ``NotAMatchError`` — which aborts that config's probe rather than declining the model, so a + sharded encoder ends up stored as ``unknown``. Every folder-layout Qwen3 encoder we ship is + sharded: the FLUX.2 Klein 4B/9B and Z-Image ``text_encoder`` downloads are 2-4 shards each. """ - keys: set[str] = set() - for file in files: - if file.suffix != ".safetensors": - continue - try: - with safe_open(file, framework="pt", device="cpu") as f: - keys.update(f.keys()) - except Exception: - continue - return _keys_look_sdnq(keys) + return safetensors_tensor_names(mod.weight_files()) + + +def _folder_is_sdnq_quantized(mod: ModelOnDisk) -> bool: + """True if the Qwen3 encoder folder at `mod.path` holds SDNQ-quantized weights. + + `Qwen3Encoder_Qwen3Encoder_Config` and `Qwen3Encoder_SDNQ_Folder_Config` must be mutually + exclusive: one rejects what the other requires. That only holds if both ask the *same* question, + so both call this. They used to ask different ones — the unquantized config looked for the marker + in `text_encoder/` as well and for keys across every shard, while the SDNQ config looked only at + the root and only at safetensors sitting directly in it. A markerless SDNQ encoder in the nested + `text_encoder/` layout was therefore rejected by *both* and stored as `unknown`. + """ + if any(folder_has_sdnq_marker(folder) for folder in (mod.path, mod.path / "text_encoder")): + return True + return safetensors_have_sdnq_keys(mod.weight_files()) def _has_t5_encoder_keys(state_dict: dict[str | int, Any]) -> bool: @@ -122,8 +120,8 @@ def _has_gemma2_keys(state_dict: dict[str | int, Any]) -> bool: return False -def _has_qwen_vl_visual_tower(state_dict: dict[str | int, Any]) -> bool: - """Check if state dict bundles a Qwen-VL vision tower (Qwen2-VL / Qwen2.5-VL / Qwen3-VL). +def _has_qwen_vl_visual_tower(tensor_names: Iterable[str | int]) -> bool: + """Check if the tensor names bundle a Qwen-VL vision tower (Qwen2-VL / Qwen2.5-VL / Qwen3-VL). VL encoders ship a visual tower alongside the language model, whereas a text-only Qwen3 encoder never does. A VL file otherwise satisfies the Qwen3 key heuristic (it has ``model.layers.*`` / @@ -135,14 +133,17 @@ def _has_qwen_vl_visual_tower(state_dict: dict[str | int, Any]) -> bool: layout that ComfyUI single-file Qwen3-VL checkpoints use. Matching only bare ``visual.blocks.*`` missed that layout, letting a single-file Qwen3-VL 4B encoder match both configs and get misrouted to the text-only Qwen3 type - silently breaking the single-file/GGUF Krea-2 encoder install path. + + Takes any iterable of names so a folder probe can pass the union of its safetensors headers; a + state dict iterates over its keys, so existing call sites are unaffected. """ - for key in state_dict.keys(): + for key in tensor_names: if isinstance(key, str) and (key.startswith(("visual.", "model.visual.")) or ".visual." in key): return True return False -def _has_qwen3_specific_keys(state_dict: dict[str | int, Any]) -> bool: +def _has_qwen3_specific_keys(tensor_names: Iterable[str | int]) -> bool: """Check for Qwen3-only QK-normalization weights (``q_norm``/``k_norm`` per attention block). Qwen3 adds an RMSNorm on the query and key projections that Qwen2 does not have. A Qwen2 @@ -150,8 +151,11 @@ def _has_qwen3_specific_keys(state_dict: dict[str | int, Any]) -> bool: / ``model.embed_tokens.weight`` layout), so this is the discriminator that keeps a Qwen2 file from being accepted as a Qwen3 encoder the loader would fail to build. Covers both the PyTorch/diffusers naming and the llama.cpp/GGUF naming. + + Takes any iterable of names, so a folder probe can pass safetensors header keys instead of a + state dict it cannot load. """ - for key in state_dict.keys(): + for key in tensor_names: if not isinstance(key, str): continue if ".self_attn.q_norm." in key or ".self_attn.k_norm." in key: @@ -372,23 +376,10 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - @classmethod def _reject_if_sdnq_quantized(cls, mod: ModelOnDisk) -> None: - # Primary signal: quantization_config.json with quant_method="sdnq" (at root or in - # text_encoder/). Fallback: SDNQ-style weight+scale key pairs in the state dict. This mirrors - # the detection in Qwen3Encoder_SDNQ_Folder_Config so the two stay mutually exclusive. - for folder in (mod.path, mod.path / "text_encoder"): - quant_config_path = folder / "quantization_config.json" - if not quant_config_path.exists(): - continue - try: - with open(quant_config_path, "r", encoding="utf-8") as f: - quant_config = json.load(f) - except (json.JSONDecodeError, OSError): - continue - if quant_config.get("quant_method") == "sdnq": - raise NotAMatchError("folder is SDNQ-quantized; use Qwen3Encoder_SDNQ_Folder_Config") - - if _files_look_sdnq_quantized(mod.weight_files()): - raise NotAMatchError("state dict looks SDNQ-quantized; use Qwen3Encoder_SDNQ_Folder_Config") + # Shared with Qwen3Encoder_SDNQ_Folder_Config so the two configs cannot both reject (or both + # accept) the same folder. + if _folder_is_sdnq_quantized(mod): + raise NotAMatchError("folder is SDNQ-quantized; use Qwen3Encoder_SDNQ_Folder_Config") @classmethod def _get_variant_from_config(cls, config_path) -> Qwen3VariantType: @@ -557,22 +548,10 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - raise_for_override_fields(cls, override_fields) - matched = False - - # Check for quantization_config.json with quant_method="sdnq" - quant_config_path = mod.path / "quantization_config.json" - if quant_config_path.exists(): - with open(quant_config_path, "r", encoding="utf-8") as f: - quant_config = json.load(f) - if quant_config.get("quant_method") == "sdnq": - matched = True - - # Fallback: check if safetensors files have SDNQ-style keys - if not matched: - if _files_look_sdnq_quantized(mod.path.glob("*.safetensors")): - matched = True - - if not matched: + # Shared with the rejection guard in Qwen3Encoder_Qwen3Encoder_Config: exactly one of the two + # configs claims a given folder. A corrupt or foreign quantization_config.json falls through + # to the key check there instead of aborting this probe with a JSONDecodeError. + if not _folder_is_sdnq_quantized(mod): raise NotAMatchError("directory does not look like an SDNQ-quantized Qwen3 encoder") # A root config.json next to tokenizer files is a complete causal LM (TextLLM), not a Qwen3 @@ -635,25 +614,22 @@ def _validate_is_qwen3_encoder(cls, mod: ModelOnDisk) -> None: "(only Qwen3ForCausalLM is supported)" ) - # Fallback for folders without a usable config.json architecture: check the state dict. The + # Fallback for folders without a usable config.json architecture: check the tensor names. The # generic Qwen keys (model.layers. / model.embed_tokens.weight) are NOT enough — Qwen2 and # Qwen2-VL folders carry exactly the same ones, and the loader reconstructs a text-only # Qwen3ForCausalLM that fails on Qwen2's missing q/k-norm params and on Qwen-VL's visual # tower. Mirror the single-file path: reject a bundled visual tower and require the # Qwen3-only q/k-norm weights. An SDNQ transformer/VAE folder has transformer_blocks. / - # decoder. keys instead and is rejected by both checks. Loading the state dict can raise for - # sharded folders (multiple weight files), so treat that as "no usable signal" rather than - # letting it abort identification. - try: - state_dict = mod.load_state_dict() - except Exception: - state_dict = {} + # decoder. keys instead and is rejected by both checks. Names come from the safetensors + # headers because a sharded folder has no single state dict to load — that path used to + # yield no signal at all and reject every sharded markerless SDNQ encoder. + tensor_names = _folder_tensor_names(mod) - if _has_qwen_vl_visual_tower(state_dict): + if _has_qwen_vl_visual_tower(tensor_names): raise NotAMatchError( "state dict bundles a Qwen-VL visual tower; this is a Qwen-VL encoder, not a text-only Qwen3 encoder" ) - if _has_qwen3_specific_keys(state_dict): + if _has_qwen3_specific_keys(tensor_names): return raise NotAMatchError( diff --git a/invokeai/backend/quantization/sdnq/detection.py b/invokeai/backend/quantization/sdnq/detection.py index c18cba46345..eadc5b4deb6 100644 --- a/invokeai/backend/quantization/sdnq/detection.py +++ b/invokeai/backend/quantization/sdnq/detection.py @@ -10,6 +10,7 @@ """ import json +from collections.abc import Iterable from pathlib import Path from safetensors import safe_open @@ -17,11 +18,44 @@ _QUANTIZATION_CONFIG_FILENAME = "quantization_config.json" -def folder_has_sdnq_keys(folder_path: Path) -> bool: - """True if the safetensors in `folder_path` carry an SDNQ ``.weight`` / ``.scale`` pair. +def safetensors_tensor_names(files: Iterable[Path]) -> set[str]: + """Union of the tensor names declared by `files`, read from their headers. + + Identification only ever needs the *names* — "is there a `.scale` next to this + `.weight`", "is there a visual tower" — never the tensor data, and the header carries them + per file. That is what makes this shard-safe, which `ModelOnDisk.load_state_dict()` is not: it + refuses to pick a file when a folder holds more than one weight file and raises `ValueError` — + not a `NotAMatchError` — so it aborts that config's probe instead of declining the model. Every + sharded encoder folder (the FLUX.2 Klein and Z-Image `text_encoder` downloads are all 2-4 shards) + hit that and fell back to `unknown`. + + Non-safetensors entries are skipped: they are the set `sdnq_sd_loader` reads, and a file whose + header cannot be read contributes nothing rather than failing the whole check. + """ + names: set[str] = set() + for file in sorted(files): + if file.suffix != ".safetensors": + continue + try: + with safe_open(file, framework="pt", device="cpu") as f: + names.update(f.keys()) + except Exception: + continue + return names + + +def safetensors_have_sdnq_keys(files: Iterable[Path]) -> bool: + """True if `files` carry an SDNQ ``.weight`` / ``.scale`` pair. The pair is resolved across the union of every shard, never within a single file: sharding splits a checkpoint by tensor order, so a weight and its scale routinely land in different files. + """ + names = safetensors_tensor_names(files) + return any(name.endswith(".weight") and f"{name[: -len('.weight')]}.scale" in names for name in names) + + +def folder_has_sdnq_keys(folder_path: Path) -> bool: + """True if the safetensors directly in `folder_path` carry an SDNQ weight/scale pair. Only safetensors are inspected, which is the same set `sdnq_sd_loader` reads — a `.bin` holding SDNQ-shaped tensors is not something we could load anyway, so calling it SDNQ would only move the @@ -30,15 +64,23 @@ def folder_has_sdnq_keys(folder_path: Path) -> bool: if not folder_path.is_dir(): return False - keys: set[str] = set() - for shard in sorted(folder_path.glob("*.safetensors")): - try: - with safe_open(shard, framework="pt", device="cpu") as f: - keys.update(f.keys()) - except Exception: - continue + return safetensors_have_sdnq_keys(folder_path.glob("*.safetensors")) + + +def folder_has_sdnq_marker(folder_path: Path) -> bool: + """True if `folder_path` holds a `quantization_config.json` naming SDNQ as the quant method. - return any(key.endswith(".weight") and f"{key[: -len('.weight')]}.scale" in keys for key in keys) + A marker that is missing, unreadable or names another method is not evidence *against* SDNQ + weights — callers fall through to the key shape rather than treating False as "not SDNQ". + """ + marker = folder_path / _QUANTIZATION_CONFIG_FILENAME + if not marker.is_file(): + return False + try: + with open(marker, "r", encoding="utf-8") as f: + return json.load(f).get("quant_method") == "sdnq" + except (json.JSONDecodeError, OSError): + return False def is_sdnq_folder(folder_path: Path) -> bool: @@ -48,15 +90,4 @@ def is_sdnq_folder(folder_path: Path) -> bool: back to the key shape. The fallback is what covers exports that ship no marker — without it such a folder reads as plain diffusers to identification and as SDNQ to nothing at all. """ - marker = folder_path / _QUANTIZATION_CONFIG_FILENAME - if marker.is_file(): - try: - with open(marker, "r", encoding="utf-8") as f: - if json.load(f).get("quant_method") == "sdnq": - return True - except (json.JSONDecodeError, OSError): - pass - # A marker that exists but names another method is not evidence *against* SDNQ keys, so fall - # through rather than returning False here. - - return folder_has_sdnq_keys(folder_path) + return folder_has_sdnq_marker(folder_path) or folder_has_sdnq_keys(folder_path) diff --git a/tests/backend/model_manager/configs/test_qwen3_encoder_config.py b/tests/backend/model_manager/configs/test_qwen3_encoder_config.py index f406a7d0227..2527e9b7779 100644 --- a/tests/backend/model_manager/configs/test_qwen3_encoder_config.py +++ b/tests/backend/model_manager/configs/test_qwen3_encoder_config.py @@ -23,6 +23,7 @@ Qwen3Encoder_Checkpoint_Config, Qwen3Encoder_GGUF_Config, Qwen3Encoder_Qwen3Encoder_Config, + Qwen3Encoder_SDNQ_Folder_Config, _has_gemma2_keys, _has_qwen_vl_visual_tower, ) @@ -195,3 +196,99 @@ def test_sharded_sdnq_encoder_is_still_rejected(self, tmp_path: Path) -> None: root = self._make_sharded_encoder(tmp_path / "klein-9b-encoder-sdnq", sdnq=True) with pytest.raises(NotAMatchError, match="SDNQ"): Qwen3Encoder_Qwen3Encoder_Config.from_model_on_disk(ModelOnDisk(root), dict(_OVERRIDE_FIELDS)) + + def test_folder_probe_never_loads_a_state_dict(self, tmp_path: Path) -> None: + """Identifying a folder encoder must not go through `load_state_dict()` at all. + + A folder holding more than one weight file has no single state dict to load: the call raises + ValueError, which is not a NotAMatchError, and the probe is abandoned. Everything this config + needs is in `config.json` and the safetensors headers, so make the absence of that call a + property of the test suite rather than of today's implementation. + """ + root = self._make_sharded_encoder(tmp_path / "klein-9b-encoder-nosd") + + def _explode(*args: object, **kwargs: object) -> None: + raise AssertionError("identification must not load the state dict of a folder model") + + with patch.object(ModelOnDisk, "load_state_dict", _explode): + config = Qwen3Encoder_Qwen3Encoder_Config.from_model_on_disk(ModelOnDisk(root), dict(_OVERRIDE_FIELDS)) + + assert config.variant.value == "qwen3_8b" + + +class TestSdnqQwen3EncoderFolder: + """`Qwen3Encoder_Qwen3Encoder_Config` and `Qwen3Encoder_SDNQ_Folder_Config` must partition folders. + + The two used to ask different questions about the same directory: the unquantized config looked + for the SDNQ marker in `text_encoder/` as well and for SDNQ keys across every shard, while the + SDNQ config looked only at the root and only at safetensors sitting directly in it. An SDNQ + encoder in the nested `text_encoder/` layout was therefore rejected by *both* - the exact shape + that lands a model in `unknown`. + """ + + @staticmethod + def _make_sdnq_encoder(root: Path, *, marker_in: str | None = None, architecture: str | None = None) -> Path: + text_encoder = root / "text_encoder" + text_encoder.mkdir(parents=True) + + config: dict[str, object] = {"hidden_size": 4096} + if architecture is not None: + config["architectures"] = [architecture] + (text_encoder / "config.json").write_text(json.dumps(config)) + + # q_norm/k_norm are the Qwen3-only marker the config falls back to when config.json declares + # no architecture; they are split across shards on purpose. + save_file( + {"model.embed_tokens.weight": torch.zeros(8, 4096, dtype=torch.uint8)}, + str(text_encoder / "model-00001-of-00002.safetensors"), + ) + save_file( + { + "model.layers.0.self_attn.q_norm.weight": torch.zeros(8, 8, dtype=torch.uint8), + "model.layers.0.self_attn.q_norm.scale": torch.zeros(8, 1, dtype=torch.float32), + }, + str(text_encoder / "model-00002-of-00002.safetensors"), + ) + + if marker_in is not None: + (root / marker_in / "quantization_config.json").write_text(json.dumps({"quant_method": "sdnq"})) + return root + + def test_nested_markerless_sdnq_encoder_is_claimed(self, tmp_path: Path) -> None: + """Detected by key shape across the shards of `text_encoder/`, with no marker file at all.""" + root = self._make_sdnq_encoder(tmp_path / "sdnq-nested", architecture="Qwen3ForCausalLM") + + config = Qwen3Encoder_SDNQ_Folder_Config.from_model_on_disk(ModelOnDisk(root), dict(_OVERRIDE_FIELDS)) + + assert config.format.value == "sdnq_quantized" + assert config.variant.value == "qwen3_8b" + # ...and the unquantized config declines the same folder, so exactly one of them matches. + with pytest.raises(NotAMatchError, match="SDNQ"): + Qwen3Encoder_Qwen3Encoder_Config.from_model_on_disk(ModelOnDisk(root), dict(_OVERRIDE_FIELDS)) + + def test_marker_in_text_encoder_subfolder_is_honored(self, tmp_path: Path) -> None: + root = self._make_sdnq_encoder( + tmp_path / "sdnq-marker-nested", marker_in="text_encoder", architecture="Qwen3ForCausalLM" + ) + + config = Qwen3Encoder_SDNQ_Folder_Config.from_model_on_disk(ModelOnDisk(root), dict(_OVERRIDE_FIELDS)) + assert config.format.value == "sdnq_quantized" + + def test_unreadable_marker_falls_through_to_the_key_check(self, tmp_path: Path) -> None: + """A corrupt `quantization_config.json` used to abort this probe with a JSONDecodeError.""" + root = self._make_sdnq_encoder(tmp_path / "sdnq-bad-marker", architecture="Qwen3ForCausalLM") + (root / "quantization_config.json").write_text("{not json") + + config = Qwen3Encoder_SDNQ_Folder_Config.from_model_on_disk(ModelOnDisk(root), dict(_OVERRIDE_FIELDS)) + assert config.format.value == "sdnq_quantized" + + def test_sharded_folder_without_declared_architecture_is_claimed(self, tmp_path: Path) -> None: + """The Qwen3-only q_norm/k_norm fallback must read shard headers, not a single state dict. + + With no `architectures` in config.json the config falls back to tensor names. Reading them + via `load_state_dict()` yielded nothing for a sharded folder, so this shape was rejected. + """ + root = self._make_sdnq_encoder(tmp_path / "sdnq-no-arch") + + config = Qwen3Encoder_SDNQ_Folder_Config.from_model_on_disk(ModelOnDisk(root), dict(_OVERRIDE_FIELDS)) + assert config.format.value == "sdnq_quantized"