Lightning-AI/pytorch-lightning · error · ValueError

Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})`

Error message

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, ...})`

What it means

XLAFSDPStrategy.load_checkpoint requires a non-empty state dict containing the objects to restore (at minimum the model). Unlike simple torch.load wrappers, FSDP loading must load sharded tensors back into existing (wrapped) module structures, so an empty/falsy state cannot be loaded into anything. The error points the user to the expected state={'model': model, ...} form.

Source

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

            storage_options=storage_options,
        )

    @override
    def load_checkpoint(
        self,
        path: _PATH,
        state: Optional[Union[Module, Optimizer, dict[str, Union[Module, Optimizer, Any]]]] = None,
        strict: bool = True,
        weights_only: Optional[bool] = None,
    ) -> 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)}

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Pass the objects to restore: fabric.load_checkpoint(path, state={'model': model, 'optimizer': optimizer})
2. If you only want to inspect checkpoint contents, use torch.load / torch_xla APIs directly instead of the strategy loader

Example fix

# before
fabric.load_checkpoint(path, state={})

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

assert state, 'load_checkpoint requires a non-empty state dict, e.g. {"model": model}'

def is_loadable_state(state) -> bool:
    return bool(state) and isinstance(state, dict)

When it happens

Trigger: Calling fabric.load_checkpoint(path, state={}) or state=None (or any falsy value) — e.g. loading a checkpoint just to inspect it, or a code path that forgot to populate the state mapping.

Common situations: Copy-pasted checkpoint-inspection code passed to load_checkpoint; refactors that build state conditionally and pass empty dicts; assuming load_checkpoint returns contents without needing target objects.

Related errors

• Loading a single module or optimizer object from a checkpoin
• The path {str(file)!r} does not point to valid sharded check
• Could not find a XLAFSDP model in the provided checkpoint st
• Found multiple XLAFSDP modules in the given state. Loading c
• The path {str(path)!r} does not point to a valid full checkp

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