fix(model loaders): ignore unexpected checkpoint keys, report them at debug only

[lstein]

Summary

Closes #9437, implementing the resolution from the 2026-08-03 comment: unexpected keys in a checkpoint are reported at DEBUG and otherwise ignored silently. They never fail a load.

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.

What changed

New invokeai/backend/util/state_dict_loading.py holds the single policy:

helper	contract
log_unexpected_keys(source, keys)	DEBUG only, never raises
load_state_dict_ignoring_extras(model, sd, ...)	drop-in for strict=True: keeps the strictness that matters (every required parameter must be filled; shape mismatches still raise from torch) and drops the strictness that only produces whack-a-mole
reject_incomplete_load(model, what=...)	the meta-device completeness sweep, generalized out of krea2.py (which now delegates to it)

~35 call sites are routed through it across anima, flux, z_image, krea2, qwen_image, vae, ideogram4, gemma2_encoder and mistral_encoder, plus pid/decode.py, quantization/sdnq/loaders.py and ideogram4/quantized_loading.py.

Every previously existing missing-key guard is preserved exactly; only the unexpected-key policy changed. All eight sites that now pass allow_missing=True were already strict=False before, so nothing was silently weakened. flux.py's bare assert len(unexpected_keys) == 0 — which carried no message and was stripped entirely under python -O — goes with it.

Two consequences worth reviewing carefully

1. The identification gate had to move with the loader. configs/pid_decoder.py rejected unexpected keys at identification time, deliberately mirroring the loader ("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"). Relaxing only pid/decode.py would have left the PiD relaxation unreachable — the installer would still refuse a file that now loads fine, so #9437 would not have been fixed for PiD at all. The gate now refuses only non-string keys, which load_state_dict genuinely cannot survive (it calls .startswith() on every key).

2. Anima's unexpected-key RuntimeError was its only hard load-time guard — missing keys were merely a warning there. Dropping it outright would have turned an incomplete checkpoint from a load-time error into a Cannot copy out of meta tensor crash mid-inference, so it is replaced with reject_incomplete_load rather than removed. Verified empirically on accelerate 1.14 that this cannot false-positive: init_empty_weights() defaults to include_buffers=False, so AnimaTransformer's 685 parameters are on meta but its three persistent=False buffers are materialized by the constructor.

Deliberately not converted

wan.py::_raise_for_incompatible_keys still raises on unexpected keys, and I'd argue it should. Wan 2.2 derivatives (Animate, S2V, Fun-Camera) are supersets of the plain transformer whose extra audio_injector / face_adapter / control_adapter branches are the entire feature the checkpoint exists for, not exporter noise — they report zero missing keys and would generate with that conditioning silently absent. That loader also strips the genuinely benign extras (bundled VAE/text-encoder weights, merged-LoRA residue) before the check, so what reaches it is a named-variant signal. #9437's own table likewise lists Wan only under missing-key handling. The new module docstring records this exception and why it isn't a counter-example. Happy to convert it if the team disagrees.

Testing

• New tests/backend/util/test_state_dict_loading.py covers all three helpers, including the two cases the old missing_keys-based reasoning got wrong: a non-persistent buffer must not be a false positive, and a persistent buffer left on meta must be caught.
• Four existing tests that pinned the old raising behaviour were inverted to pin the new one, each paired with a "silent above DEBUG" case so the log level is load-bearing, not just the log line.
• Mutation-checked: logging at WARNING instead of DEBUG fails 4 tests; re-rejecting unexpected keys at PiD identification fails 2.
• Full suite: 5874 passed / 5 failed, and those 5 are identical to what unmodified main produces on my rig (4 test_flux2_working_memory, 1 test_pid_chunked_equivalence — both ROCm-environment, unrelated to this change).

Adversarial review

The diff was put through a fresh-context adversarial review whose brief was to prove the "missing-key strictness is preserved exactly" invariant false. It confirmed the invariant across all 32 converted call sites (tied weights, non-persistent buffers, GGUF deliberate-later-fill, init_empty_weights + assign=True, return-value comparisons, error-message matching, import cycles) and surfaced the two issues above, both of which are fixed here.

🤖 Generated with Claude Code

https://claude.ai/code/session_01V4nLTdWBQaaDUFVLx9TyZH

[mohithvardhan002]

I found one potential issue with the new state-dict loading flow. Please take a look at the duplicate load mentioned above.