Lightning-AI/pytorch-lightning · error · ValueError

Could not find a XLAFSDP model in the provided checkpoint st

Error message

Could not find a XLAFSDP model in the provided checkpoint state. Please provide the model as part of the state like so: `save_checkpoint(..., state={'model': model, ...})`. Make sure you set up the model (and optimizers if any) through the strategy before saving the checkpoint.

What it means

Raised by XLAFSDPStrategy.save_checkpoint when no XlaFullyShardedDataParallel-wrapped module exists in the state dict passed to save_checkpoint. The XLA FSDP save path relies on torch_xla's sharded state-dict APIs, which operate on XLAFSDP-wrapped modules, so an unwrapped or missing model cannot be serialized. Lightning therefore requires the model to be both wrapped via the strategy (fabric.setup(model)) and included under a key in the state mapping.

Source

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

        storage_options: Optional[Any] = None,
        filter: Optional[dict[str, Callable[[str, Any], bool]]] = None,
    ) -> None:
        """Save model, optimizer, and other state in the provided checkpoint directory.

        If the user specifies sharded checkpointing, the directory will contain one file per process, with model- and
        optimizer shards stored per file. If the user specifies full checkpointing, the directory will contain a
        consolidated checkpoint combining all of the sharded checkpoints.

        """
        # broadcast the path from rank 0 to ensure all the states are saved in a common path
        path = Path(self.broadcast(path))
        if path.is_dir() and any(path.iterdir()):
            raise FileExistsError(f"The checkpoint directory already exists and is not empty: {path}")
        from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP

        modules = [module for module in state.values() if isinstance(module, XLAFSDP)]
        if len(modules) == 0:
            raise ValueError(
                "Could not find a XLAFSDP model in the provided checkpoint state. Please provide the model as"
                " part of the state like so: `save_checkpoint(..., state={'model': model, ...})`. Make sure"
                " you set up the model (and optimizers if any) through the strategy before saving the checkpoint."
            )
        if len(modules) > 1:
            raise ValueError(
                "Found multiple XLAFSDP modules in the given state. Saving checkpoints with FSDP is"
                " currently limited to a single model per checkpoint. To save multiple models, call the"
                " save method for each model separately with a different path."
            )
        import torch_xla.core.xla_model as xm

        # ensure model parameters are updated
        xm.mark_step()

        parallel_devices = self.parallel_devices
        assert parallel_devices is not None
        if self._sequential_save:

View on GitHub (pinned to 9fed5c27d2)

Solutions

Ensure the model is set up through the strategy/fabric before saving: model = fabric.setup(model) so it gets wrapped in XLAFSDP
Pass the wrapped module in the state: fabric.save_checkpoint(path, state={'model': model, 'optimizer': optimizer})
If you only have a raw state dict, either wrap the model via strategy.setup() first or use torch.save/torch_xla APIs directly instead of the strategy's save_checkpoint

Example fix

// before
fabric.save_checkpoint(path, state={'model': raw_model})

// after
model = fabric.setup(raw_model)  # wraps in XlaFullyShardedDataParallel
optimizer = fabric.setup_optimizer(optimizer)
fabric.save_checkpoint(path, state={'model': model, 'optimizer': optimizer})

Defensive patterns

Strategy: validation

Validate before calling

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP
modules = [v for v in state.values() if isinstance(v, XLAFSDP)]
assert modules, 'state must contain an XLAFSDP-wrapped model; call fabric.setup(model) first'

Type guard

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP
from torch.nn import Module

def is_xlafsdp_module(obj) -> bool:
    return isinstance(obj, XLAFSDP)

def state_has_wrapped_model(state: dict) -> bool:
    return any(is_xlafsdp_module(v) for v in state.values())

Prevention

Always call model = fabric.setup(model) before building the save state
Keep a single code path that constructs state={'model': model, ...} right before save_checkpoint
Add an assertion on the wrapped type before saving in debug builds

When it happens

Trigger: Calling fabric.save_checkpoint(path, state={...}) (or strategy.save_checkpoint) where no value in the state dict is an instance of torch_xla.distributed.fsdp.XlaFullyShardedDataParallel — e.g. passing the raw nn.Module, only optimizer state, or an empty dict.

Common situations: Developer forgets to run model = fabric.setup(model) (which applies the XLAFSDP wrapper) before saving; saving raw state dicts collected before setup; refactoring code so the model key is dropped; passing state={'model': model.state_dict()} instead of the module itself.

Related errors

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