headroomlabs-ai/headroom · critical · RuntimeError

{model_id} model.safetensors: state_dict mismatch against He

Error message

{model_id} model.safetensors: state_dict mismatch against HeadroomCompressorModel (missing={list(missing)[:5]}, unexpected={list(unexpected)[:5]}). Refusing to run with unloaded weights.

What it means

Raised by _load_plain_state_dict (the pre-v2 / non-PEFT fallback path) when loading model.safetensors into the full HeadroomCompressorModel with strict=False still leaves missing or unexpected keys. Unlike the merged.pt path, this is the final weights source, so the loader refuses to run a model with partially unloaded weights rather than emitting garbage compressions — the message says exactly that ('Refusing to run with unloaded weights').

Source

Thrown at headroom/transforms/kompress_compressor.py:817

        ("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."
        )


def _load_pytorch_weights(model: Any, model_id: str, *, allow_download: bool) -> None:
    """Load PyTorch weights into ``model``, preferring the merged v2 checkpoint.

    ``merged.pt`` (when the repo ships one) holds LoRA-merged sub-state-dicts
    keyed by submodule name. In a PEFT-trained repo, ``model.safetensors`` is
    the *unmerged* adapter checkpoint (encoder keys prefixed
    ``encoder.base_model.model...``) and does not map onto this module tree at
    all, so it is only used as a fallback for repos that never shipped a
    merged checkpoint (e.g. the original non-LoRA kompress-base).

    In cache-only mode (``allow_download=False``) a ``merged.pt`` cache miss is
    ambiguous: it could mean the repo has no merged checkpoint (safe to use the

View on GitHub (pinned to 322425c43b)

Solutions

Use a model_id whose model.safetensors matches your headroom version (check the repo's export script version).
If the repo is PEFT-format, ensure a merged.pt exists (or merge and upload one) so the v2 path is used instead of the plain path.
Clear the local cache and re-download in case of a corrupted/stale artifact before assuming architecture mismatch.

Example fix

# before: adapter-format model.safetensors routed to plain loader -> RuntimeError
# after: export a merged v2 checkpoint so the correct loader is used
# python scripts/export_kompress_v2_onnx.py --model <model_id> --emit-merged-pt
# then reload: load_kompress_model('<model_id>')

Defensive patterns

Strategy: try-catch

Try / catch

try:
    load_kompress_model(model_id)
except RuntimeError as e:
    if "Refusing to run with unloaded weights" in str(e):
        disable_kompress_and_use_builtin_compressors()
    else:
        raise

Prevention

Never load a PEFT/adapter-format model.safetensors via the plain path — ship a merged.pt.
Export checkpoints with the current export script (scripts/export_kompress_v2_onnx.py).
Treat this error as hard-stop: partially loaded weights yield silent garbage, so do not catch-and-continue into compression.

When it happens

Trigger: A repo without merged.pt falls back to model.safetensors, and that safetensors file was exported against a different HeadroomCompressorModel architecture (key names/shapes drifted), so strict=False still reports mismatches.

Common situations: Loading a PEFT-format repo's unmerged adapter safetensors through the plain path; architecture refactor in headroom outdating an old safetensors artifact; fork repos with renamed modules.

Related errors

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