Lightning-AI/pytorch-lightning · error · ValueError

Could not find a FSDP model in the provided checkpoint state

Error message

Could not find a FSDP 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

When loading, the strategy scans the state dict for entries containing FSDP modules; if none are found it cannot determine where to restore weights. This happens when the model was never wrapped in FSDP (not set up through the strategy) or the state dict has no model at all.

Source

Thrown at src/lightning/fabric/strategies/fsdp.py:556

        path = _resolve_path(self.broadcast(path))

        if isinstance(state, Module):
            from lightning.fabric.strategies.model_parallel import _load_raw_module_state_from_path

            _load_raw_module_state_from_path(path, module=state, world_size=self.world_size, strict=strict)
            return {}

        if isinstance(state, Optimizer):
            raise NotImplementedError(
                "Loading a single optimizer object from a checkpoint is not supported yet with the FSDP strategy."
            )

        from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

        modules = {key: module for key, module in state.items() if _has_fsdp_modules(module)}
        if len(modules) == 0:
            raise ValueError(
                "Could not find a FSDP 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 FSDP models in the given state. Loading checkpoints with FSDP 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_ctx = _get_sharded_state_dict_context(module)

            with state_dict_ctx:
                module_state = {module_key: module.state_dict()}

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Create the model and call fabric.setup(model) (which wraps it in FSDP) before load_checkpoint
2. Pass the wrapped model explicitly: state={'model': model} and verify isinstance(model, FullyShardedDataParallel) or that it contains FSDP submodules
3. If you wrapped FSDP manually, make sure the wrapper is on the object you pass in, not a discarded copy

Example fix

# before
model = MyModel()
strategy.load_checkpoint(path, state={'model': model})
# after
model = MyModel()
model, optimizer = fabric.setup(model, optimizer)
strategy.load_checkpoint(path, state={'model': model, 'optimizer': optimizer})

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
assert any(isinstance(m, FSDP) for v in state.values() if hasattr(v, 'modules') for m in v.modules()), 'set up the model through fabric/strategy first'

def state_has_fsdp(state: dict) -> bool:
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
    return any(isinstance(m, FSDP) for v in state.values() if hasattr(v, 'modules') for m in v.modules())

When it happens

Trigger: Calling load_checkpoint(path, {'model': plain_module}) where plain_module was not passed through fabric.setup/strategy.setup; or passing only non-module metadata like {'step': 1000}.

Common situations: Forgetting fabric.setup(model) before load; loading into a fresh, unwrapped module; applying FSDP wrapping manually after calling load; the model being wrapped in a non-FSDP parallel wrapper.

Related errors

• Got FSDPStrategy.load_checkpoint(..., state={state!r}) but a
• Found multiple FSDP models in the given state. Loading check
• Found multiple FSDP models in the given state. Saving checkp
• Loading a single optimizer object from a checkpoint is not s
• The path {str(path)!r} does not point to a valid checkpoint.

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