Lightning-AI/pytorch-lightning · error · NotImplementedError

Loading a single optimizer object from a checkpoint is not s

Error message

Loading a single optimizer object from a checkpoint is not supported yet with the FSDP strategy.

What it means

The FSDP strategy can load a single bare module (raw module state) and can load optimizers when they are part of a state dict, but loading one standalone Optimizer object is not implemented because sharded optimizer state needs its paired module context.

Source

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

    ) -> dict[str, Any]:
        """Load the contents from a checkpoint and restore the state of the given objects."""
        if not state:
            raise ValueError(
                f"Got FSDPStrategy.load_checkpoint(..., state={state!r}) but a state with at least "
                f" a model instance to reload is required. Pass it in like so:"
                " FSDPStrategy.load_checkpoint(..., state={'model': model, ...})"
            )
        # broadcast the path from rank 0 to ensure all the states are loaded from a common path
        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"

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Pass model and optimizer together: load_checkpoint(path, {'model': model, 'optimizer': optimizer})
2. Load the module first with load_checkpoint(path, model), then restore optimizer state via torch.distributed.checkpoint.optimizer.load_sharded_optimizer_state_dict manually
3. Watch the repo for the upstream feature that implements single-optimizer loading

Example fix

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

from torch.optim import Optimizer

def is_single_optimizer(state) -> bool:
    return isinstance(state, Optimizer)

try:
    strategy.load_checkpoint(path, state)
except NotImplementedError:
    # fall back to combined model+optimizer state
    strategy.load_checkpoint(path, state={'model': model, 'optimizer': state})

When it happens

Trigger: Calling strategy.load_checkpoint(path, optimizer) where optimizer is a torch Optimizer instance rather than a Module or a dict.

Common situations: Resuming only optimizer state (e.g. for frozen-feature fine-tuning); porting resume logic from DDPStrategy where a lone optimizer was acceptable.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• The `{type(self).__name__}` does not use the `CheckpointIO`
• The `{type(self).__name__}` does not support setting a `Chec
• 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

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