headroomlabs-ai/headroom · error · RuntimeError

merged.pt missing '{key}'. Found: {sorted(ckpt)}. This scrip

Error message

merged.pt missing '{key}'. Found: {sorted(ckpt)}. This script targets the v2 'merged' checkpoint format.

What it means

The ONNX export script downloads `merged.pt` from the Hugging Face hub and requires the keys `encoder_state_dict`, `token_head_state_dict`, and `span_conv_state_dict` at the top level. Missing keys mean the checkpoint is not the v2 'merged' format this script targets (e.g., an older sharded checkpoint or a raw state_dict saved by a different training run).

Source

Thrown at scripts/export_kompress_v2_onnx.py:74

    adapters), which does not map onto ``HeadroomCompressorModel``. The
    canonical artifact is ``merged.pt`` — a structured checkpoint with already
    LoRA-merged sub-state-dicts:

        {"encoder_state_dict", "token_head_state_dict",
         "span_conv_state_dict", "config", "checkpoint_kind"}

    Each loads cleanly (0 missing / 0 unexpected) into the encoder + heads.
    """
    import torch
    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")

View on GitHub (pinned to 322425c43b)

Solutions

1. Load the checkpoint locally and print `sorted(ckpt.keys())` to see which format it actually is.
2. Point the script at the repo/revision containing the v2 merged checkpoint (use a pinned `revision=` if the hub file changed).
3. If you own the checkpoint pipeline, re-export a merged.pt with the three expected state_dict keys.
4. Do not try to shim old formats here — the error message is explicit that only v2 is supported.

import torch

REQUIRED_KEYS = ("encoder_state_dict", "token_head_state_dict", "span_conv_state_dict")

def is_v2_merged_checkpoint(path: str) -> bool:
    ckpt = torch.load(path, map_location="cpu")
    return isinstance(ckpt, dict) and all(k in ckpt for k in REQUIRED_KEYS)

def assert_v2_merged(ckpt: object) -> dict:
    keys = ("encoder_state_dict", "token_head_state_dict", "span_conv_state_dict")
    if not isinstance(ckpt, dict) or not all(k in ckpt for k in keys):
        raise TypeError(f"not a v2 merged checkpoint; keys={sorted(ckpt) if isinstance(ckpt, dict) else type(ckpt)}")
    return ckpt

When it happens

Trigger: hf_hub_download fetching a `merged.pt` from a model repo whose checkpoint was replaced with a newer/older format; pointing `model_id` at the wrong repo; the hub file being a full pickled model object instead of the dict of state_dicts.

Common situations: Checkpoint format migration on the model hub (v1 → v2) after the export script was written; passing a staging repo that has not been rebuilt with the v2 export path.

Related errors

• No loadable ONNX artifact in {model_id}; tried {_onnx_filena
• merged.pt for {model_id} is missing {missing_sections}; foun
• Kompress requires onnxruntime or torch. Install with: pip in
• {label}: state_dict mismatch (missing={list(missing)[:5]}, u
• HuggingFace datasets required for this loader. Install with:

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