Lightning-AI/pytorch-lightning · error · ValueError

Found a XLAFSDP model in the provided checkpoint state. Plea

Error message

Found a XLAFSDP model in the provided checkpoint state. Please provide the model without any XLAFSDP wrapper.

What it means

In the state_dict_type='full' load path, the state must contain the model unwrapped — the loader calls model.load_state_dict(full_ckpt['model']) on the raw module, which would fail on an XLAFSDP wrapper (mismatched keys like flat parameters). Therefore the strategy explicitly rejects states that already contain an XlaFullyShardedDataParallel module and asks for the underlying module.

Source

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

            # remove "shard_metadata" that is loaded in
            if "shard_metadata" in metadata:
                metadata.pop("shard_metadata")

            return metadata

        if self._state_dict_type == "full":
            if not path.is_file():
                raise ValueError(
                    f"The path {str(path)!r} does not point to a valid full checkpoint. Make sure the path points to a"
                    " directory with a full XLAFSDP checkpoint."
                )
            if len(optimizers) > 0 or len(state.keys() - modules.keys() - optimizers.keys()) > 0:
                rank_zero_warn(
                    "Loading a full checkpoint will only load the full model."
                    " The optimizer and any additional metadata are not included."
                )
            if len(modules) > 0:
                raise ValueError(
                    "Found a XLAFSDP model in the provided checkpoint state."
                    " Please provide the model without any XLAFSDP wrapper."
                )
            if "model" not in state or not isinstance(model := state["model"], torch.nn.Module):
                raise NotImplementedError("XLAFSDP only supports a single model instance with 'model' as the key.")
            full_ckpt = torch.load(path, weights_only=weights_only)
            model.load_state_dict(full_ckpt.pop("model"), strict=strict)
            return full_ckpt

        raise ValueError(f"Unknown state_dict_type: {self._state_dict_type}")

    @classmethod
    @override
    def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
        strategy_registry.register("xla_fsdp", cls, description=cls.__name__)

    def _parse_fsdp_kwargs(self) -> dict:
        # this needs to be delayed because `self.precision` isn't available at init

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Pass the unwrapped module (e.g. the original nn.Module before fabric.setup, or wrapper.module) in the state
2. Alternatively use state_dict_type='sharded' with the wrapped module if you intend to continue FSDP training

Example fix

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

# after
raw = MyModel()
fabric.load_checkpoint('full.ckpt', state={'model': raw})  # unwrapped
model = fabric.setup(raw)

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP
assert not any(isinstance(v, XLAFSDP) for v in state.values()), 'pass the unwrapped module for full-checkpoint loading'

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP

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

When it happens

Trigger: Calling fabric.load_checkpoint(path, state={'model': wrapped_model}) with state_dict_type='full' where wrapped_model is the fabric.setup() output (XLAFSDP instance).

Common situations: Reusing the same wrapped-model state dict for both sharded resume and full-checkpoint evaluation; loading a consolidated checkpoint into a model that stays wrapped for continued distributed training.

Related errors

• Could not find a XLAFSDP model in the provided checkpoint st
• The path {str(path)!r} does not point to a valid full checkp
• XLAFSDP only supports a single model instance with 'model' a
• Could not find a XLAFSDP model in the provided checkpoint st
• Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})`

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