Lightning-AI/pytorch-lightning · error · KeyError

The model contains a key '{full_param_name}' that does not e

Error message

The model contains a key '{full_param_name}' that does not exist in the loaded checkpoint. To disable strict loading, set `strict=False`.

What it means

With strict=True (default), _load_raw_module_state verifies every parameter/buffer name in the model exists in the loaded checkpoint state dict. A missing key — typically caused by saving from a model with different layer names (prefix mismatches, architecture changes, or a submodule added after saving) — raises this KeyError.

Source

Thrown at src/lightning/fabric/strategies/model_parallel.py:581

    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

    if _has_dtensor_modules(module):
        from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict

        state_dict_options = StateDictOptions(
            broadcast_from_rank0=True,
            full_state_dict=True,
            # must be set False to allow loading each param separately below
            strict=False,
        )

        for submodule_name, submodule in module.named_modules():
            for param_name, _ in _named_parameters_and_buffers_to_load(submodule):
                full_param_name = f"{submodule_name}{'.' if submodule_name else ''}{param_name}"
                if full_param_name not in state_dict:
                    if not strict:
                        continue
                    raise KeyError(
                        f"The model contains a key '{full_param_name}' that does not exist in the loaded checkpoint."
                        " To disable strict loading, set `strict=False`."
                    )
                local_state_dict = {param_name: state_dict[full_param_name]}
                set_model_state_dict(submodule, local_state_dict, options=state_dict_options)

    elif isinstance(module, FSDP):
        with _get_full_state_dict_context(module, world_size=world_size, rank0_only=False):
            module.load_state_dict(state_dict, strict=strict)
    else:
        module.load_state_dict(state_dict, strict=strict)


def _named_parameters_and_buffers_to_load(module: Module) -> Generator:
    """Returns parameters and buffers, with non-persistent buffers excluded."""
    for param_name, param in itertools.chain(
        module.named_buffers(recurse=False),
        module.named_parameters(recurse=False),

View on GitHub (pinned to 9fed5c27d2)

Solutions

Align key names: wrap/unwrap the module identically at save and load time, or remap keys with a prefix-strip before loading
If missing keys are expected (new layers), pass strict=False so they keep their initialization
Inspect mismatch: sorted(set(k for _, _ in ...) ) vs state_dict keys to find the offending names
Re-save the checkpoint from the current architecture

Example fix

# before
strategy.load_checkpoint(path, state={'model': model})  # raises KeyError: 'model.head.weight'
# after
strategy.load_checkpoint(path, state={'model': model}, strict=False)

Defensive patterns

Strategy: fallback

Validate before calling

ckpt = torch.load(path, map_location='cpu')
missing = [n for n, _ in model.named_parameters() if n not in ckpt]
if missing:
    print('missing keys:', missing[:10])

Try / catch

try:
    strategy.load_checkpoint(path, state={'model': model})
except KeyError as e:
    if 'does not exist in the loaded checkpoint' in str(e):
        strategy.load_checkpoint(path, state={'model': model}, strict=False)
    else:
        raise

Prevention

Use identical wrapping/naming at save and load time
Pass strict=False when fine-tuning across architecture variants

When it happens

Trigger: Loading a checkpoint saved from a differently-named/architected model (e.g. model wrapped or renamed between save and load), loading a sharded checkpoint missing a tensor, or a key prefix like 'model.' vs '' differing between save and load.

Common situations: Fine-tuning from a checkpoint of a slightly different architecture; module was wrapped in a container before saving but not when loading (or vice versa); resuming after adding a new layer (new head, extra embedding).

Related errors

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/0d180b09a75fe927. Report an issue: GitHub.