headroomlabs-ai/headroom · error · RuntimeError

{model_id} {section}: state_dict mismatch against {type(subm

Error message

{model_id} {section}: state_dict mismatch against {type(submodule).__name__} (missing={list(missing)[:5]}, unexpected={list(unexpected)[:5]}). The checkpoint no longer matches HeadroomCompressorModel's architecture.

What it means

Raised by _load_merged_state_dict when submodule.load_state_dict(..., strict=False) reports missing or unexpected keys for one of the three sections — the checkpoint's tensor names no longer line up with HeadroomCompressorModel's submodules (encoder/token_head/span_conv). It surfaces at most 5 keys each way and states plainly that the checkpoint no longer matches the model architecture, typically after a code-side architecture change versus an older checkpoint.

Source

Thrown at headroom/transforms/kompress_compressor.py:803

    """
    import torch

    ckpt = torch.load(ckpt_path, map_location="cpu")
    missing_sections = [k for k in _MERGED_CHECKPOINT_KEYS if k not in ckpt]
    if missing_sections:
        raise RuntimeError(
            f"merged.pt for {model_id} is missing {missing_sections}; found keys: "
            f"{sorted(ckpt)}. This checkpoint format is not what the loader expects."
        )

    for section, submodule in (
        ("encoder_state_dict", model.encoder),
        ("token_head_state_dict", model.token_head),
        ("span_conv_state_dict", model.span_conv),
    ):
        missing, unexpected = submodule.load_state_dict(ckpt[section], strict=False)
        if missing or unexpected:
            raise RuntimeError(
                f"{model_id} {section}: state_dict mismatch against {type(submodule).__name__} "
                f"(missing={list(missing)[:5]}, unexpected={list(unexpected)[:5]}). "
                "The checkpoint no longer matches HeadroomCompressorModel's architecture."
            )


def _load_plain_state_dict(model: Any, weights_path: str, model_id: str) -> None:
    """Load a plain, already-merged full state-dict (the pre-v2 / non-PEFT format)."""
    from safetensors.torch import load_file

    state_dict = load_file(weights_path)
    missing, unexpected = model.load_state_dict(state_dict, strict=False)
    if missing or unexpected:
        raise RuntimeError(
            f"{model_id} model.safetensors: state_dict mismatch against "
            f"HeadroomCompressorModel (missing={list(missing)[:5]}, "
            f"unexpected={list(unexpected)[:5]}). Refusing to run with unloaded weights."
        )

View on GitHub (pinned to 322425c43b)

Solutions

Pin the headroom version that matches the model repo's export, or update the model repo/checkpoint to the current architecture.
Clear the HF cache for that model_id and re-download so checkpoint and code revisions align.
Inspect the listed missing/unexpected keys (LoRA prefix drift is the classic cause) and re-export merged.pt with matching names.

Example fix

# before: package upgraded, cached merged.pt from old arch -> RuntimeError mismatch
# shell fix:
rm -rf "$HF_HOME/hub/models--<org>--<kompress-model>"
python -c "from headroom.transforms.kompress_compressor import load_kompress_model; load_kompress_model('<model_id>')"

Defensive patterns

Strategy: try-catch

Try / catch

try:
    _load_pytorch_weights(model, model_id, allow_download=allow_download)
except RuntimeError as e:
    if "state_dict mismatch" in str(e):
        logger.error("checkpoint/model skew for %s; pin versions", model_id)
        raise ModelVersionSkew(model_id) from e
    raise

Prevention

Keep headroom package version and model-repo revision in lockstep.
Clear the HF cache for the model after upgrading headroom.
Watch for LoRA key-prefix drift when using forked/fine-tuned checkpoints.

When it happens

Trigger: Loading merged.pt exported against an older (or newer) HeadroomCompressorModel definition — e.g. package upgraded but HF cache still holds the pre-upgrade checkpoint, or vice versa.

Common situations: Version skew between the headroom package and the downloaded model repo; a fork fine-tuned from a different base exporting mismatched key names (LoRA prefixes, renamed modules).

Related errors

AI-assisted analysis of headroomlabs-ai/headroom@322425c43b (2026-08-15). Data as JSON: /api/errors/a85d5245c572fd5e. Report an issue: GitHub.