{"record":{"id":"987c2b290cd6a684","repo":"headroomlabs-ai/headroom","slug":"merged-pt-missing-key-found-sorted-ckpt","errorCode":null,"errorMessage":"merged.pt missing '{key}'. Found: {sorted(ckpt)}. This script targets the v2 'merged' checkpoint format.","messagePattern":"merged\\.pt missing '(.+?)'\\. Found: (.+?)\\. This script targets the v2 'merged' checkpoint format\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"scripts/export_kompress_v2_onnx.py","lineNumber":74,"sourceCode":"    adapters), which does not map onto ``HeadroomCompressorModel``. The\n    canonical artifact is ``merged.pt`` — a structured checkpoint with already\n    LoRA-merged sub-state-dicts:\n\n        {\"encoder_state_dict\", \"token_head_state_dict\",\n         \"span_conv_state_dict\", \"config\", \"checkpoint_kind\"}\n\n    Each loads cleanly (0 missing / 0 unexpected) into the encoder + heads.\n    \"\"\"\n    import torch\n    from huggingface_hub import hf_hub_download\n\n    from headroom.transforms.kompress_compressor import _get_model_class\n\n    ckpt_path = hf_hub_download(model_id, \"merged.pt\")\n    ckpt = torch.load(ckpt_path, map_location=\"cpu\")\n    for key in (\"encoder_state_dict\", \"token_head_state_dict\", \"span_conv_state_dict\"):\n        if key not in ckpt:\n            raise RuntimeError(\n                f\"merged.pt missing '{key}'. Found: {sorted(ckpt)}. \"\n                \"This script targets the v2 'merged' checkpoint format.\"\n            )\n\n    core = _get_model_class()(model_name=BASE_MODEL)\n\n    def _strict_load(module, sd, label: str) -> None:\n        missing, unexpected = module.load_state_dict(sd, strict=False)\n        if missing or unexpected:\n            raise RuntimeError(\n                f\"{label}: state_dict mismatch (missing={list(missing)[:5]}, \"\n                f\"unexpected={list(unexpected)[:5]}). Architecture drifted from the checkpoint.\"\n            )\n        logger.info(\"  %s loaded (%d tensors, exact match)\", label, len(sd))\n\n    logger.info(\"Loading merged.pt (checkpoint_kind=%s)\", ckpt.get(\"checkpoint_kind\"))\n    _strict_load(core.encoder, ckpt[\"encoder_state_dict\"], \"encoder\")\n    _strict_load(core.token_head, ckpt[\"token_head_state_dict\"], \"token_head\")","sourceCodeStart":56,"sourceCodeEnd":92,"githubUrl":"https://github.com/headroomlabs-ai/headroom/blob/322425c43bffde1ed0b64fecf3cf5951565dd82b/scripts/export_kompress_v2_onnx.py#L56-L92","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["Load the checkpoint locally and print `sorted(ckpt.keys())` to see which format it actually is.","Point the script at the repo/revision containing the v2 merged checkpoint (use a pinned `revision=` if the hub file changed).","If you own the checkpoint pipeline, re-export a merged.pt with the three expected state_dict keys.","Do not try to shim old formats here — the error message is explicit that only v2 is supported."],"exampleFix":null,"handlingStrategy":"validation","validationCode":"import torch\n\nREQUIRED_KEYS = (\"encoder_state_dict\", \"token_head_state_dict\", \"span_conv_state_dict\")\n\ndef is_v2_merged_checkpoint(path: str) -> bool:\n    ckpt = torch.load(path, map_location=\"cpu\")\n    return isinstance(ckpt, dict) and all(k in ckpt for k in REQUIRED_KEYS)","typeGuard":"def assert_v2_merged(ckpt: object) -> dict:\n    keys = (\"encoder_state_dict\", \"token_head_state_dict\", \"span_conv_state_dict\")\n    if not isinstance(ckpt, dict) or not all(k in ckpt for k in keys):\n        raise TypeError(f\"not a v2 merged checkpoint; keys={sorted(ckpt) if isinstance(ckpt, dict) else type(ckpt)}\")\n    return ckpt","tryCatchPattern":null,"preventionTips":["Pin `revision=` when downloading checkpoints from the hub so format changes cannot surprise you.","Version-stamp checkpoints (checkpoint_kind) and assert it before export.","Keep the export script and the checkpoint-producing code in the same repo revision."],"tags":["pytorch","checkpoint","huggingface","onnx","export"],"backgroundTag":null,"analyzedSha":"322425c43bffde1ed0b64fecf3cf5951565dd82b","analyzedAt":"2026-08-15T01:03:05.481Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}