{"record":{"id":"a82ccba68d84c740","repo":"xai-org/grok-1","slug":"parameters-in-the-code-are-not-matching-checkpoint","errorCode":null,"errorMessage":"Parameters in the code are not matching checkpoint parameters.\nParams missing in checkpoint: {}\nParams missing in code: {}","messagePattern":"Parameters in the code are not matching checkpoint parameters\\.\nParams missing in checkpoint: (.+?)\nParams missing in code: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"critical","filePath":"checkpoint.py","lineNumber":207,"sourceCode":"    ckpt_path = os.path.join(checkpoint_path, \"ckpt-0\")\n\n    rank_logger.info(\"Loading checkpoint at {}\".format(ckpt_path))\n    ckpt_shapes = state_shapes\n    ckpt_shapes_with_path, structure = jax.tree_util.tree_flatten_with_path(ckpt_shapes)\n\n    ckpt_shapes_flat = [elem[1] for elem in ckpt_shapes_with_path]\n    loaded_tensors = load_tensors(ckpt_shapes_flat, ckpt_path, between_hosts_config)\n\n    state = jax.tree_util.tree_unflatten(structure, loaded_tensors)\n\n    # Sanity check to give a better error message.\n    ckpt_keys = set(state.params.keys())\n    code_keys = set(state_sharding.params.keys())\n\n    if ckpt_keys != code_keys and init_state is None:\n        missing_in_ckpt = code_keys - ckpt_keys\n        missing_locally = ckpt_keys - code_keys\n        raise ValueError(\n            \"Parameters in the code are not matching checkpoint parameters.\\n\"\n            \"Params missing in checkpoint: {}\\nParams missing in code: {}\".format(\n                missing_in_ckpt, missing_locally\n            )\n        )\n    state_sharding = jax.tree_util.tree_map(\n        lambda x: jax.sharding.PartitionSpec() if x is None else x,\n        state_sharding,\n        is_leaf=lambda x: x is None,\n    )\n    state = multihost_utils.host_local_array_to_global_array(state, mesh, state_sharding)\n    if params_only:\n        state = state.params\n    return state\n","sourceCodeStart":189,"sourceCodeEnd":222,"githubUrl":"https://github.com/xai-org/grok-1/blob/7050ed204b8206bb8645c7b7bbef7252f79561b0/checkpoint.py#L189-L222","documentation":"Raised by load_checkpoint in checkpoint.py:207 when the parameter tree produced by the model code (state_sharding.params, built from your ModelConfig) does not have exactly the same top-level parameter keys as the tensors loaded from the UL2 checkpoint shards. It is a sanity check fired after load_tensors and tree_unflatten, so the checkpoint was already read from disk/TCP and only the key-set comparison failed. The message lists the two set differences: keys your code expects but the checkpoint lacks (missing_in_ckpt) and keys the checkpoint has but your code does not define (missing_locally).","triggerScenarios":"Calling load or the run.py path that ends in this check with init_state=None, while (a) ModelConfig in model.py (e.g. vocabulary size 131072, 64 layers, 8 heads, MoE width) has been edited from the released Grok-1 values, (b) a different/newer checkpoint format (e.g. single 'model' tensor files from the 2024-03-29 update) is loaded with old code expecting the old shard layout, or (c) the wrong checkpoints directory (partial download, e.g. only some of the 604 shards present) is passed via --checkpoint-path.","commonSituations":"Fine-tuning Grok-1 and changing hyperparameters (num_layers, embedding size, MoE config) so hk.get_parameter names no longer line up; mixing checkpoint versions (the repo was updated to consolidate ~604 small files into ~8 large 'model-*.tensor' files, so old downloads + new code mismatch); a corrupted or truncated gs://grok-1/download because a download was interrupted; running a modified model.py whose top-level params dict ('model', 'model_layer_norm', etc.) differs.","solutions":["Re-read the two sets printed in the error: if BOTH are non-empty you are loading a structurally different model — restore ModelConfig and model.py to the released Grok-1 values (vocabulary=131072, num_layers=64, key_size=128, num_experts=8, etc.).","If keys look like raw shard names vs module names (e.g. 'model' vs 'embedding'), your checkpoint version does not match this code revision — re-download the checkpoint from gs://grok-1 and git pull the matching xai-org/grok-1 commit.","Verify the checkpoint directory is complete: every expected file in ckpt_path exists and is non-empty; the load in load_checkpoint validates sizes for a reason, so finish/restart an interrupted download (gsutil -m cp -r gs://grok-1 .).","If you intentionally changed the architecture, pass init_state (instead of None) so the mismatch branch is skipped and the loader takes the partial-load path that only reads intersecting params.","Pass params_only=True and inspect sorted(state.params.keys()) vs sorted(state_sharding.params.keys()) in a scratch script to see exactly which module names diverge before re-running."],"exampleFix":"// before (config edited for a smaller model)\n@dataclass\nclass ModelConfig:\n    vocabulary: int = 32000\n    num_layers: int = 8\n    ...\n// after (released Grok-1 values that match the checkpoint)\n@dataclass\nclass ModelConfig:\n    vocabulary: int = 131072\n    num_layers: int = 64\n    num_attention_heads: int = 48\n    num_experts: int = 8\n    ...","handlingStrategy":"validation","validationCode":"import numpy as np, glob, pickle\nfrom model import ModelConfig\n\n# 1) checkpoint side: inspect what keys the shards actually contain\n#    (run once, before load_checkpoint)\nfor f in sorted(glob.glob(ckpt_path + '/*'))[:1]:\n    with open(f, 'rb') as fh:\n        head = pickle.load(fh)\n    print(type(head), getattr(head, 'name', None))\n\n# 2) code side: dry-run the model and collect expected top-level params\nimport haiku as hk, jax\nfrom model import model_config\n\ndef _fwd(tokens, *, rng, pad):  # minimal signature from run.py\n    from model import Grok1  # whatever entry model.py exposes\n    ...  # build exactly as in run.py's forward fn\n\nexpected = {'model', 'model_layer_norm'}  # sanity anchor set\ntree = jax.eval_shape(lambda: hk.transform(_fwd).init(rng, tokens))\n# compare against the names printed in step 1 before calling load_checkpoint","typeGuard":"def params_compatible(ckpt_key_set: set[str], code_tree_params) -> bool:\n    \"\"\"True when the model's top-level params match the checkpoint key set.\"\"\"\n    code_keys = set(code_tree_params.keys())\n    return ckpt_key_set == code_keys","tryCatchPattern":"try:\n    state = load_checkpoint(...)\nexcept ValueError as e:\n    if 'Params missing in checkpoint' in str(e):\n        # structural mismatch: never retry as-is; fix config or checkpoint source\n        logging.error('config/checkpoint mismatch: %s', e)\n        raise\n    raise","preventionTips":["Pin the xai-org/grok-1 commit whose model.py matches your downloaded checkpoint generation (small-shard vs consolidated 'model-*.tensor' layout).","Keep ModelConfig byte-identical to the release when you only want inference; any edit to num_layers/vocabulary/MoE settings will trip this check.","After gsutil -m cp of the checkpoint, verify file count and total size against gs://grok-1 metadata before loading.","When intentionally changing the architecture, pass init_state explicitly and treat the loader's partial behavior as the contract, not an error path."],"tags":["grok-1","checkpoint","jax","haiku","config-mismatch","model-loading"],"backgroundTag":null,"analyzedSha":"7050ed204b8206bb8645c7b7bbef7252f79561b0","analyzedAt":"2026-08-15T04:18:25.087Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}