{"record":{"id":"62a8c4e3d77cdbef","repo":"Lightning-AI/pytorch-lightning","slug":"loading-a-single-module-or-optimizer-object-from-a","errorCode":null,"errorMessage":"Loading a single module or optimizer object from a checkpoint is not supported yet with the XLAFSDP strategy.","messagePattern":"Loading a single module or optimizer object from a checkpoint is not supported yet with the XLAFSDP strategy\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"src/lightning/fabric/strategies/xla_fsdp.py","lineNumber":538,"sourceCode":"    ) -> dict[str, Any]:\n        \"\"\"Given a folder, load the contents from a checkpoint and restore the state of the given objects.\n\n        The strategy currently only supports saving and loading sharded checkpoints which are stored in form of a\n        directory of multiple files rather than a single file.\n\n        \"\"\"\n        if not state:\n            raise ValueError(\n                f\"Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})` but a state with at least \"\n                \" a model instance to reload is required. Pass it in like so:\"\n                \" `FSDPStrategy.load_checkpoint(..., state={'model': model, ...})`\"\n            )\n\n        # broadcast the path from rank 0 to ensure all the states are loaded from a common path\n        path = Path(self.broadcast(path))\n\n        if isinstance(state, (Module, Optimizer)):\n            raise NotImplementedError(\n                \"Loading a single module or optimizer object from a checkpoint\"\n                \" is not supported yet with the XLAFSDP strategy.\"\n            )\n\n        from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\n\n        modules = {key: module for key, module in state.items() if isinstance(module, XLAFSDP)}\n        optimizers = {key: optim for key, optim in state.items() if isinstance(optim, Optimizer)}\n        if self._state_dict_type == \"sharded\":\n            file = path / f\"checkpoint_rank-{self.global_rank:08d}-of-{self.world_size:08d}.pth\"\n            if not file.is_file():\n                raise ValueError(\n                    f\"The path {str(file)!r} does not point to valid sharded checkpoints. Make sure the path points to\"\n                    \" a directory with XLAFSDP checkpoint shards.\"\n                )\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\"","sourceCodeStart":520,"sourceCodeEnd":556,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/strategies/xla_fsdp.py#L520-L556","documentation":"XLAFSDPStrategy.load_checkpoint does not accept a bare torch.nn.Module or torch Optimizer as the state argument (unlike some other Lightning strategies). FSDP loading needs key names to map checkpoint entries to objects, so only a dict-form state is supported in this strategy. A NotImplementedError signals a known API asymmetry rather than misuse.","triggerScenarios":"Calling fabric.load_checkpoint(path, state=model) or fabric.load_checkpoint(path, optimizer) directly with a single object instead of a mapping.","commonSituations":"Porting code between strategies (e.g. DDPStrategy/Fabric with single-object state) to XLAFSDP; following tutorials that use the shorthand load form.","solutions":["Wrap the object in a dict: fabric.load_checkpoint(path, state={'model': model})","For optimizers, include both the model and optimizer in the state dict since FSDP restoring often requires both"],"exampleFix":"# before\nfabric.load_checkpoint(path, state=model)\n\n# after\nfabric.load_checkpoint(path, state={'model': model, 'optimizer': optimizer})","handlingStrategy":"type-guard","validationCode":"if not isinstance(state, dict):\n    state = {'model': state}  # normalize single objects to the dict form","typeGuard":"from collections.abc import Mapping\n\ndef is_dict_state(state) -> bool:\n    return isinstance(state, Mapping)","tryCatchPattern":null,"preventionTips":["Use the dict form state={'model': model, ...} uniformly across strategies","Wrap single-object loads in a small helper that normalizes to dict form"],"tags":["xla","fsdp","checkpoint","load","not-implemented","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"}