Lightning-AI/pytorch-lightning · error · NotImplementedError

Loading a single module or optimizer object from a checkpoin

Error message

Loading a single module or optimizer object from a checkpoint is not supported yet with the XLAFSDP strategy.

What it means

XLAFSDPStrategy.load_checkpoint does not accept a bare torch.nn.Module or torch Optimizer as the state argument (unlike some other Lightning strategies). FSDP loading needs key names to map checkpoint entries to objects, so only a dict-form state is supported in this strategy. A NotImplementedError signals a known API asymmetry rather than misuse.

Source

Thrown at src/lightning/fabric/strategies/xla_fsdp.py:538

    ) -> dict[str, Any]:
        """Given a folder, load the contents from a checkpoint and restore the state of the given objects.

        The strategy currently only supports saving and loading sharded checkpoints which are stored in form of a
        directory of multiple files rather than a single file.

        """
        if not state:
            raise ValueError(
                f"Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})` but a state with at least "
                " 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 = Path(self.broadcast(path))

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

        from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP

        modules = {key: module for key, module in state.items() if isinstance(module, XLAFSDP)}
        optimizers = {key: optim for key, optim in state.items() if isinstance(optim, Optimizer)}
        if self._state_dict_type == "sharded":
            file = path / f"checkpoint_rank-{self.global_rank:08d}-of-{self.world_size:08d}.pth"
            if not file.is_file():
                raise ValueError(
                    f"The path {str(file)!r} does not point to valid sharded checkpoints. Make sure the path points to"
                    " a directory with XLAFSDP checkpoint shards."
                )
            if len(modules) == 0:
                raise ValueError(
                    "Could not find a XLAFSDP model in the provided checkpoint state. Please provide the model as"

View on GitHub (pinned to 9fed5c27d2)

Solutions

Wrap the object in a dict: fabric.load_checkpoint(path, state={'model': model})
For optimizers, include both the model and optimizer in the state dict since FSDP restoring often requires both

Example fix

# before
fabric.load_checkpoint(path, state=model)

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

Defensive patterns

Strategy: type-guard

Validate before calling

if not isinstance(state, dict):
    state = {'model': state}  # normalize single objects to the dict form

Type guard

from collections.abc import Mapping

def is_dict_state(state) -> bool:
    return isinstance(state, Mapping)

Prevention

Use the dict form state={'model': model, ...} uniformly across strategies
Wrap single-object loads in a small helper that normalizes to dict form

When it happens

Trigger: Calling fabric.load_checkpoint(path, state=model) or fabric.load_checkpoint(path, optimizer) directly with a single object instead of a mapping.

Common situations: Porting code between strategies (e.g. DDPStrategy/Fabric with single-object state) to XLAFSDP; following tutorials that use the shorthand load form.

Related errors

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