From b4002534ccbd2b6b341270e197be5c359a627907 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 7 Sep 2026 13:46:38 -0400 Subject: [PATCH 1/2] fix(model loaders): ignore unexpected checkpoint keys, report them at debug only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-file loaders disagreed, loader by loader, about what an unexpected key from `load_state_dict(strict=False)` means: some raised, one warned, most ignored it silently, and the rest used `strict=True` and let torch raise. The ones that hard-failed turned any harmless extra tensor an exporter happened to serialize into a user-facing crash that needed a code change and a release — Anima went through this twice (#9201, #9402) for tensors the model does not need and the official checkpoint does not contain. Per the team decision on #9437, extra keys are now reported at DEBUG and otherwise ignored everywhere. New `invokeai/backend/util/state_dict_loading.py` holds the single policy: - `log_unexpected_keys()` — DEBUG only, never raises. - `load_state_dict_ignoring_extras()` — a drop-in for `strict=True` that keeps the strictness that matters (every required parameter must be filled, shape mismatches still raise) and drops the strictness that only produces whack-a-mole. - `reject_incomplete_load()` — the meta-device completeness sweep, generalized out of krea2. Stronger than `missing_keys` for models built under `init_empty_weights()`: immune to non-persistent buffers and tied weights. Every previously existing missing-key guard is preserved exactly; only the unexpected-key policy changed. `flux.py`'s bare `assert len(unexpected_keys) == 0` — which carried no message and was stripped entirely under `python -O` — is gone with it. Two consequences worth calling out: - `configs/pid_decoder.py` rejected unexpected keys at *identification* time, deliberately mirroring the loader ("both are fatal there"). Left alone, the PiD relaxation would have been unreachable and the installer would refuse a file that now loads fine. It keeps refusing non-string keys, which `load_state_dict` genuinely cannot survive. - Anima's unexpected-key `RuntimeError` was its only hard load-time guard, so it is replaced with the meta-device sweep rather than dropped — otherwise an incomplete checkpoint would fail mid-inference with "Cannot copy out of meta tensor" instead of at load time. `wan.py::_raise_for_incompatible_keys` deliberately keeps raising: Wan derivatives (Animate, S2V, Fun-Camera) are supersets whose extra branches are the feature the checkpoint exists for, not exporter noise, and it strips the benign extras before that check. Closes #9437 Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01V4nLTdWBQaaDUFVLx9TyZH --- .../backend/ideogram4/quantized_loading.py | 10 +- .../model_manager/configs/pid_decoder.py | 21 +- .../model_manager/load/model_loaders/anima.py | 19 +- .../model_manager/load/model_loaders/flux.py | 47 ++-- .../load/model_loaders/gemma2_encoder.py | 4 +- .../load/model_loaders/ideogram4.py | 14 +- .../model_manager/load/model_loaders/krea2.py | 29 +-- .../load/model_loaders/mistral_encoder.py | 7 +- .../load/model_loaders/qwen_image.py | 15 +- .../model_manager/load/model_loaders/vae.py | 5 +- .../load/model_loaders/z_image.py | 24 ++- invokeai/backend/pid/decode.py | 16 +- invokeai/backend/quantization/sdnq/loaders.py | 18 +- invokeai/backend/util/state_dict_loading.py | 145 +++++++++++++ .../ideogram4/test_quantized_loading.py | 15 +- .../configs/test_pid_decoder_config.py | 35 ++- ...encoder_sdnq_single_file_identification.py | 7 +- .../load/test_gemma2_encoder_gguf_loader.py | 27 ++- tests/backend/pid/test_pid_decode.py | 18 +- .../backend/pid/test_pid_state_dict_utils.py | 3 +- .../quantization/sdnq/test_sdnq_loader.py | 13 +- tests/backend/util/test_state_dict_loading.py | 203 ++++++++++++++++++ 22 files changed, 558 insertions(+), 137 deletions(-) create mode 100644 invokeai/backend/util/state_dict_loading.py create mode 100644 tests/backend/util/test_state_dict_loading.py diff --git a/invokeai/backend/ideogram4/quantized_loading.py b/invokeai/backend/ideogram4/quantized_loading.py index 1b710293bde..326a44e04c5 100644 --- a/invokeai/backend/ideogram4/quantized_loading.py +++ b/invokeai/backend/ideogram4/quantized_loading.py @@ -7,6 +7,8 @@ import torch.nn as nn import torch.nn.functional as F +from invokeai.backend.util.state_dict_loading import log_unexpected_keys + if TYPE_CHECKING: pass @@ -104,8 +106,7 @@ def load_bnb4bit_state_dict( real_missing = [m for m in missing if m not in consumed] if real_missing: raise RuntimeError(f"missing keys after quantized load: {real_missing[:10]}") - if unexpected: - raise RuntimeError(f"unexpected keys after quantized load: {unexpected[:10]}") + log_unexpected_keys("Ideogram 4 quantized checkpoint", unexpected) for p in model.parameters(): if isinstance(p, bnb.nn.Params4bit): @@ -258,7 +259,7 @@ def load_fp8_state_dict( the caller must have already put the unquantized params in ``dtype``. ``strict=False`` downgrades missing keys to a warning (e.g. tied weights that a - ``transformers`` model resolves itself); unexpected keys always raise. + ``transformers`` model resolves itself); unexpected keys are logged at DEBUG and ignored. """ prepared: dict[str, torch.Tensor] = {} for k, v in state_dict.items(): @@ -272,8 +273,7 @@ def load_fp8_state_dict( prepared[k] = v.to(device=device) missing, unexpected = model.load_state_dict(prepared, strict=False, assign=assign) - if unexpected: - raise RuntimeError(f"unexpected keys after fp8 load: {unexpected[:10]}") + log_unexpected_keys("Ideogram 4 fp8 checkpoint", unexpected) if missing: if strict: raise RuntimeError(f"missing keys after fp8 load: {missing[:10]}") diff --git a/invokeai/backend/model_manager/configs/pid_decoder.py b/invokeai/backend/model_manager/configs/pid_decoder.py index dcc783953ad..a34dfa7aa8c 100644 --- a/invokeai/backend/model_manager/configs/pid_decoder.py +++ b/invokeai/backend/model_manager/configs/pid_decoder.py @@ -137,9 +137,14 @@ def _raise_if_pid_net_contract_unmet(shapes: _Shapes, contract: Mapping[str, tup guarantee — loaders run under `skip_torch_weight_init()`, so a weight the checkpoint does not supply is uninitialised memory rather than a default. - Missing *and* unexpected keys are fatal here because both are fatal there, which is what makes - installation and loading accept the same set of files. A stricter installer cannot reject a file - that would have loaded: the loader already refuses everything rejected here. + Missing keys are fatal here because they are fatal there, which is what makes installation and + loading accept the same set of files. A stricter installer cannot reject a file that would have + loaded: the loader already refuses everything rejected here. + + Extra keys are *not* fatal — `load_pid_decoder` ignores them (issue #9437), so rejecting them + here would refuse to install a file that loads fine. The one kind of extra key the loader still + cannot survive is a non-string one, which makes `nn.Module.load_state_dict` raise from inside + torch, so that is the extra this check keeps. `_LATENT_PROJ_KEY` is excluded from the shape comparison, and only from that: it is the one parameter whose shape legitimately varies by backbone, and its variable dimensions each have a @@ -151,18 +156,18 @@ def _raise_if_pid_net_contract_unmet(shapes: _Shapes, contract: Mapping[str, tup # Both sorts take `key=str`: a bare checkpoint's keys need not all be strings (see # `strip_net_prefix`), and sorting a mixed set raises TypeError — which the factory answers with # the `Unknown_Config` registration these checks exist to prevent, so the crash fails as a silent - # accept rather than loudly. Only `unexpected` can hold one today; sorting both the same way keeps - # that from depending on which set is on which side of the subtraction. + # accept rather than loudly. if missing := sorted(contract.keys() - shapes.keys(), key=str): raise InvalidMatchError( f"PiD checkpoint is missing {len(missing)} of the weights required by PidNet; the file is " f"incomplete and cannot be used as a PiD decoder: {missing[:5]}{_and_more(missing)}" ) - if unexpected := sorted(shapes.keys() - contract.keys(), key=str): + if not_strings := sorted((k for k in shapes if not isinstance(k, str)), key=str): raise InvalidMatchError( - f"PiD checkpoint has {len(unexpected)} keys PidNet does not expect, which `load_pid_decoder` " - f"rejects too: {unexpected[:5]}{_and_more(unexpected)}" + f"PiD checkpoint has {len(not_strings)} keys that are not strings and so cannot name a " + f"PidNet parameter, which `load_pid_decoder` rejects too: " + f"{not_strings[:5]}{_and_more(not_strings)}" ) mismatched = [(k, shapes[k], want) for k, want in contract.items() if k != _LATENT_PROJ_KEY and shapes[k] != want] diff --git a/invokeai/backend/model_manager/load/model_loaders/anima.py b/invokeai/backend/model_manager/load/model_loaders/anima.py index 97782b07ffd..5689665c11d 100644 --- a/invokeai/backend/model_manager/load/model_loaders/anima.py +++ b/invokeai/backend/model_manager/load/model_loaders/anima.py @@ -21,6 +21,7 @@ ) from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.util.state_dict_loading import log_unexpected_keys, reject_incomplete_load logger = InvokeAILogger.get_logger(__name__) @@ -175,17 +176,13 @@ def _load_from_singlefile( sd[k] = sd[k].to(model_dtype) load_result = model.load_state_dict(sd, assign=True, strict=False) - if load_result.unexpected_keys: - raise RuntimeError( - f"Checkpoint contains {len(load_result.unexpected_keys)} unexpected keys. " - f"This may indicate a corrupted or incompatible checkpoint. " - f"First 5 unexpected keys: {load_result.unexpected_keys[:5]}" - ) - if load_result.missing_keys: - logger.warning( - f"Checkpoint is missing {len(load_result.missing_keys)} keys " - f"(expected for inv_freq buffers). First 5: {load_result.missing_keys[:5]}" - ) + log_unexpected_keys("Anima transformer checkpoint", load_result.unexpected_keys) + # `missing_keys` alone cannot police completeness here: AnimaTransformer's only three buffers + # are registered `persistent=False`, so they never appear in it (the old warning claiming + # otherwise was misleading). Sweep for tensors the checkpoint left on the meta device instead + # — that is the failure worth catching, and it is what the removed unexpected-key + # `RuntimeError` was really standing in for. + reject_incomplete_load(model, what="Anima transformer checkpoint") # Without this the `fp8_storage` toggle is shown for Anima models but does nothing. The # state dict was cast to a single `model_dtype` above, so the layerwise cast has one diff --git a/invokeai/backend/model_manager/load/model_loaders/flux.py b/invokeai/backend/model_manager/load/model_loaders/flux.py index 8ad9459735a..f31f58710c9 100644 --- a/invokeai/backend/model_manager/load/model_loaders/flux.py +++ b/invokeai/backend/model_manager/load/model_loaders/flux.py @@ -86,6 +86,7 @@ from invokeai.backend.quantization.sdnq.loaders import raise_on_incomplete_sdnq_load, sdnq_sd_loader from invokeai.backend.util.logging import InvokeAILogger from invokeai.backend.util.silence_warnings import SilenceWarnings +from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras logger = InvokeAILogger.get_logger(__name__) @@ -117,7 +118,7 @@ def _load_model( with accelerate.init_empty_weights(): model = AutoEncoder(get_flux_ae_params()) sd = load_file(model_path) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="FLUX VAE checkpoint", assign=True) # VAE is broken in float16, which mps defaults to if self._torch_dtype == torch.float16: try: @@ -241,7 +242,7 @@ def _load_model( for k in sd.keys(): sd[k] = sd[k].to(torch.bfloat16) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="FLUX.2 VAE checkpoint", assign=True) # VAE is broken in float16, which mps defaults to if self._torch_dtype == torch.float16: @@ -322,9 +323,13 @@ def _load_state_dict_into_t5(cls, model: T5EncoderModel, state_dict: dict[str, t # There is a shared reference to a single weight tensor in the model. # Both "encoder.embed_tokens.weight" and "shared.weight" refer to the same tensor, so only the latter should # be present in the state_dict. - missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False, assign=True) - assert len(unexpected_keys) == 0 - assert set(missing_keys) == {"encoder.embed_tokens.weight"} + load_state_dict_ignoring_extras( + model, + state_dict, + source="FLUX bnb-int8 T5 encoder", + assign=True, + allowed_missing={"encoder.embed_tokens.weight"}, + ) # Re-tie shared weights. In transformers 5.x, weight tying is implemented at the # parameter level (via _tie_weights / tie_weights) rather than as a Python object # alias. load_state_dict(assign=True) replaces parameters in-place, which severs @@ -467,7 +472,7 @@ def _load_from_gguf(self, config: T5Encoder_GGUF_Config) -> AnyModel: model = T5EncoderModel(t5_config) # Leave transformer Linear weights as GGMLTensors; the autocast cache handles them. - model.load_state_dict(sd, strict=False, assign=True) + load_state_dict_ignoring_extras(model, sd, source="FLUX GGUF T5 encoder", assign=True, allow_missing=True) # Embedding lookups can't run on quantized GGMLTensors, so dequantize the token embeddings and # re-tie the encoder's embed_tokens to the shared embedding. @@ -720,7 +725,7 @@ def _load_from_singlefile( for k in sd.keys(): # We need to cast to bfloat16 due to it being the only currently supported dtype for inference sd[k] = sd[k].to(torch.bfloat16) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="FLUX transformer checkpoint", assign=True) return model @@ -766,7 +771,7 @@ def _load_from_singlefile( img_in_weight.quantized_data = img_in_weight.quantized_data.view(expected_img_in_weight_shape) img_in_weight.tensor_shape = expected_img_in_weight_shape - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="FLUX GGUF transformer checkpoint", assign=True) return model @@ -808,7 +813,7 @@ def _load_from_singlefile( sd = load_file(model_path) if "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale" in sd: sd = convert_bundle_to_flux_transformer_checkpoint(sd) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="FLUX nf4 transformer checkpoint", assign=True) return model @@ -1068,7 +1073,7 @@ def _load_from_singlefile( converted_sd[k] = converted_sd[k].to(torch.bfloat16) # Load the state dict - guidance weights were already initialized above if missing - model.load_state_dict(converted_sd, assign=True) + load_state_dict_ignoring_extras(model, converted_sd, source="FLUX.2 transformer checkpoint", assign=True) return model @@ -1223,11 +1228,9 @@ def _load_text_encoder(self, config: Main_SDNQ_Diffusers_Flux2_Config) -> AnyMod model = Qwen3ForCausalLM(te_config) sd = sdnq_sd_loader(te_dir, compute_dtype=torch.bfloat16) - missing, unexpected = model.load_state_dict(sd, assign=True, strict=False) - if unexpected: - raise ValueError(f"Unexpected keys loading SDNQ Qwen3 text encoder: {unexpected}") - if missing and missing != ["lm_head.weight"]: - raise ValueError(f"Unexpected missing keys loading SDNQ Qwen3 text encoder: {missing}") + missing = load_state_dict_ignoring_extras( + model, sd, source="SDNQ Qwen3 text encoder", assign=True, allowed_missing={"lm_head.weight"} + ) if missing == ["lm_head.weight"]: model.lm_head.weight = model.model.embed_tokens.weight return model @@ -1316,7 +1319,7 @@ def _load_from_singlefile(self, config: Main_SDNQ_Flux2_Config) -> AnyModel: out2, in2, dtype=torch.bfloat16 ) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="SDNQ FLUX.2 transformer checkpoint", assign=True) return model @@ -1463,7 +1466,7 @@ def _load_from_singlefile( out_features2, in_features2, dtype=torch.bfloat16 ) - model.load_state_dict(converted_sd, assign=True) + load_state_dict_ignoring_extras(model, converted_sd, source="FLUX.2 GGUF transformer checkpoint", assign=True) return model @@ -1500,7 +1503,7 @@ def _load_xlabs_controlnet(self, sd: dict[str, torch.Tensor]) -> AnyModel: # HACK(ryand): Is it safe to assume dev here? model = XLabsControlNetFlux(get_flux_transformers_params(FluxVariantType.Dev)) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="FLUX XLabs ControlNet checkpoint", assign=True) return model def _load_instantx_controlnet(self, sd: dict[str, torch.Tensor]) -> AnyModel: @@ -1511,7 +1514,7 @@ def _load_instantx_controlnet(self, sd: dict[str, torch.Tensor]) -> AnyModel: with accelerate.init_empty_weights(): model = InstantXControlNetFlux(flux_params, num_control_modes) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="FLUX InstantX ControlNet checkpoint", assign=True) return model @@ -1555,7 +1558,7 @@ def _load_model( with accelerate.init_empty_weights(): model = FluxReduxModel() - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="FLUX Redux checkpoint", assign=True) model.to(dtype=torch.bfloat16) return model @@ -1635,7 +1638,7 @@ def _load_sdnq_transformer_checkpoint(self, config: Main_SDNQ_FLUX_Config) -> An if "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale" in sd: sd = convert_bundle_to_flux_transformer_checkpoint(sd) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="SDNQ FLUX transformer checkpoint", assign=True) return model def _load_sdnq_transformer(self, transformer_path: Path, config: Main_SDNQ_Diffusers_FLUX_Config) -> AnyModel: @@ -1649,7 +1652,7 @@ def _load_sdnq_transformer(self, transformer_path: Path, config: Main_SDNQ_Diffu # Convert from diffusers format to BFL format sd = self._convert_diffusers_sd_to_bfl(sd) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="SDNQ FLUX transformer", assign=True) return model def _convert_diffusers_sd_to_bfl(self, sd: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: diff --git a/invokeai/backend/model_manager/load/model_loaders/gemma2_encoder.py b/invokeai/backend/model_manager/load/model_loaders/gemma2_encoder.py index 940c5ae9d08..c11e9524b3c 100644 --- a/invokeai/backend/model_manager/load/model_loaders/gemma2_encoder.py +++ b/invokeai/backend/model_manager/load/model_loaders/gemma2_encoder.py @@ -23,6 +23,7 @@ from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.taxonomy import AnyModel, BaseModelType, ModelFormat, ModelType, SubModelType from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.state_dict_loading import log_unexpected_keys # llama.cpp GGUF tensor-name component -> Gemma2Model (decoder-only) component. PiD consumes only the # decoder stack, so we target Gemma2Model directly: no `model.` prefix and no lm_head. Gemma-2 has no @@ -168,8 +169,7 @@ def load_gemma2_model_from_gguf(gguf_path: Path, compute_dtype: "torch.dtype") - model = Gemma2Model(gemma_config) _missing, unexpected = model.load_state_dict(sd, strict=False, assign=True) - if unexpected: - raise RuntimeError(f"Unexpected keys loading Gemma-2 GGUF encoder: {unexpected[:10]}") + log_unexpected_keys("Gemma-2 GGUF encoder", unexpected) # Materialize the weights that cannot remain quantized: # - the token embedding, because nn.Embedding needs indexed access, and diff --git a/invokeai/backend/model_manager/load/model_loaders/ideogram4.py b/invokeai/backend/model_manager/load/model_loaders/ideogram4.py index 3b430dc0b36..a467b400ce0 100644 --- a/invokeai/backend/model_manager/load/model_loaders/ideogram4.py +++ b/invokeai/backend/model_manager/load/model_loaders/ideogram4.py @@ -32,6 +32,7 @@ SubModelType, ) from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras, log_unexpected_keys def _load_local_state_dict(folder: Path, basename: str) -> dict[str, torch.Tensor]: @@ -133,7 +134,7 @@ def _load_one_transformer(self, folder: Path) -> torch.nn.Module: with accelerate.init_empty_weights(): model: torch.nn.Module = Ideogram4Transformer(Ideogram4Config()) model = quantize_model_nf4(model, modules_to_not_convert=set(), compute_dtype=compute_dtype) - model.load_state_dict(sd, strict=True, assign=True) + load_state_dict_ignoring_extras(model, sd, source="Ideogram 4 nf4 transformer", assign=True) return model if is_fp8_state_dict(sd): @@ -148,7 +149,7 @@ def _load_one_transformer(self, folder: Path) -> torch.nn.Module: # Unquantized fallback. with accelerate.init_empty_weights(): model = Ideogram4Transformer(Ideogram4Config()) - model.load_state_dict(sd, strict=True, assign=True) + load_state_dict_ignoring_extras(model, sd, source="Ideogram 4 transformer", assign=True) return model.to(compute_dtype) def _load_text_encoder(self, model_path: Path) -> AnyModel: @@ -201,11 +202,10 @@ def _load_text_encoder(self, model_path: Path) -> AnyModel: model = quantize_model_nf4(model, modules_to_not_convert=set(), compute_dtype=compute_dtype) _, unexpected = model.load_state_dict(sd, strict=False, assign=True) - # Unexpected keys signal a wrong or contaminated checkpoint and must hard-fail. Missing keys are - # acceptable only for tied weights (resolved by _verify_encoder_fully_materialized via + # Extra keys are exporter noise, not a correctness signal - log them and move on. Missing keys + # are acceptable only for tied weights (resolved by _verify_encoder_fully_materialized via # tie_weights); any genuinely missing non-tied weight is caught there as a leftover meta tensor. - if unexpected: - raise RuntimeError(f"unexpected keys loading Ideogram 4 text encoder: {unexpected[:10]}") + log_unexpected_keys("Ideogram 4 text encoder", unexpected) _verify_encoder_fully_materialized(model, context="Ideogram 4 text encoder") if not is_bnb_nf4: model = model.to(compute_dtype) @@ -225,6 +225,6 @@ def _load_vae(self, model_path: Path) -> AnyModel: sd = load_file(model_path / "vae" / "diffusion_pytorch_model.safetensors") sd = convert_diffusers_state_dict(sd) ae = AutoEncoder(AutoEncoderParams()) - ae.load_state_dict(sd) + load_state_dict_ignoring_extras(ae, sd, source="Ideogram 4 VAE") ae.eval() return ae.to(model_dtype) diff --git a/invokeai/backend/model_manager/load/model_loaders/krea2.py b/invokeai/backend/model_manager/load/model_loaders/krea2.py index fbd681c8025..e6b0b55475c 100644 --- a/invokeai/backend/model_manager/load/model_loaders/krea2.py +++ b/invokeai/backend/model_manager/load/model_loaders/krea2.py @@ -26,6 +26,7 @@ ) from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras, reject_incomplete_load if TYPE_CHECKING: # torch is imported lazily inside the helpers below; this is annotations-only. @@ -368,7 +369,9 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: for k in sd.keys(): sd[k] = sd[k].to(model_dtype) - model.load_state_dict(sd, assign=True, strict=False) + load_state_dict_ignoring_extras( + model, sd, source="Krea-2 single-file checkpoint", assign=True, allow_missing=True + ) _reject_incomplete_load(model, what="Krea-2 single-file checkpoint") # `assign=True` aliases every param to its `sd` tensor. Drop the dict's references before # the FP8 cast, or each param's `model_dtype` original stays reachable while its fp8 copy is @@ -421,7 +424,7 @@ def _load_from_gguf(self, config: AnyModelConfig) -> AnyModel: with accelerate.init_empty_weights(): model = Krea2Transformer2DModel(**KREA2_TRANSFORMER_CONFIG) - model.load_state_dict(sd, assign=True, strict=False) + load_state_dict_ignoring_extras(model, sd, source="Krea-2 GGUF checkpoint", assign=True, allow_missing=True) # Reject GGUF layouts that don't fully populate the diffusers Krea2Transformer2DModel (city96/ # ComfyUI GGUFs may use key names needing conversion). Failing here beats a confusing meta-tensor # crash mid-inference. @@ -504,28 +507,14 @@ def _remap_qwen3vl_singlefile_keys(sd: dict[str, Any]) -> dict[str, Any]: def _reject_incomplete_load(model: Any, *, what: str) -> None: - """Raise if a ``load_state_dict(strict=False)`` left required tensors on the meta device. + """Krea-2's alias for the shared meta-device completeness sweep. ``strict=False`` is used to tolerate benign extra/renamed keys, but it also silently accepts a checkpoint that omits required weights — those tensors stay on the meta device and only fail much later during inference. Reject such loads here, naming the offending tensors, so an incomplete, misidentified, or differently-converted checkpoint fails at load time with an actionable message. - - Both parameters *and persistent buffers* are checked: ``accelerate.init_empty_weights()`` places - buffers on the meta device too, so a native/GGUF checkpoint that omits a persistent buffer would - slip past a parameters-only guard and fail mid-inference instead of at load time. """ - still_meta = [ - name - for name, tensor in (*model.named_parameters(), *model.named_buffers()) - if getattr(tensor, "is_meta", False) - ] - if still_meta: - raise RuntimeError( - f"{what} is incomplete: {len(still_meta)} tensor(s) were not provided by the checkpoint " - f"and remain uninitialized (meta device). First few: {still_meta[:8]}. The file is likely " - "incomplete, misidentified, or uses a key layout that needs conversion." - ) + reject_incomplete_load(model, what=what) @ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.Qwen3VLEncoder, format=ModelFormat.Checkpoint) @@ -605,7 +594,9 @@ def _load_text_encoder(self, config: Qwen3VLEncoder_Checkpoint_Config) -> AnyMod for k in sd.keys(): sd[k] = sd[k].to(model_dtype) - model.load_state_dict(sd, assign=True, strict=False) + load_state_dict_ignoring_extras( + model, sd, source="Qwen3-VL encoder checkpoint", assign=True, allow_missing=True + ) _reject_incomplete_load(model, what="Qwen3-VL encoder checkpoint") # Keep an fp8 encoder running in fp8 (storage=float8_e4m3fn, per-layer upcast to the compute diff --git a/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py b/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py index 6aad71a71c3..b933aebd6ad 100644 --- a/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py +++ b/invokeai/backend/model_manager/load/model_loaders/mistral_encoder.py @@ -41,6 +41,7 @@ from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.util.devices import TorchDevice from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.util.state_dict_loading import log_unexpected_keys # Architecture constants for the 30-layer cow-mistral3-small distillation. # Sourced from BFL's FLUX.2-dev ``text_encoder/config.json`` (text-model side of @@ -953,8 +954,7 @@ def _load_text_encoder(self, config: MistralEncoder_Checkpoint_Config) -> AnyMod model = MistralModel(mistral_config) missing, unexpected = model.load_state_dict(sd, strict=False, assign=True) - if unexpected: - logger.debug(f"Mistral encoder: ignored {len(unexpected)} unexpected keys") + log_unexpected_keys("Mistral encoder checkpoint", unexpected) if missing: # Re-initialize any RMSNorm weights that may have been pruned during repackaging. for name in missing: @@ -1057,8 +1057,7 @@ def _load_from_gguf(self, config: MistralEncoder_GGUF_Config) -> AnyModel: model = MistralModel(mistral_config) missing, unexpected = model.load_state_dict(sd, strict=False, assign=True) - if unexpected: - logger.debug(f"Mistral encoder (GGUF): ignored {len(unexpected)} unexpected keys") + log_unexpected_keys("Mistral GGUF encoder", unexpected) if missing: logger.debug( f"Mistral encoder (GGUF): {len(missing)} keys missing from state dict (first 5: {missing[:5]})" diff --git a/invokeai/backend/model_manager/load/model_loaders/qwen_image.py b/invokeai/backend/model_manager/load/model_loaders/qwen_image.py index 1342cbe1075..f0b70263545 100644 --- a/invokeai/backend/model_manager/load/model_loaders/qwen_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/qwen_image.py @@ -33,6 +33,7 @@ from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras, log_unexpected_keys def _remap_qwen_vl_checkpoint_keys(sd: dict) -> dict: @@ -218,7 +219,9 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: with accelerate.init_empty_weights(): model = QwenImageTransformer2DModel(**model_config) - model.load_state_dict(sd, strict=False, assign=True) + load_state_dict_ignoring_extras( + model, sd, source="Qwen-Image transformer checkpoint", assign=True, allow_missing=True + ) return model @@ -285,7 +288,9 @@ def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: new_sd_size = sum(t.nelement() * t.element_size() for t in sd.values()) self._ram_cache.make_room(new_sd_size) - model.load_state_dict(sd, strict=False, assign=True) + load_state_dict_ignoring_extras( + model, sd, source="Qwen-Image transformer checkpoint", assign=True, allow_missing=True + ) return model @@ -443,11 +448,7 @@ def _load_text_encoder_from_singlefile(self, config: QwenVLEncoder_Checkpoint_Co # Load weights; allow missing keys for tied lm_head and re-initialised buffers. load_result = model.load_state_dict(sd, strict=False, assign=True) - if load_result.unexpected_keys: - logger.warning( - f"{len(load_result.unexpected_keys)} unexpected keys in checkpoint, " - f"first 5: {load_result.unexpected_keys[:5]}" - ) + log_unexpected_keys("Qwen2.5-VL text encoder checkpoint", load_result.unexpected_keys) # Tie lm_head ↔ embed_tokens if config requires it and lm_head wasn't loaded if getattr(qwen_config, "tie_word_embeddings", False): diff --git a/invokeai/backend/model_manager/load/model_loaders/vae.py b/invokeai/backend/model_manager/load/model_loaders/vae.py index 51497e725a0..947af36b5e0 100644 --- a/invokeai/backend/model_manager/load/model_loaders/vae.py +++ b/invokeai/backend/model_manager/load/model_loaders/vae.py @@ -26,6 +26,7 @@ ) from invokeai.backend.quantization.sdnq.detection import is_sdnq_folder from invokeai.backend.quantization.sdnq.loaders import raise_on_incomplete_sdnq_load, sdnq_sd_loader +from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras def _is_sdnq_vae_folder(path: Path) -> bool: @@ -247,7 +248,7 @@ def _load_wan_vae(self, config: VAE_Checkpoint_Wan_Config) -> AnyModel: with accelerate.init_empty_weights(): model = AutoencoderKLWan(**init_kwargs) - model.load_state_dict(sd, strict=True, assign=True) + load_state_dict_ignoring_extras(model, sd, source="Wan VAE checkpoint", assign=True) model.eval() return model @@ -302,7 +303,7 @@ def _load_qwen_image_vae(self, config: VAE_Checkpoint_QwenImage_Config) -> AnyMo with accelerate.init_empty_weights(): model = AutoencoderKLQwenImage() - model.load_state_dict(sd, strict=True, assign=True) + load_state_dict_ignoring_extras(model, sd, source="Qwen-Image VAE checkpoint", assign=True) model.eval() return model diff --git a/invokeai/backend/model_manager/load/model_loaders/z_image.py b/invokeai/backend/model_manager/load/model_loaders/z_image.py index ffda6eaf9cd..2a5f3c5dfe6 100644 --- a/invokeai/backend/model_manager/load/model_loaders/z_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/z_image.py @@ -39,6 +39,7 @@ from invokeai.backend.quantization.sdnq.loaders import raise_on_incomplete_sdnq_load, sdnq_sd_loader from invokeai.backend.qwen3.qwen3_tokenizer import load_bundled_qwen3_tokenizer from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.state_dict_loading import load_state_dict_ignoring_extras, log_unexpected_keys def _convert_z_image_gguf_to_diffusers(sd: dict[str, Any]) -> dict[str, Any]: @@ -357,7 +358,7 @@ def _load_sdnq_transformer(self, transformer_path: Path) -> AnyModel: axes_lens=[1024, 512, 512], ) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="Z-Image transformer", assign=True) return model @@ -481,7 +482,7 @@ def _load_from_singlefile( for k in sd.keys(): sd[k] = sd[k].to(model_dtype) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="Z-Image transformer checkpoint", assign=True) # `assign=True` aliases every param to its `sd` tensor, so the dict keeps the whole model # alive a second time. The FP8 cast below allocates the fp8 copy per param while the # `model_dtype` original is still reachable through `sd`, pushing peak RAM to ~1.5x what @@ -580,7 +581,7 @@ def _load_from_singlefile( axes_lens=[1024, 512, 512], ) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="Z-Image GGUF transformer checkpoint", assign=True) return model @@ -643,11 +644,9 @@ def _load_text_encoder(self, config: Main_SDNQ_Diffusers_ZImage_Config) -> AnyMo sd = sdnq_sd_loader(te_dir, compute_dtype=compute_dtype) # Qwen3ForCausalLM may share lm_head.weight with model.embed_tokens.weight; missing keys # for that tie are expected and handled by re-sharing post-load. - missing, unexpected = model.load_state_dict(sd, assign=True, strict=False) - if unexpected: - raise ValueError(f"Unexpected keys loading SDNQ Qwen3 text encoder: {unexpected}") - if missing and missing != ["lm_head.weight"]: - raise ValueError(f"Unexpected missing keys loading SDNQ Qwen3 text encoder: {missing}") + missing = load_state_dict_ignoring_extras( + model, sd, source="SDNQ Qwen3 text encoder", assign=True, allowed_missing={"lm_head.weight"} + ) if missing == ["lm_head.weight"]: model.lm_head.weight = model.model.embed_tokens.weight return model @@ -717,7 +716,7 @@ def _load_from_singlefile( axes_lens=[1024, 512, 512], ) - model.load_state_dict(sd, assign=True) + load_state_dict_ignoring_extras(model, sd, source="SDNQ Z-Image transformer checkpoint", assign=True) return model def _load_from_diffusers_folder( @@ -899,6 +898,7 @@ def _load_control_adapter( # Load state dict with strict=False to handle missing keys like x_pad_token # Some control adapters may not include x_pad_token in their checkpoint missing_keys, unexpected_keys = model.load_state_dict(sd, assign=True, strict=False) + log_unexpected_keys("Z-Image ControlNet checkpoint", unexpected_keys) # Initialize x_pad_token if it was missing from the checkpoint if "x_pad_token" in missing_keys: @@ -1124,7 +1124,9 @@ def _load_from_singlefile( # Load the text model weights from checkpoint # assign=True replaces meta tensors with real ones from state dict - model.load_state_dict(sd, strict=False, assign=True) + load_state_dict_ignoring_extras( + model, sd, source="Qwen3 text encoder checkpoint", assign=True, allow_missing=True + ) # Handle tied weights: lm_head shares weight with embed_tokens when tie_word_embeddings=True # This doesn't work automatically with init_empty_weights, so we need to manually tie them @@ -1324,7 +1326,7 @@ def _load_from_gguf( # Load the GGUF weights with assign=True # GGMLTensor wrappers will be dequantized on-the-fly during inference - model.load_state_dict(sd, strict=False, assign=True) + load_state_dict_ignoring_extras(model, sd, source="Qwen3 GGUF text encoder", assign=True, allow_missing=True) # Dequantize embed_tokens weight - embedding lookups require indexed access # which quantized GGMLTensors can't efficiently provide (no __torch_dispatch__ for embedding) diff --git a/invokeai/backend/pid/decode.py b/invokeai/backend/pid/decode.py index 4a9f4abe56d..55d7fa94ca8 100644 --- a/invokeai/backend/pid/decode.py +++ b/invokeai/backend/pid/decode.py @@ -31,6 +31,7 @@ from invokeai.backend.model_manager.taxonomy import BaseModelType from invokeai.backend.pid._src.networks.pid_net import PidNet from invokeai.backend.util.logging import InvokeAILogger +from invokeai.backend.util.state_dict_loading import log_unexpected_keys _PID_ACTIVATION_CHUNK_SIZE = 1024 @@ -268,16 +269,13 @@ def load_pid_decoder(state_dict: dict[Any, Tensor], backbone: BaseModelType) -> + (f" (+ {len(not_strings) - 5} more)" if len(not_strings) > 5 else "") ) - # strict=False so we can report missing and unexpected keys separately; both are fatal. The model - # cache builds loaders under `skip_torch_weight_init()`, which no-ops every `reset_parameters()`, - # so a key the checkpoint does not supply is left as uninitialised memory rather than a sane - # default — a partial checkpoint would decode to garbage / NaNs instead of failing. + # strict=False so we can report missing and unexpected keys separately. Missing keys are fatal: + # the model cache builds loaders under `skip_torch_weight_init()`, which no-ops every + # `reset_parameters()`, so a key the checkpoint does not supply is left as uninitialised memory + # rather than a sane default — a partial checkpoint would decode to garbage / NaNs instead of + # failing. Unexpected keys are exporter noise and are only logged (see `log_unexpected_keys`). missing, unexpected = net.load_state_dict(state_dict, strict=False) - if unexpected: - raise RuntimeError( - f"PiD checkpoint has unexpected keys not present in PidNet: {unexpected[:5]}" - + (f" (+ {len(unexpected) - 5} more)" if len(unexpected) > 5 else "") - ) + log_unexpected_keys("PiD checkpoint", unexpected) if missing: lq = [k for k in missing if k.startswith("lq_proj.")] detail = ( diff --git a/invokeai/backend/quantization/sdnq/loaders.py b/invokeai/backend/quantization/sdnq/loaders.py index 252ea8a6208..cfd0fd7c487 100644 --- a/invokeai/backend/quantization/sdnq/loaders.py +++ b/invokeai/backend/quantization/sdnq/loaders.py @@ -12,6 +12,7 @@ from invokeai.backend.quantization.sdnq.sdnq_tensor import SDNQTensor from invokeai.backend.quantization.sdnq.utils import SDNQQuantizationType +from invokeai.backend.util.state_dict_loading import log_unexpected_keys logger = logging.getLogger(__name__) @@ -25,23 +26,24 @@ def raise_on_incomplete_sdnq_load( """Fail fast when ``load_state_dict(..., strict=False)`` left an SDNQ model incomplete. ``load_state_dict`` with ``strict=False`` silently ignores the missing/unexpected key lists. - For SDNQ folder loads that is dangerous: a partial export, missing shard key or architecture - mismatch leaves required parameters on the meta device and returns a model that fails much later - during device movement or inference, far from the real cause. This raises with the offending - keys instead. + For SDNQ folder loads a missing key is dangerous: a partial export, missing shard key or + architecture mismatch leaves required parameters on the meta device and returns a model that + fails much later during device movement or inference, far from the real cause. This raises with + the offending keys instead. + + Unexpected keys are *not* an error - they are exporter noise rather than a correctness signal, + so they are logged at DEBUG and ignored (see ``log_unexpected_keys``). Args: model_name: Human-readable name for the error message (e.g. "SDNQ Z-Image transformer"). missing_keys: The ``missing_keys`` returned by ``load_state_dict``. - unexpected_keys: The ``unexpected_keys`` returned by ``load_state_dict``. + unexpected_keys: The ``unexpected_keys`` returned by ``load_state_dict``; logged, not fatal. allowed_missing: Keys that are expected to be absent (e.g. tied weights the caller re-shares after load), which must not trigger a failure. """ allowed = set(allowed_missing) real_missing = [k for k in missing_keys if k not in allowed] - unexpected = list(unexpected_keys) - if unexpected: - raise ValueError(f"Unexpected keys loading {model_name}: {unexpected}") + log_unexpected_keys(model_name, unexpected_keys) if real_missing: raise ValueError(f"Missing keys loading {model_name} (required parameters left on meta): {real_missing}") diff --git a/invokeai/backend/util/state_dict_loading.py b/invokeai/backend/util/state_dict_loading.py new file mode 100644 index 00000000000..82438a71faa --- /dev/null +++ b/invokeai/backend/util/state_dict_loading.py @@ -0,0 +1,145 @@ +"""One shared policy for the key lists that ``nn.Module.load_state_dict`` reports. + +Checkpoints in the wild routinely carry tensors the model has nowhere to put: runtime-derived +buffers an exporter happened to serialize, bookkeeping the training rig wrote out, leftovers from a +merge. None of them are weights the model needs, and none of them say anything about whether the +weights that *are* present are correct — the official checkpoint for the same architecture loads +fine without them. + +Historically each single-file loader made its own call about that, and the loaders that hard-failed +turned every such extra tensor into a user-facing crash that could only be cleared by adding one +more entry to an allowlist and cutting a release (see issue #9437). So the policy here is: + +* **Unexpected keys are ignored.** They are logged at ``DEBUG`` — visible when someone turns the log + level up to diagnose a load, silent otherwise — and never raise. + + The one deliberate exception in the codebase is + ``model_loaders/wan.py::_raise_for_incompatible_keys``, and it is worth knowing why it is not a + counter-example to the policy: several Wan 2.2 derivatives (Animate, S2V, Fun-Camera) are + *supersets* of the plain transformer, so their extra ``audio_injector``/``face_adapter``/ + ``control_adapter`` branches are not exporter noise but the entire feature the checkpoint exists + for, and running one without them generates silently degraded output. That loader strips the + genuinely benign extras first, so what reaches its check is a named-variant signal, not junk. +* **Missing keys are still an error** where the caller says the state dict is supposed to be + complete. A parameter the checkpoint never filled is a real defect: it stays on the meta device + (or, under ``skip_torch_weight_init()``, holds uninitialised memory) and blows up mid-inference, + far from the cause. +""" + +from collections.abc import Iterable, Mapping +from typing import Any + +import torch + +from invokeai.backend.util.logging import InvokeAILogger + +logger = InvokeAILogger.get_logger(__name__) + +# Enough keys to recognize what the extras are without dumping hundreds of lines for a bundled VAE. +MAX_REPORTED_KEYS = 10 + + +def _format_keys(keys: list[str]) -> str: + shown = keys[:MAX_REPORTED_KEYS] + suffix = f" (+{len(keys) - len(shown)} more)" if len(keys) > len(shown) else "" + return f"{shown}{suffix}" + + +def log_unexpected_keys(source: str, unexpected_keys: Iterable[Any]) -> None: + """Report keys that ``load_state_dict`` had nowhere to put, at DEBUG level only. + + Args: + source: Human-readable name of what was being loaded, e.g. "Anima transformer checkpoint". + unexpected_keys: The ``unexpected_keys`` reported by ``load_state_dict(strict=False)``. + """ + # Not every key is guaranteed to be a string: a `.pth` unpickles to whatever it contains. + keys = sorted(str(key) for key in unexpected_keys) + if not keys: + return + logger.debug(f"{source}: ignoring {len(keys)} key(s) not present in the model: {_format_keys(keys)}") + + +def reject_incomplete_load(model: torch.nn.Module, *, what: str) -> None: + """Raise if a ``load_state_dict(strict=False)`` left required tensors on the meta device. + + The completeness check for loaders that cannot use ``load_state_dict_ignoring_extras``'s + missing-key check — because they legitimately tolerate *some* missing keys, or because the key + list cannot see the problem. It is strictly stronger than ``missing_keys`` for models built + under ``accelerate.init_empty_weights()``: it is immune to non-persistent buffers (which never + appear in ``missing_keys`` at all) and to tied weights (materialized by ``tie_weights()`` after + the load), and it catches the real failure — a required tensor that was never filled and would + otherwise blow up mid-inference with "Cannot copy out of meta tensor" instead of at load time. + + Buffers are checked as well as parameters. ``accelerate.init_empty_weights()`` defaults to + ``include_buffers=False``, so a module's constructor normally materializes them — but a loader + that opts into ``include_buffers=True``, or that builds the module with ``to_empty()``, leaves + persistent buffers on meta, and a checkpoint omitting one would slip past a parameters-only + guard. + + Args: + model: The module that was just loaded into. + what: Human-readable name of what was loaded, used in the error message. + + Raises: + RuntimeError: If any parameter or buffer is still on the meta device. + """ + still_meta = [ + name + for name, tensor in (*model.named_parameters(), *model.named_buffers()) + if getattr(tensor, "is_meta", False) + ] + if still_meta: + raise RuntimeError( + f"{what} is incomplete: {len(still_meta)} tensor(s) were not provided by the checkpoint " + f"and remain uninitialized (meta device). First few: {still_meta[:8]}. The file is likely " + "incomplete, misidentified, or uses a key layout that needs conversion." + ) + + +def load_state_dict_ignoring_extras( + model: torch.nn.Module, + state_dict: Mapping[str, Any], + *, + source: str, + assign: bool = False, + allow_missing: bool = False, + allowed_missing: Iterable[str] = (), +) -> list[str]: + """Load ``state_dict`` into ``model``, ignoring keys the model has nowhere to put. + + A drop-in replacement for ``model.load_state_dict(state_dict, strict=True)`` that keeps the + strictness that matters (every required parameter must be filled) and drops the strictness that + only produces whack-a-mole (the checkpoint must contain nothing else). Shape mismatches still + raise from torch, exactly as they do under ``strict=True``. + + Args: + model: The module to load into. + state_dict: The state dict to load. + source: Human-readable name of what is being loaded, used in the log line and the error. + assign: Passed through to ``load_state_dict``. + allow_missing: When True, missing keys are tolerated too — for callers that fill them in + afterwards (tied weights, re-initialised buffers) or that run their own completeness + check, such as a sweep for tensors left on the meta device. + allowed_missing: Specific keys that are expected to be absent and must not raise. Ignored + when ``allow_missing`` is True. + + Returns: + The missing keys, so callers can materialize what they said they would. + + Raises: + RuntimeError: If a required key was missing from ``state_dict``. + """ + incompatible_keys = model.load_state_dict(state_dict, strict=False, assign=assign) + log_unexpected_keys(source, incompatible_keys.unexpected_keys) + + missing_keys = [str(key) for key in incompatible_keys.missing_keys] + if not allow_missing: + allowed = set(allowed_missing) + required_missing = sorted(key for key in missing_keys if key not in allowed) + if required_missing: + raise RuntimeError( + f"{source} is missing {len(required_missing)} parameter(s) that the model requires: " + f"{_format_keys(required_missing)}. The checkpoint is likely incomplete, misidentified, " + "or uses a key layout that needs conversion." + ) + return missing_keys diff --git a/tests/backend/ideogram4/test_quantized_loading.py b/tests/backend/ideogram4/test_quantized_loading.py index fa5e9ad154c..75b405ccb5f 100644 --- a/tests/backend/ideogram4/test_quantized_loading.py +++ b/tests/backend/ideogram4/test_quantized_loading.py @@ -7,6 +7,8 @@ tiny CPU model so the fp8 path has regression coverage without a multi-GB checkpoint. """ +import logging + import accelerate import pytest import torch @@ -102,8 +104,9 @@ def test_fp8_load_matches_loader_pattern() -> None: assert torch.allclose(out, expected, atol=1e-5, rtol=1e-4) -def test_fp8_load_rejects_unexpected_keys() -> None: - """A key the model has no home for must fail loudly rather than load silently.""" +def test_fp8_load_ignores_unexpected_keys(caplog: pytest.LogCaptureFixture) -> None: + """A key the model has no home for is exporter noise: the load succeeds and only DEBUG says + anything about it (issue #9437).""" torch.manual_seed(1) compute_dtype = torch.float32 ref = _TinyEncoder().to(compute_dtype).eval() @@ -112,9 +115,15 @@ def test_fp8_load_rejects_unexpected_keys() -> None: model = _TinyEncoder().to(compute_dtype) swap_linears_to_fp8(model, sd, compute_dtype=compute_dtype) - with pytest.raises(RuntimeError, match="unexpected keys"): + with caplog.at_level(logging.DEBUG, logger="invokeai.backend.util.state_dict_loading"): load_fp8_state_dict(model, sd, device=torch.device("cpu"), dtype=compute_dtype, strict=False) + assert "lin1.bogus_extra" in caplog.text + + x = torch.randn(2, 8, dtype=compute_dtype) + with torch.no_grad(): + assert torch.allclose(model(x), _dequant_reference(ref, sd, x), atol=1e-5, rtol=1e-4) + def test_fp8_missing_key_strictness() -> None: """strict=True raises on a missing weight; strict=False downgrades it to a warning.""" diff --git a/tests/backend/model_manager/configs/test_pid_decoder_config.py b/tests/backend/model_manager/configs/test_pid_decoder_config.py index 62d1045aabd..4449cd526e1 100644 --- a/tests/backend/model_manager/configs/test_pid_decoder_config.py +++ b/tests/backend/model_manager/configs/test_pid_decoder_config.py @@ -173,14 +173,29 @@ def test_a_missing_backbone_weight_is_rejected(self, dropped: str) -> None: with pytest.raises(InvalidMatchError, match="missing 1 of the weights required by PidNet"): PiDDecoder_Checkpoint_FLUX_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) - def test_an_unexpected_key_is_rejected(self) -> None: - """`load_pid_decoder` refuses these too, so accepting them here would install a file that - cannot load.""" + def test_an_unexpected_key_is_accepted(self) -> None: + """Since issue #9437 `load_pid_decoder` ignores extra keys, so rejecting them here would + refuse to install a file that loads fine — the installer must track the loader in both + directions, not just the strict one.""" sd = _pid_state_dict() sd[f"{_NET_PREFIX}not_a_pid_key"] = _FakeShapeTensor(1) with TemporaryDirectory() as tmpdir: mod = _mock_mod(Path(tmpdir), sd) - with pytest.raises(InvalidMatchError, match="1 keys PidNet does not expect"): + config = PiDDecoder_Checkpoint_FLUX_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) + assert config.base is BaseModelType.Flux + + def test_a_non_string_key_is_rejected(self) -> None: + """The one extra key the loader genuinely cannot survive: `nn.Module.load_state_dict` calls + `.startswith()` on every key, so a non-string one raises from inside torch. Identification + has to keep refusing these even though it now accepts ordinary extras. + + Built bare, without the `net.` prefix: `strip_net_prefix` drops non-string keys when it has a + prefix to strip, so only a bare checkpoint can carry one this far.""" + sd: dict[Any, Any] = {k: _FakeShapeTensor(*shape) for k, shape in required_pid_net_shapes().items()} + sd[1] = _FakeShapeTensor(1) + with TemporaryDirectory() as tmpdir: + mod = _mock_mod(Path(tmpdir), sd) + with pytest.raises(InvalidMatchError, match="1 keys that are not strings"): PiDDecoder_Checkpoint_FLUX_Config.from_model_on_disk(mod, dict(_OVERRIDE_FIELDS)) def test_a_wrong_shaped_weight_is_rejected(self) -> None: @@ -312,10 +327,12 @@ def _bare_with_a_non_string_key(self) -> dict[Any, object]: """A complete *bare* contract plus two keys PidNet does not expect, one of them not a string. A bare checkpoint is passed through `strip_net_prefix` untouched, so a `.pth` can hand - identification whatever it was pickled with. Reporting the unexpected keys sorts them, and - sorting `{1, "not_a_pid_key"}` raises TypeError — which the factory catches as a generic - candidate failure and answers with the Unknown_Config registration this class is about. A - crash in an unusability check therefore does not fail loudly; it fails as a silent accept. + identification whatever it was pickled with. Since issue #9437 the plain `not_a_pid_key` + is accepted (the loader ignores it), but `1` is not: `load_state_dict` calls `.startswith()` + on every key. Reporting it sorts the offenders, and sorting a mixed set raises TypeError — + which the factory catches as a generic candidate failure and answers with the Unknown_Config + registration this class is about. A crash in an unusability check therefore does not fail + loudly; it fails as a silent accept, so the sort stays `key=str`. """ scalar = torch.zeros(()) sd: dict[Any, object] = {k: scalar.expand(shape) for k, shape in required_pid_net_shapes().items()} @@ -332,7 +349,7 @@ def _bare_with_a_non_string_key(self) -> dict[Any, object]: ("_intact_v1_5", "lq_proj hidden dim 1024"), ("_unsupported_latent_channels", "32 latent channels"), ("_malformed_discriminator", "malformed lq_proj.latent_proj.0.weight"), - ("_bare_with_a_non_string_key", "2 keys PidNet does not expect"), + ("_bare_with_a_non_string_key", "1 keys that are not strings"), ], ) def test_factory_returns_no_config_even_with_allow_unknown(self, case: str, expected_reason: str) -> None: diff --git a/tests/backend/model_manager/configs/test_qwen3_encoder_sdnq_single_file_identification.py b/tests/backend/model_manager/configs/test_qwen3_encoder_sdnq_single_file_identification.py index 3d5f443c3f8..c6a39e17585 100644 --- a/tests/backend/model_manager/configs/test_qwen3_encoder_sdnq_single_file_identification.py +++ b/tests/backend/model_manager/configs/test_qwen3_encoder_sdnq_single_file_identification.py @@ -6,9 +6,14 @@ `Qwen3ForCausalLM`, identification must reject anything it cannot load: - a Qwen2 causal LM (no Qwen3 QK-norm params -> missing weights), and -- a Qwen-VL model (bundles a visual tower -> unexpected weights), +- a Qwen-VL model, whose `visual.*` tower is accompanied by a Qwen2-style attention block that is + likewise missing the Qwen3 QK-norm params. while still accepting a genuine Qwen3 checkpoint carrying its q_norm/k_norm parameters. + +Note that the visual tower itself is no longer a rejection signal: since issue #9437 the loader +ignores unexpected keys, so identification must not reject on them either. Missing weights are what +both the loader and identification still refuse. """ from pathlib import Path diff --git a/tests/backend/model_manager/load/test_gemma2_encoder_gguf_loader.py b/tests/backend/model_manager/load/test_gemma2_encoder_gguf_loader.py index 21b5ace1393..beeef8c155b 100644 --- a/tests/backend/model_manager/load/test_gemma2_encoder_gguf_loader.py +++ b/tests/backend/model_manager/load/test_gemma2_encoder_gguf_loader.py @@ -12,6 +12,7 @@ """ import importlib +import logging import os from pathlib import Path from typing import Any @@ -239,19 +240,37 @@ def test_forward_runs_with_quantized_weights(self, tiny_gguf: dict[str, GGMLTens # The projections were not dequantized in place by running the model. assert isinstance(model.layers[0].self_attn.q_proj.weight, GGMLTensor) - def test_unexpected_tensor_is_rejected( - self, monkeypatch: pytest.MonkeyPatch, tiny_gguf: dict[str, GGMLTensor] + def test_unexpected_tensor_is_ignored_and_logged_at_debug( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, tiny_gguf: dict[str, GGMLTensor] ) -> None: """A tensor that maps cleanly but has no home in the configured model (here: a block beyond - `num_hidden_layers`) must fail loudly instead of being dropped.""" + `num_hidden_layers`) is exporter noise, not a broken checkpoint: the load must succeed and + say so only at DEBUG (issue #9437).""" sd = dict(tiny_gguf) | {f"blk.{_LAYERS}.attn_q.weight": _q8(_HEADS * _HEAD_DIM, _HIDDEN)} monkeypatch.setattr( invokeai.backend.quantization.gguf.loaders, "gguf_sd_loader", lambda path, compute_dtype: sd ) - with pytest.raises(RuntimeError, match="Unexpected keys"): + with caplog.at_level(logging.DEBUG, logger="invokeai.backend.util.state_dict_loading"): + model = load_gemma2_model_from_gguf(Path("unused.gguf"), torch.float32) + + assert not any(p.is_meta for p in model.parameters()) + assert f"layers.{_LAYERS}.self_attn.q_proj.weight" in caplog.text + + def test_unexpected_tensor_is_silent_above_debug( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, tiny_gguf: dict[str, GGMLTensor] + ) -> None: + """...and nothing at all at the default log level.""" + sd = dict(tiny_gguf) | {f"blk.{_LAYERS}.attn_q.weight": _q8(_HEADS * _HEAD_DIM, _HIDDEN)} + monkeypatch.setattr( + invokeai.backend.quantization.gguf.loaders, "gguf_sd_loader", lambda path, compute_dtype: sd + ) + + with caplog.at_level(logging.INFO, logger="invokeai.backend.util.state_dict_loading"): load_gemma2_model_from_gguf(Path("unused.gguf"), torch.float32) + assert caplog.text == "" + @pytest.mark.skipif(not os.environ.get(_LOCAL_GGUF_ENV_VAR), reason=f"set {_LOCAL_GGUF_ENV_VAR} to a Gemma-2-2b GGUF") def test_native_gguf_load_keeps_projections_quantized_and_matches_reference() -> None: diff --git a/tests/backend/pid/test_pid_decode.py b/tests/backend/pid/test_pid_decode.py index d051e32d4c3..fb901a8e2dd 100644 --- a/tests/backend/pid/test_pid_decode.py +++ b/tests/backend/pid/test_pid_decode.py @@ -1,5 +1,6 @@ """Regression tests for the PiD distill schedule, decoder/base validation and checkpoint completeness.""" +import logging import math from typing import Any from unittest.mock import patch @@ -331,10 +332,23 @@ def test_missing_backbone_keys_are_rejected(self, tiny_net: torch.nn.Module) -> with pytest.raises(RuntimeError, match="missing 1 keys"): load_pid_decoder(sd, BaseModelType.Flux) - def test_unexpected_keys_are_rejected(self, tiny_net: torch.nn.Module) -> None: + def test_unexpected_keys_are_ignored_and_logged_at_debug( + self, tiny_net: torch.nn.Module, caplog: pytest.LogCaptureFixture + ) -> None: + """An extra key says nothing about whether the weights that *are* present are right, so the + checkpoint still loads and the extras are reported only at DEBUG (issue #9437).""" sd = dict(tiny_net.state_dict()) | {"not_a_pid_key": torch.zeros(1)} - with pytest.raises(RuntimeError, match="unexpected keys"): + with caplog.at_level(logging.DEBUG, logger="invokeai.backend.util.state_dict_loading"): + assert load_pid_decoder(sd, BaseModelType.Flux) is tiny_net + assert "not_a_pid_key" in caplog.text + + def test_unexpected_keys_are_silent_above_debug( + self, tiny_net: torch.nn.Module, caplog: pytest.LogCaptureFixture + ) -> None: + sd = dict(tiny_net.state_dict()) | {"not_a_pid_key": torch.zeros(1)} + with caplog.at_level(logging.INFO, logger="invokeai.backend.util.state_dict_loading"): load_pid_decoder(sd, BaseModelType.Flux) + assert caplog.text == "" def test_non_string_keys_are_rejected_before_torch_sees_them(self, tiny_net: torch.nn.Module) -> None: """A bare checkpoint keeps whatever keys the `.pth` was pickled with (see `strip_net_prefix`), diff --git a/tests/backend/pid/test_pid_state_dict_utils.py b/tests/backend/pid/test_pid_state_dict_utils.py index df161edd537..4aca9930522 100644 --- a/tests/backend/pid/test_pid_state_dict_utils.py +++ b/tests/backend/pid/test_pid_state_dict_utils.py @@ -28,7 +28,8 @@ def test_distill_only_submodules_are_dropped(self) -> None: def test_a_bare_pid_net_checkpoint_is_untouched(self) -> None: """Without the prefix there is no evidence this is a distill serialisation, so nothing is - filtered — a stray key should reach the loader's "unexpected keys" check, not vanish.""" + filtered — a stray key should reach the loader, which reports it at DEBUG and ignores it + (issue #9437), rather than vanish here where nothing could report it at all.""" sd = {"lq_proj.a": torch.zeros(1), "discriminator.y": torch.zeros(1)} assert strip_net_prefix(sd) is sd diff --git a/tests/backend/quantization/sdnq/test_sdnq_loader.py b/tests/backend/quantization/sdnq/test_sdnq_loader.py index a02816ce6a4..af5594e0375 100644 --- a/tests/backend/quantization/sdnq/test_sdnq_loader.py +++ b/tests/backend/quantization/sdnq/test_sdnq_loader.py @@ -1,5 +1,6 @@ """Integration tests for SDNQ state dict loader.""" +import logging from pathlib import Path import pytest @@ -237,9 +238,17 @@ def test_missing_required_transformer_weight_raises(self): with pytest.raises(ValueError, match="bias"): raise_on_incomplete_sdnq_load("SDNQ test transformer", missing, unexpected) - def test_unexpected_key_raises(self): - with pytest.raises(ValueError, match="not_a_real_param"): + def test_unexpected_key_is_logged_not_raised(self, caplog): + """Extra keys are exporter noise, so they are reported at DEBUG and never fail the load + (issue #9437). Only completeness — a required parameter left on meta — is fatal.""" + with caplog.at_level(logging.DEBUG, logger="invokeai.backend.util.state_dict_loading"): raise_on_incomplete_sdnq_load("SDNQ test", missing_keys=[], unexpected_keys=["not_a_real_param"]) + assert "not_a_real_param" in caplog.text + + def test_unexpected_key_is_silent_above_debug(self, caplog): + with caplog.at_level(logging.INFO, logger="invokeai.backend.util.state_dict_loading"): + raise_on_incomplete_sdnq_load("SDNQ test", missing_keys=[], unexpected_keys=["not_a_real_param"]) + assert caplog.text == "" def test_allowed_missing_is_tolerated(self): # Tied/re-shared weights (e.g. T5 encoder.embed_tokens, Qwen3 lm_head) must not raise. diff --git a/tests/backend/util/test_state_dict_loading.py b/tests/backend/util/test_state_dict_loading.py new file mode 100644 index 00000000000..3ef6e0a0202 --- /dev/null +++ b/tests/backend/util/test_state_dict_loading.py @@ -0,0 +1,203 @@ +"""Tests for the shared unexpected/missing key policy used by the single-file model loaders. + +Issue #9437: loaders used to disagree about what an extra key in a checkpoint means, and the ones +that hard-failed turned exporter noise (`model_sampling.sigmas`, `pos_embedder.seq`, ...) into a +release-blocking crash. The policy pinned here is: extra keys are reported at DEBUG and ignored; +a *missing* required parameter is still an error. +""" + +import logging + +import pytest +import torch + +from invokeai.backend.util.state_dict_loading import ( + MAX_REPORTED_KEYS, + load_state_dict_ignoring_extras, + log_unexpected_keys, + reject_incomplete_load, +) + +LOGGER_NAME = "invokeai.backend.util.state_dict_loading" + + +class _TinyNet(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin = torch.nn.Linear(4, 4) + self.register_buffer("scale", torch.ones(4)) + + +@pytest.fixture +def net() -> _TinyNet: + return _TinyNet() + + +@pytest.fixture +def full_sd(net: _TinyNet) -> dict[str, torch.Tensor]: + return {k: v.clone() for k, v in net.state_dict().items()} + + +class TestLogUnexpectedKeys: + def test_reports_at_debug(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.DEBUG, logger=LOGGER_NAME): + log_unexpected_keys("tiny checkpoint", ["model_sampling.sigmas"]) + + assert len(caplog.records) == 1 + assert caplog.records[0].levelno == logging.DEBUG + assert "tiny checkpoint" in caplog.text + assert "model_sampling.sigmas" in caplog.text + + def test_silent_above_debug(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.INFO, logger=LOGGER_NAME): + log_unexpected_keys("tiny checkpoint", ["model_sampling.sigmas"]) + + assert caplog.records == [] + + def test_no_keys_logs_nothing(self, caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.DEBUG, logger=LOGGER_NAME): + log_unexpected_keys("tiny checkpoint", []) + + assert caplog.records == [] + + def test_long_lists_are_truncated(self, caplog: pytest.LogCaptureFixture) -> None: + """A bundled VAE contributes hundreds of extras; the log line must stay readable.""" + keys = [f"vae.decoder.block_{i:03d}.weight" for i in range(MAX_REPORTED_KEYS + 25)] + with caplog.at_level(logging.DEBUG, logger=LOGGER_NAME): + log_unexpected_keys("tiny checkpoint", keys) + + assert f"ignoring {len(keys)} key(s)" in caplog.text + assert "(+25 more)" in caplog.text + assert keys[MAX_REPORTED_KEYS] not in caplog.text + + def test_non_string_keys_do_not_crash_the_report(self, caplog: pytest.LogCaptureFixture) -> None: + """A `.pth` unpickles to whatever it contains, so a key need not be a string.""" + with caplog.at_level(logging.DEBUG, logger=LOGGER_NAME): + log_unexpected_keys("tiny checkpoint", [1, "real.key"]) + + assert "real.key" in caplog.text + + +class TestLoadStateDictIgnoringExtras: + def test_extra_key_loads_and_only_logs( + self, net: _TinyNet, full_sd: dict[str, torch.Tensor], caplog: pytest.LogCaptureFixture + ) -> None: + """The Anima case: an exporter serialized a tensor the model has no slot for.""" + sd = full_sd | {"model_sampling.sigmas": torch.zeros(3)} + + with caplog.at_level(logging.DEBUG, logger=LOGGER_NAME): + missing = load_state_dict_ignoring_extras(net, sd, source="tiny checkpoint") + + assert missing == [] + assert "model_sampling.sigmas" in caplog.text + torch.testing.assert_close(net.lin.weight, full_sd["lin.weight"]) + + def test_missing_key_still_raises(self, net: _TinyNet, full_sd: dict[str, torch.Tensor]) -> None: + """Completeness is the invariant worth enforcing — a required parameter never filled would + otherwise blow up mid-inference instead of at load time.""" + del full_sd["lin.bias"] + + with pytest.raises(RuntimeError, match="lin.bias"): + load_state_dict_ignoring_extras(net, full_sd, source="tiny checkpoint") + + def test_allowed_missing_is_tolerated(self, net: _TinyNet, full_sd: dict[str, torch.Tensor]) -> None: + """Tied weights the caller re-shares after the load (T5's encoder.embed_tokens, Qwen3's + lm_head) are legitimately absent from the file.""" + del full_sd["lin.bias"] + + missing = load_state_dict_ignoring_extras(net, full_sd, source="tiny checkpoint", allowed_missing={"lin.bias"}) + + assert missing == ["lin.bias"] + + def test_allow_missing_defers_to_the_caller(self, net: _TinyNet, full_sd: dict[str, torch.Tensor]) -> None: + """Callers that run their own completeness check (a sweep for tensors left on the meta + device) opt out of the missing-key error entirely.""" + del full_sd["lin.bias"] + + missing = load_state_dict_ignoring_extras(net, full_sd, source="tiny checkpoint", allow_missing=True) + + assert missing == ["lin.bias"] + + def test_shape_mismatch_still_raises(self, net: _TinyNet, full_sd: dict[str, torch.Tensor]) -> None: + """Dropping `strict=True` must not cost the size check that came with it.""" + full_sd["lin.weight"] = torch.zeros(8, 8) + + with pytest.raises(RuntimeError, match="size mismatch"): + load_state_dict_ignoring_extras(net, full_sd, source="tiny checkpoint") + + def test_assign_is_passed_through(self, full_sd: dict[str, torch.Tensor]) -> None: + """`assign=True` is how every meta-device loader materializes its parameters.""" + import accelerate + + with accelerate.init_empty_weights(): + model = _TinyNet() + assert model.lin.weight.is_meta + + load_state_dict_ignoring_extras(model, full_sd, source="tiny checkpoint", assign=True) + + assert not any(t.is_meta for t in (*model.parameters(), *model.buffers())) + + +class TestRejectIncompleteLoad: + """The completeness half of the policy, for loaders that cannot use the missing-key check — + either because they legitimately tolerate some missing keys, or because the key list cannot see + the problem at all.""" + + def test_raises_when_a_parameter_is_left_on_meta(self) -> None: + import accelerate + + with accelerate.init_empty_weights(): + model = _TinyNet() + + with pytest.raises(RuntimeError, match="tiny checkpoint is incomplete"): + reject_incomplete_load(model, what="tiny checkpoint") + + def test_names_the_tensors_the_checkpoint_did_not_fill(self, full_sd: dict[str, torch.Tensor]) -> None: + import accelerate + + with accelerate.init_empty_weights(): + model = _TinyNet() + del full_sd["lin.bias"] + model.load_state_dict(full_sd, strict=False, assign=True) + + with pytest.raises(RuntimeError, match="lin.bias") as exc_info: + reject_incomplete_load(model, what="tiny checkpoint") + assert "1 tensor(s)" in str(exc_info.value) + + def test_does_not_raise_for_a_fully_materialized_model(self, full_sd: dict[str, torch.Tensor]) -> None: + import accelerate + + with accelerate.init_empty_weights(): + model = _TinyNet() + model.load_state_dict(full_sd, strict=False, assign=True) + + reject_incomplete_load(model, what="tiny checkpoint") + + def test_a_persistent_buffer_left_on_meta_is_caught(self) -> None: + """A parameters-only sweep would miss this; `strict=False` reports the buffer as missing but + several callers deliberately tolerate missing keys.""" + model = _TinyNet() + model.scale = torch.empty(4, device="meta") + + with pytest.raises(RuntimeError, match="scale"): + reject_incomplete_load(model, what="tiny checkpoint") + + def test_a_non_persistent_buffer_is_not_a_false_positive(self) -> None: + """The case `missing_keys` cannot see and this sweep must not invent: a buffer registered + `persistent=False` (Anima has three) never appears in the state dict, and its module\'s + constructor materializes it because `init_empty_weights` defaults to `include_buffers=False`. + """ + import accelerate + + class _WithNonPersistentBuffer(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin = torch.nn.Linear(4, 4) + self.register_buffer("inv_freq", torch.ones(2), persistent=False) + + with accelerate.init_empty_weights(): + model = _WithNonPersistentBuffer() + assert "inv_freq" not in model.state_dict() + model.load_state_dict({k: torch.zeros(v.shape) for k, v in model.state_dict().items()}, assign=True) + + reject_incomplete_load(model, what="tiny checkpoint") From 20a0d3a50fe457feddb82314a3ad4ad93081e9c3 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 14 Sep 2026 21:12:37 -0400 Subject: [PATCH 2/2] docs(model loaders): fix two stale unexpected-key comments Review follow-up for #9581. The ideogram4 fp8 encoder branch and the z_image SDNQ Qwen3 loader still described unexpected checkpoint keys as fatal, but both now go through helpers that log them at DEBUG and ignore them (`load_fp8_state_dict`, `raise_on_incomplete_sdnq_load`). Update the comments to match the code; missing-key strictness is unchanged. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01NxYxw5RUtuRTBRWTn79M9d --- .../backend/model_manager/load/model_loaders/ideogram4.py | 3 ++- .../backend/model_manager/load/model_loaders/z_image.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/invokeai/backend/model_manager/load/model_loaders/ideogram4.py b/invokeai/backend/model_manager/load/model_loaders/ideogram4.py index a467b400ce0..24fffa99ffe 100644 --- a/invokeai/backend/model_manager/load/model_loaders/ideogram4.py +++ b/invokeai/backend/model_manager/load/model_loaders/ideogram4.py @@ -185,7 +185,8 @@ def _load_text_encoder(self, model_path: Path) -> AnyModel: # Weight-only fp8 (e4m3): build the empty architecture, swap the quantized Linears for # Fp8Linear (gated on a saved per-row scale), then load. Mirrors the transformer fp8 branch; # runs on any device. strict=False tolerates the tied embed weights transformers resolves - # itself; unexpected keys still raise. assign=True fills the meta params directly. + # itself; unexpected keys are logged at DEBUG and ignored (see `load_fp8_state_dict`). + # assign=True fills the meta params directly. with accelerate.init_empty_weights(): model: torch.nn.Module = AutoModel.from_config(cfg) swap_linears_to_fp8(model, sd, compute_dtype=compute_dtype) diff --git a/invokeai/backend/model_manager/load/model_loaders/z_image.py b/invokeai/backend/model_manager/load/model_loaders/z_image.py index 2a5f3c5dfe6..a01c854ccf0 100644 --- a/invokeai/backend/model_manager/load/model_loaders/z_image.py +++ b/invokeai/backend/model_manager/load/model_loaders/z_image.py @@ -1592,9 +1592,9 @@ def _load_from_sdnq( model = Qwen3ForCausalLM(qwen_config) # Load the SDNQ weights with assign=True. lm_head is tied to embed_tokens (re-shared below), - # so it is expected to be missing; any other missing key or any unexpected key (e.g. from an - # incompatible or contaminated export) must fail here. The later meta-parameter guard only - # catches missing required params, not unexpected ones. + # so it is expected to be missing; any other missing key (e.g. from a partial or incompatible + # export) must fail here. Unexpected keys are exporter noise and are only logged at DEBUG + # (see `raise_on_incomplete_sdnq_load`). missing, unexpected = model.load_state_dict(sd, strict=False, assign=True) raise_on_incomplete_sdnq_load("SDNQ Qwen3 encoder", missing, unexpected, allowed_missing={"lm_head.weight"})