headroomlabs-ai/headroom · error · RuntimeError

merged.pt for {model_id} is missing {missing_sections}; foun

Error message

merged.pt for {model_id} is missing {missing_sections}; found keys: {sorted(ckpt)}. This checkpoint format is not what the loader expects.

What it means

Raised by _load_merged_state_dict when the downloaded merged.pt checkpoint lacks one or more of the three expected sub-state-dict sections (encoder_state_dict, token_head_state_dict, span_conv_state_dict). The v2 loader expects a dict-of-state-dicts keyed by submodule (see scripts/export_kompress_v2_onnx.py), not a flat tensor dict; the error enumerates what is missing and what keys were actually found so you can tell a wrong-format or stale file at a glance.

Source

Thrown at headroom/transforms/kompress_compressor.py:791

# scripts/export_kompress_v2_onnx.py, which this mirrors).
_MERGED_CHECKPOINT_KEYS = ("encoder_state_dict", "token_head_state_dict", "span_conv_state_dict")


def _load_merged_state_dict(model: Any, ckpt_path: str, model_id: str) -> None:
    """Load a merged v2-style checkpoint (LoRA already folded into the encoder).

    The checkpoint is a dict of per-submodule state-dicts
    (``encoder_state_dict`` / ``token_head_state_dict`` / ``span_conv_state_dict``)
    rather than a single flat state-dict, so each piece is loaded into its
    matching submodule directly instead of via a single ``load_state_dict``
    call on the whole model.
    """
    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."
            )

View on GitHub (pinned to 322425c43b)

Solutions

Clear the cached merged.pt (huggingface-cli delete ... or remove it under HF_HOME) and re-download the current revision.
Compare 'found keys' in the message against the expected three sections to identify the format mismatch.
If the repo genuinely ships no v2 merged checkpoint, use a model_id that does, or rely on the plain model.safetensors fallback path.

Example fix

# before: stale/corrupt merged.pt in cache -> RuntimeError on missing sections
# shell fix:
huggingface-cli delete <model_id> merged.pt
python -c "from headroom.transforms.kompress_compressor import load_kompress_model; load_kompress_model('<model_id>')"

Defensive patterns

Strategy: retry

Try / catch

try:
    _load_merged_state_dict(model, ckpt_path, model_id)
except RuntimeError as e:
    if "merged.pt" in str(e) and "missing" in str(e):
        evict_hf_file(model_id, "merged.pt")  # then retry once with a fresh download
        _load_merged_state_dict(model, re_download(model_id, "merged.pt"), model_id)
    else:
        raise

Prevention

Evict and re-download checkpoints after headroom upgrades instead of trusting stale caches.
Pin model repos to a known-good revision.
Check the 'found keys' list in the message before assuming corruption.

When it happens

Trigger: Loading a Kompress model whose cached merged.pt is from an older export format, was replaced by a flat checkpoint, or is truncated/corrupt — i.e. torch.load succeeds but the top-level keys don't include all three sections.

Common situations: HF repo revision changed the checkpoint layout; a stale local cache from a previous model version; someone pointed model_id at a fork whose merged.pt is the pre-v2 format.

Related errors

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