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: `load_checkpoint(..., state={'model': model, ...})`. Make sure you set up the model (and optimizers if any) through the strategy before loading the checkpoint.

What it means

Mirror of the save-side check: when loading a sharded XLAFSDP checkpoint, the state dict passed to load_checkpoint must contain at least one XlaFullyShardedDataParallel-wrapped module so sharded tensors can be loaded back into the correct wrapped structures. Without a wrapped model there is no target for the shards.

Source

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

        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"
                    " 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."
                )
            if len(modules) > 1:
                raise ValueError(
                    "Found multiple XLAFSDP modules 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 = list(modules.items())[0]
            sharded_ckpt = torch.load(file)

            module.load_state_dict(sharded_ckpt["model"], strict=strict)
            for opt_key, opt in optimizers.items():
                opt.load_state_dict(sharded_ckpt[opt_key])

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Set up the model through fabric/strategy first (model = fabric.setup(model)) so it is XLAFSDP-wrapped, then call load_checkpoint with that wrapped module
2. Ensure the state dict includes the model under some key: state={'model': model, ...}

Example fix

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

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

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP
assert any(isinstance(v, XLAFSDP) for v in state.values()), 'load with an XLAFSDP-wrapped model from fabric.setup()'

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP

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

When it happens

Trigger: Calling fabric.load_checkpoint(path, state={'model': raw_model}) where the model was not passed through fabric.setup/strategy.setup (thus not XLAFSDP-wrapped), or passing a state with only non-module entries.

Common situations: Restore-before-setup ordering bugs: loading the checkpoint before wrapping the model; refactors that moved fabric.setup after load_checkpoint; mixing wrapped and unwrapped model references.

Related errors

• Found a XLAFSDP model in the provided checkpoint state. Plea
• Could not find a XLAFSDP model in the provided checkpoint st
• Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})`
• Loading a single module or optimizer object from a checkpoin
• The path {str(file)!r} does not point to valid sharded check

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