Skip to content

fix(mm): stop the SDNQ guards from aborting folder-encoder identification - #9561

Open
Pfannkuchensack wants to merge 4 commits into
invoke-ai:mainfrom
Pfannkuchensack:fix/qwen3_encoder_sharded_probe
Open

fix(mm): stop the SDNQ guards from aborting folder-encoder identification#9561
Pfannkuchensack wants to merge 4 commits into
invoke-ai:mainfrom
Pfannkuchensack:fix/qwen3_encoder_sharded_probe

Conversation

@Pfannkuchensack

@Pfannkuchensack Pfannkuchensack commented Aug 30, 2026

Copy link
Copy Markdown
Member

Summary

fix(mm): stop the SDNQ guards from aborting folder-encoder identification

The SDNQ support in #9228 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 and the config never matches. With no candidate left, the model falls
back to Unknown_Config.

Every folder-layout Qwen3 encoder we ship as a starter model is sharded, so all three install as
"Unable to identify model":

Starter model Source Shards
FLUX.2 Klein Qwen3 8B Encoder black-forest-labs/FLUX.2-klein-9B::text_encoder+tokenizer 4
FLUX.2 Klein Qwen3 4B Encoder black-forest-labs/FLUX.2-klein-4B::text_encoder+tokenizer 2
Z-Image Qwen3 Text Encoder Tongyi-MAI/Z-Image-Turbo::text_encoder+tokenizer 3

What changed

  1. The SDNQ checks in identification read tensor names from safetensors headers.
    safetensors_tensor_names() / safetensors_have_sdnq_keys() live in
    backend/quantization/sdnq/detection.py — the module the SDNQ PR created precisely so this
    question has one implementation — and folder_has_sdnq_keys() now delegates to them. 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.

  2. The two Qwen3 folder configs now ask the same question. Qwen3Encoder_Qwen3Encoder_Config
    (rejects SDNQ) and Qwen3Encoder_SDNQ_Folder_Config (requires it) must partition folders, but they
    looked at different places: the unquantized one checked the marker in text_encoder/ too and keys
    across every shard, the SDNQ one only the root and only safetensors sitting directly in it. A
    markerless SDNQ encoder in the nested text_encoder/ layout was therefore rejected by both
    the exact shape that lands a model in unknown. Both now call one shared predicate.

  3. A corrupt quantization_config.json no longer aborts the probe.
    Qwen3Encoder_SDNQ_Folder_Config called json.load() unguarded; it now falls through to the key
    check, like every other marker read.

  4. The Qwen3-only q_norm/k_norm fallback works for sharded folders. It read the state dict
    inside a bare except, so a sharded folder yielded no signal at all and was rejected. It reads
    header names now.

No behaviour changes for single-file checkpoints or GGUFs — those paths still use the state dict,
which is correct for them.

Related Issues / Discussions

Closes #9567

QA Instructions

Verified against a real black-forest-labs/FLUX.2-klein-9B::text_encoder+tokenizer install (4 shards)
on Windows:

  • Before: ModelConfigFactory.from_model_on_disk()Unknown_Config, with
    Qwen3Encoder_Qwen3Encoder_Config: ValueError: Multiple weight files found for this model in the
    details.
  • After: Qwen3Encoder_Qwen3Encoder_Config, variant=qwen3_8b, and it is selectable for FLUX.2
    Klein 9B.
  • With ModelOnDisk.load_state_dict patched to raise, identification still succeeds.

To reproduce in the UI: Model Manager → Starter Models → install "FLUX.2 Klein Qwen3 8B Encoder"
(or the 4B / Z-Image encoders). It should register as a Qwen3 encoder, not as Unknown.

New tests:

  • a sharded folder identifies, and a sharded SDNQ folder is still rejected by the unquantized config
  • a folder probe that fails if load_state_dict() is called at all
  • a markerless nested SDNQ encoder is claimed by exactly one of the two configs
  • an SDNQ marker in text_encoder/ is honoured
  • a corrupt marker falls through to the key check
  • a sharded folder with no declared architecture is matched via header q_norm keys

Merge Plan

Normal merge. Supersedes the narrower first pass on this branch.

Checklist

  • The PR has a short but descriptive title, suitable for a changelog
  • Tests added / updated (if applicable)
  • ❗Changes to a redux slice have a corresponding migration
  • Documentation added / updated (if applicable)
  • Updated What's New copy (if doing a release after this PR)

…ck 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) <noreply@anthropic.com>
@github-actions github-actions Bot added python PRs that change python files backend PRs that change backend files python-tests PRs that change python tests labels Aug 30, 2026
…tion

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 invoke-ai#9567
@Pfannkuchensack Pfannkuchensack changed the title fix(mm): identify sharded Qwen3 encoder folders instead of falling back to unknown fix(mm): stop the SDNQ guards from aborting folder-encoder identification Sep 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

6.14.2 backend PRs that change backend files python PRs that change python files python-tests PRs that change python tests

Projects

Status: 6.14.2

Development

Successfully merging this pull request may close these issues.

[bug]: Official FLUX.2 Klein Qwen3 8B Encoder installs as Unknown; metadata update fails on macOS

2 participants