Lightning-AI/pytorch-lightning · error · ValueError

Found multiple distributed models in the given state. Loadin

Error message

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.

What it means

ModelParallelStrategy's distributed checkpoint loading supports at most one distributed model per checkpoint. When more than one entry in the `state` dict contains modules with DTensor parameters, this ValueError is raised because the sharded checkpoint format ties shards to a single model's state.

Source

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

    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)
        module.load_state_dict(module_state[module_key], strict=strict)

        # the optimizer states must be loaded separately
        for optim_key, optim in optimizers.items():
            optim_state = {optim_key: get_optimizer_state_dict(module, optim)}
            _distributed_checkpoint_load(optim_state, path)
            set_optimizer_state_dict(module, optim, optim_state_dict=optim_state[optim_key], options=state_dict_options)

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Split into separate calls, one per model, each with its own checkpoint path: load_checkpoint(path1, state={'model1': m1}); load_checkpoint(path2, state={'model2': m2})
2. Only include one distributed model in the state dict and load other (non-distributed) objects separately

Example fix

# before
strategy.load_checkpoint(path, state={'model1': m1, 'model2': m2})
# after
strategy.load_checkpoint(path1, state={'model1': m1})
strategy.load_checkpoint(path2, state={'model2': m2})

distributed = [k for k, v in state.items() if is_distributed_module(v)]
assert len(distributed) <= 1, f'multiple distributed models: {distributed}; load one per checkpoint'

When it happens

Trigger: load_checkpoint(path, state={'model1': m1, 'model2': m2}) where both m1 and m2 have DTensor parameters (both set up through the parallel strategy).

Common situations: Ensembling / multi-model pipelines under ModelParallel or TensorParallel plugins; refactoring code that previously loaded several models in one call under a different strategy.

Related errors

• Could not find a distributed model in the provided checkpoin
• 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/bda0d1491bb067a5. Report an issue: GitHub.