{"record":{"id":"1134a446489f655e","repo":"headroomlabs-ai/headroom","slug":"label-state-dict-mismatch-missing-list-missin","errorCode":null,"errorMessage":"{label}: state_dict mismatch (missing={list(missing)[:5]}, unexpected={list(unexpected)[:5]}). Architecture drifted from the checkpoint.","messagePattern":"(.+?): state_dict mismatch \\(missing=(.+?), unexpected=(.+?)\\)\\. Architecture drifted from the checkpoint\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"scripts/export_kompress_v2_onnx.py","lineNumber":84,"sourceCode":"    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\")\n    _strict_load(core.span_conv, ckpt[\"span_conv_state_dict\"], \"span_conv\")\n\n    core.eval()\n    return core\n\n\ndef _export_wrapper(core):\n    \"\"\"Wrap the dual head so forward() returns `final_scores` (== get_scores).\"\"\"\n    import torch\n    import torch.nn as nn","sourceCodeStart":66,"sourceCodeEnd":102,"githubUrl":"https://github.com/headroomlabs-ai/headroom/blob/322425c43bffde1ed0b64fecf3cf5951565dd82b/scripts/export_kompress_v2_onnx.py#L66-L102","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"# before\nmissing, unexpected = module.load_state_dict(sd, strict=False)\nif missing or unexpected:\n    raise RuntimeError(f\"{label}: state_dict mismatch ...\")\n\n# after: remap renamed keys, then demand an exact load\nsd = {f\"layers.{k.removeprefix('blocks.')}\" if k.startswith(\"blocks.\") else k: v for k, v in sd.items()}\nmodule.load_state_dict(sd, strict=True)","handlingStrategy":"validation","validationCode":"def keys_compatible(module, sd: dict) -> tuple[list[str], list[str]]:\n    model_keys = set(module.state_dict().keys())\n    sd_keys = set(sd.keys())\n    return sorted(model_keys - sd_keys), sorted(sd_keys - model_keys)\n# before strict load: assert both lists are empty","typeGuard":"def has_matching_state(module, sd: dict) -> bool:\n    model_keys = set(module.state_dict().keys())\n    sd_keys = set(sd.keys())\n    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)","tryCatchPattern":null,"preventionTips":["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."],"tags":["pytorch","state-dict","model-loading","onnx"],"backgroundTag":null,"analyzedSha":"322425c43bffde1ed0b64fecf3cf5951565dd82b","analyzedAt":"2026-08-15T01:03:05.481Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}