xai-org/grok-1 · critical · ValueError

Parameters in the code are not matching checkpoint parameter

Error message

Parameters in the code are not matching checkpoint parameters.
Params missing in checkpoint: {}
Params missing in code: {}

What it means

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).

Source

Thrown at checkpoint.py:207

    ckpt_path = os.path.join(checkpoint_path, "ckpt-0")

    rank_logger.info("Loading checkpoint at {}".format(ckpt_path))
    ckpt_shapes = state_shapes
    ckpt_shapes_with_path, structure = jax.tree_util.tree_flatten_with_path(ckpt_shapes)

    ckpt_shapes_flat = [elem[1] for elem in ckpt_shapes_with_path]
    loaded_tensors = load_tensors(ckpt_shapes_flat, ckpt_path, between_hosts_config)

    state = jax.tree_util.tree_unflatten(structure, loaded_tensors)

    # Sanity check to give a better error message.
    ckpt_keys = set(state.params.keys())
    code_keys = set(state_sharding.params.keys())

    if ckpt_keys != code_keys and init_state is None:
        missing_in_ckpt = code_keys - ckpt_keys
        missing_locally = ckpt_keys - code_keys
        raise ValueError(
            "Parameters in the code are not matching checkpoint parameters.\n"
            "Params missing in checkpoint: {}\nParams missing in code: {}".format(
                missing_in_ckpt, missing_locally
            )
        )
    state_sharding = jax.tree_util.tree_map(
        lambda x: jax.sharding.PartitionSpec() if x is None else x,
        state_sharding,
        is_leaf=lambda x: x is None,
    )
    state = multihost_utils.host_local_array_to_global_array(state, mesh, state_sharding)
    if params_only:
        state = state.params
    return state

View on GitHub (pinned to 7050ed204b)

Solutions

1. 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.).
2. 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.
3. 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 .).
4. 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.
5. 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.

Example fix

// before (config edited for a smaller model)
@dataclass
class ModelConfig:
    vocabulary: int = 32000
    num_layers: int = 8
    ...
// after (released Grok-1 values that match the checkpoint)
@dataclass
class ModelConfig:
    vocabulary: int = 131072
    num_layers: int = 64
    num_attention_heads: int = 48
    num_experts: int = 8
    ...

import numpy as np, glob, pickle
from model import ModelConfig

# 1) checkpoint side: inspect what keys the shards actually contain
#    (run once, before load_checkpoint)
for f in sorted(glob.glob(ckpt_path + '/*'))[:1]:
    with open(f, 'rb') as fh:
        head = pickle.load(fh)
    print(type(head), getattr(head, 'name', None))

# 2) code side: dry-run the model and collect expected top-level params
import haiku as hk, jax
from model import model_config

def _fwd(tokens, *, rng, pad):  # minimal signature from run.py
    from model import Grok1  # whatever entry model.py exposes
    ...  # build exactly as in run.py's forward fn

expected = {'model', 'model_layer_norm'}  # sanity anchor set
tree = jax.eval_shape(lambda: hk.transform(_fwd).init(rng, tokens))
# compare against the names printed in step 1 before calling load_checkpoint

def params_compatible(ckpt_key_set: set[str], code_tree_params) -> bool:
    """True when the model's top-level params match the checkpoint key set."""
    code_keys = set(code_tree_params.keys())
    return ckpt_key_set == code_keys

try:
    state = load_checkpoint(...)
except ValueError as e:
    if 'Params missing in checkpoint' in str(e):
        # structural mismatch: never retry as-is; fix config or checkpoint source
        logging.error('config/checkpoint mismatch: %s', e)
        raise
    raise

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Input must not be scalar.
• Mask dimensionality {mask.ndim} must match logits dimensiona

AI-assisted analysis of xai-org/grok-1@7050ed204b (2026-08-15). Data as JSON: /api/errors/a82ccba68d84c740. Report an issue: GitHub.