Lightning-AI/pytorch-lightning · error · ValueError

Failed to load checkpoint directly into the model. The given

Error message

Failed to load checkpoint directly into the model. The given path must be a single file containing the full state dict: {path}

What it means

_load_raw_module_state_from_path loads a checkpoint directly into a module and requires the path to be a single file containing the full state dict (a 'full checkpoint'), not a directory of shards. _is_full_checkpoint returned False, so the load is rejected before torch.load.

Source

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

            f" ({world_size})."
        )
    return init_device_mesh(
        device_type=device.type,
        mesh_shape=(data_parallel_size, tensor_parallel_size),
        mesh_dim_names=("data_parallel", "tensor_parallel"),
    )


def _has_dtensor_modules(module: object) -> TypeGuard[Module]:
    from torch.distributed._tensor import DTensor

    return isinstance(module, Module) and any(isinstance(t, DTensor) for t in module.parameters())


def _load_raw_module_state_from_path(path: _PATH, module: Module, world_size: int, strict: bool = True) -> None:
    """Loads the state dict from a file path into the FSDP module."""
    if not _is_full_checkpoint(path):
        raise ValueError(
            "Failed to load checkpoint directly into the model. The given path must be a single file containing the"
            f" full state dict: {path}"
        )
    if _is_local_file_protocol(str(path)):
        # Use `mmap` to avoid storing a copy of the full checkpoint per rank
        state_dict = torch.load(path, mmap=True, map_location="cpu")
    else:
        state_dict = _load(path, map_location="cpu")
    _load_raw_module_state(state_dict=state_dict, module=module, world_size=world_size, strict=strict)


def _load_raw_module_state(
    state_dict: dict[str, Any], module: Module, world_size: int = 1, strict: bool = True
) -> None:
    """Loads the state dict into the module by gathering all weights first and then and writing back to each shard."""
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

    if _has_dtensor_modules(module):

View on GitHub (pinned to 9fed5c27d2)

Solutions

Consolidate the sharded checkpoint into a single file first (e.g. torch.distributed.checkpoint.format_utils.dcp_to_torch_save) and point at that file
Or load through the distributed path: pass state={'model': model} with the module set up via the strategy so _load_checkpoint handles the shard directory
Verify the path is a regular file containing a state dict

Example fix

# before
strategy.load_checkpoint('ckpt_dir', state={'other': metadata})
# after (consolidate shards into one file)
# python -m torch.distributed.checkpoint.format_utils --dcp_to_torch_save ckpt_dir consolidated.pt
strategy.load_checkpoint('consolidated.pt', state={'other': metadata})

Defensive patterns

Strategy: validation

Validate before calling

p = Path(path)
assert p.is_file(), f'{path} must be a single full-checkpoint file, not a shard directory'

Prevention

Keep track of whether a checkpoint is sharded or consolidated
Consolidate shard dirs with dcp_to_torch_save before raw loading

When it happens

Trigger: Calling the raw load path (e.g. load_checkpoint with no distributed model in state, which delegates to _load_raw_module_state_from_path) with a path that is a directory of distributed shards or an otherwise non-file/non-checkpoint location.

Common situations: Mixing checkpoint formats: a sharded checkpoint saved by ModelParallelStrategy later loaded via a path/code path expecting one consolidated file; remote URLs where the file-existence/full-checkpoint probe fails; resuming from a .distcp directory.

Related errors

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