headroomlabs-ai/headroom · error · RuntimeError

{label}: state_dict mismatch (missing={list(missing)[:5]}, u

Error message

{label}: state_dict mismatch (missing={list(missing)[:5]}, unexpected={list(unexpected)[:5]}). Architecture drifted from the checkpoint.

What it means

During ONNX export, each submodel (encoder, token_head, span_conv) is loaded with `load_state_dict(strict=False)` and then manually checked: any missing or unexpected tensor names mean the model class in the current codebase does not match the architecture the checkpoint was trained with. The message lists up to five names from each side to identify the drift.

Source

Thrown at scripts/export_kompress_v2_onnx.py:84

    from huggingface_hub import hf_hub_download

    from headroom.transforms.kompress_compressor import _get_model_class

    ckpt_path = hf_hub_download(model_id, "merged.pt")
    ckpt = torch.load(ckpt_path, map_location="cpu")
    for key in ("encoder_state_dict", "token_head_state_dict", "span_conv_state_dict"):
        if key not in ckpt:
            raise RuntimeError(
                f"merged.pt missing '{key}'. Found: {sorted(ckpt)}. "
                "This script targets the v2 'merged' checkpoint format."
            )

    core = _get_model_class()(model_name=BASE_MODEL)

    def _strict_load(module, sd, label: str) -> None:
        missing, unexpected = module.load_state_dict(sd, strict=False)
        if missing or unexpected:
            raise RuntimeError(
                f"{label}: state_dict mismatch (missing={list(missing)[:5]}, "
                f"unexpected={list(unexpected)[:5]}). Architecture drifted from the checkpoint."
            )
        logger.info("  %s loaded (%d tensors, exact match)", label, len(sd))

    logger.info("Loading merged.pt (checkpoint_kind=%s)", ckpt.get("checkpoint_kind"))
    _strict_load(core.encoder, ckpt["encoder_state_dict"], "encoder")
    _strict_load(core.token_head, ckpt["token_head_state_dict"], "token_head")
    _strict_load(core.span_conv, ckpt["span_conv_state_dict"], "span_conv")

    core.eval()
    return core


def _export_wrapper(core):
    """Wrap the dual head so forward() returns `final_scores` (== get_scores)."""
    import torch
    import torch.nn as nn

View on GitHub (pinned to 322425c43b)

Solutions

Read the printed missing/unexpected names — a pure prefix rename means the code changed names; a full mismatch means the architecture or base model changed.
Checkout the code revision the checkpoint was trained with, or re-export a merged.pt from the current code.
If the rename is intentional, write a key-mapping shim that rewrites state_dict keys before load_state_dict (then load strict).
Confirm BASE_MODEL matches the model the checkpoint was fine-tuned from.

Example fix

# before
missing, unexpected = module.load_state_dict(sd, strict=False)
if missing or unexpected:
    raise RuntimeError(f"{label}: state_dict mismatch ...")

# after: remap renamed keys, then demand an exact load
sd = {f"layers.{k.removeprefix('blocks.')}" if k.startswith("blocks.") else k: v for k, v in sd.items()}
module.load_state_dict(sd, strict=True)

Defensive patterns

Strategy: validation

Validate before calling

def keys_compatible(module, sd: dict) -> tuple[list[str], list[str]]:
    model_keys = set(module.state_dict().keys())
    sd_keys = set(sd.keys())
    return sorted(model_keys - sd_keys), sorted(sd_keys - model_keys)
# before strict load: assert both lists are empty

Type guard

def has_matching_state(module, sd: dict) -> bool:
    model_keys = set(module.state_dict().keys())
    sd_keys = set(sd.keys())
    return model_keys == sd_keys and all(module.state_dict()[k].shape == v.shape for k, v in sd.items() if k in model_keys)

Prevention

Compare checkpoint keys against module.state_dict().keys() before loading.
Tie checkpoint and code versions together (same commit or model registry tag).
When renaming layers intentionally, write an explicit key-remap and keep it versioned next to the model.

When it happens

Trigger: Renaming layers/modules in the model code after the checkpoint was trained; changing config defaults that alter layer shapes or counts (hidden size, number of layers); loading a v2 checkpoint into a newer architecture revision.

Common situations: Pulling new model code but an old hub checkpoint; a base model name change (BASE_MODEL) altering the transformer architecture; refactors that rename `encoder.*`, `token_head.*`, or `span_conv.*` parameter prefixes.

Related errors

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