{"record":{"id":"08d5d2dc365d1e84","repo":"Lightning-AI/pytorch-lightning","slug":"could-not-find-a-xlafsdp-model-in-the-provided-che","errorCode":null,"errorMessage":"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.","messagePattern":"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=(.+?)\\)`\\. Make sure you set up the model \\(and optimizers if any\\) through the strategy before saving the checkpoint\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"src/lightning/fabric/strategies/xla_fsdp.py","lineNumber":432,"sourceCode":"        storage_options: Optional[Any] = None,\n        filter: Optional[dict[str, Callable[[str, Any], bool]]] = None,\n    ) -> None:\n        \"\"\"Save model, optimizer, and other state in the provided checkpoint directory.\n\n        If the user specifies sharded checkpointing, the directory will contain one file per process, with model- and\n        optimizer shards stored per file. If the user specifies full checkpointing, the directory will contain a\n        consolidated checkpoint combining all of the sharded checkpoints.\n\n        \"\"\"\n        # broadcast the path from rank 0 to ensure all the states are saved in a common path\n        path = Path(self.broadcast(path))\n        if path.is_dir() and any(path.iterdir()):\n            raise FileExistsError(f\"The checkpoint directory already exists and is not empty: {path}\")\n        from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\n\n        modules = [module for module in state.values() if isinstance(module, XLAFSDP)]\n        if len(modules) == 0:\n            raise ValueError(\n                \"Could not find a XLAFSDP model in the provided checkpoint state. Please provide the model as\"\n                \" part of the state like so: `save_checkpoint(..., state={'model': model, ...})`. Make sure\"\n                \" you set up the model (and optimizers if any) through the strategy before saving the checkpoint.\"\n            )\n        if len(modules) > 1:\n            raise ValueError(\n                \"Found multiple XLAFSDP modules in the given state. Saving checkpoints with FSDP is\"\n                \" currently limited to a single model per checkpoint. To save multiple models, call the\"\n                \" save method for each model separately with a different path.\"\n            )\n        import torch_xla.core.xla_model as xm\n\n        # ensure model parameters are updated\n        xm.mark_step()\n\n        parallel_devices = self.parallel_devices\n        assert parallel_devices is not None\n        if self._sequential_save:","sourceCodeStart":414,"sourceCodeEnd":450,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/strategies/xla_fsdp.py#L414-L450","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\nfabric.save_checkpoint(path, state={'model': raw_model})\n\n// after\nmodel = fabric.setup(raw_model)  # wraps in XlaFullyShardedDataParallel\noptimizer = fabric.setup_optimizer(optimizer)\nfabric.save_checkpoint(path, state={'model': model, 'optimizer': optimizer})","handlingStrategy":"validation","validationCode":"from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\nmodules = [v for v in state.values() if isinstance(v, XLAFSDP)]\nassert modules, 'state must contain an XLAFSDP-wrapped model; call fabric.setup(model) first'","typeGuard":"from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\nfrom torch.nn import Module\n\ndef is_xlafsdp_module(obj) -> bool:\n    return isinstance(obj, XLAFSDP)\n\ndef state_has_wrapped_model(state: dict) -> bool:\n    return any(is_xlafsdp_module(v) for v in state.values())","tryCatchPattern":null,"preventionTips":["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"],"tags":["xla","fsdp","checkpoint","save","model-wrapping","lightning-fabric"],"backgroundTag":"checkpoint-state-validation-failed","analyzedSha":"9fed5c27d2a62ff0efd6c3573599921d6ff67c14","analyzedAt":"2026-08-28T11:52:41.083Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}