Lightning-AI/pytorch-lightning · error · ValueError

Found multiple FSDP models in the given state. Loading check

Error message

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.

What it means

Symmetric to the save-side restriction: FSDP checkpoint loading supports exactly one FSDP model per load call because the sharded checkpoint layout identifies one module. A state containing multiple FSDP modules is ambiguous and rejected.

Source

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

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

                if optimizers:
                    # TODO: replace with newer APIs
                    # https://github.com/pytorch/pytorch/issues/119800#issuecomment-1942156271
                    reader = _get_distributed_checkpoint_reader(path)

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Load each model separately: load_checkpoint(path1, {'model': model_a}) then load_checkpoint(path2, {'model': model_b})
2. Restructure training so only one model is FSDP-wrapped and the others use a different strategy
3. If checkpoints must share a path, use distinct subdirectories per model

Example fix

# before
strategy.load_checkpoint('ckpt', state={'unet': unet, 'vae': vae})
# after
strategy.load_checkpoint('ckpt/unet', state={'unet': unet})
strategy.load_checkpoint('ckpt/vae', state={'vae': vae})

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
mods = [k for k, v in state.items() if hasattr(v, 'modules') and any(isinstance(m, FSDP) for m in v.modules())]
assert len(mods) <= 1, f'multiple FSDP models: {mods}'

def single_fsdp_keys(state: dict) -> list[str]:
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
    return [k for k, v in state.items() if hasattr(v, 'modules') and any(isinstance(m, FSDP) for m in v.modules())]

When it happens

Trigger: Calling load_checkpoint(path, {'model': model_a, 'vae': model_b}) where both entries contain FSDP-wrapped modules.

Common situations: Restoring a multi-component generative model (UNet + VAE + text encoder) from what used to be a single combined checkpoint; resuming knowledge-distillation setups with student and teacher both sharded.

Related errors

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

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