Lightning-AI/pytorch-lightning · error · ValueError

Could not find a distributed model in the provided checkpoin

Error message

Could not find a distributed model in the provided checkpoint state. Please provide the model as part of the state like so: `load_checkpoint(..., state={'model': model, ...})`. Make sure you set up the model (and optimizers if any) through the strategy before loading the checkpoint.

What it means

Raised by ModelParallelStrategy._load_checkpoint when none of the objects passed in the `state` dict of load_checkpoint contain DTensor parameters. The strategy's distributed checkpoint loading only works on modules that were parallelized (their parameters became DTensors) by the strategy itself. If the model wasn't set up through the strategy (or isn't distributed at all), no candidate module is found and loading is aborted.

Source

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

def _load_checkpoint(
    path: _PATH,
    state: dict[str, Union[Module, Optimizer, Any]],
    strict: bool = True,
    optimizer_states_from_list: bool = False,
    weights_only: Optional[bool] = None,
) -> dict[str, Any]:
    from torch.distributed.checkpoint.state_dict import (
        StateDictOptions,
        get_model_state_dict,
        get_optimizer_state_dict,
        set_optimizer_state_dict,
    )

    modules = {key: module for key, module in state.items() if _has_dtensor_modules(module)}
    if len(modules) == 0:
        raise ValueError(
            "Could not find a distributed model in the provided checkpoint state. Please provide the model as"
            " part of the state like so: `load_checkpoint(..., state={'model': model, ...})`. Make sure"
            " you set up the model (and optimizers if any) through the strategy before loading the checkpoint."
        )
    optimizers = {key: optim for key, optim in state.items() if isinstance(optim, Optimizer)}
    if len(modules) > 1:
        raise ValueError(
            "Found multiple distributed models in the given state. Loading distributed checkpoints is"
            " currently limited to a single model per checkpoint. To load multiple models, call the"
            " load method for each model separately with a different path."
        )
    module_key, module = list(modules.items())[0]

    if _is_sharded_checkpoint(path):
        state_dict_options = StateDictOptions(cpu_offload=True)

        module_state = {module_key: get_model_state_dict(module)}
        _distributed_checkpoint_load(module_state, path)

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Ensure the model is added to the state dict under a key, e.g. load_checkpoint(path, state={'model': model})
2. Set up the model (and optimizers) through the strategy first: strategy.setup_module(model) / fabric.setup() so parameters become DTensors
3. Verify the module actually contains DTensors: any(isinstance(p, DTensor) for p in model.parameters())
4. If the checkpoint is a plain full-file state dict, load it with torch.load + model.load_state_dict instead of the distributed path

Example fix

# before
strategy.load_checkpoint(path, state={})
# after
model = strategy.setup_module(model)
strategy.load_checkpoint(path, state={'model': model})

from torch.distributed.tensor import DTensor

def has_dtensor(module) -> bool:
    return any(isinstance(p, DTensor) for p in module.parameters())

assert 'model' in state and has_dtensor(state['model']), 'set up the model through the strategy first'

from torch import nn
from torch.distributed.tensor import DTensor

def is_distributed_module(obj) -> bool:
    return isinstance(obj, nn.Module) and any(isinstance(p, DTensor) for p in obj.parameters())

When it happens

Trigger: Calling fabric.strategy.load_checkpoint(path, state={'model': model, ...}) where `model` has no DTensor parameters — e.g. the model was not wrapped via strategy.setup_module/model_parallel setup, or a plain CPU/SingleDevice strategy checkpoint is being loaded through ModelParallelStrategy, or the state only contains optimizers/raw tensors.

Common situations: User calls load_checkpoint before setup_environment/setup_module; uses ModelParallel strategy with a model that has no parallelized layers; passes a raw state_dict instead of the module; mismatch between the strategy used to save and the one used to load.

Related errors

• Found multiple distributed models in the given state. Loadin
• The path {str(path)!r} does not point to a valid checkpoint.
• The sizes `data_parallel_size={data_parallel_size}` and `ten
• Failed to load checkpoint directly into the model. The given
• The model contains a key '{full_param_name}' that does not e

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