{"record":{"id":"1e5bbbbd9ab94e3c","repo":"Lightning-AI/pytorch-lightning","slug":"found-a-xlafsdp-model-in-the-provided-checkpoint-s","errorCode":null,"errorMessage":"Found a XLAFSDP model in the provided checkpoint state. Please provide the model without any XLAFSDP wrapper.","messagePattern":"Found a XLAFSDP model in the provided checkpoint state\\. Please provide the model without any XLAFSDP wrapper\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"src/lightning/fabric/strategies/xla_fsdp.py","lineNumber":606,"sourceCode":"            # remove \"shard_metadata\" that is loaded in\n            if \"shard_metadata\" in metadata:\n                metadata.pop(\"shard_metadata\")\n\n            return metadata\n\n        if self._state_dict_type == \"full\":\n            if not path.is_file():\n                raise ValueError(\n                    f\"The path {str(path)!r} does not point to a valid full checkpoint. Make sure the path points to a\"\n                    \" directory with a full XLAFSDP checkpoint.\"\n                )\n            if len(optimizers) > 0 or len(state.keys() - modules.keys() - optimizers.keys()) > 0:\n                rank_zero_warn(\n                    \"Loading a full checkpoint will only load the full model.\"\n                    \" The optimizer and any additional metadata are not included.\"\n                )\n            if len(modules) > 0:\n                raise ValueError(\n                    \"Found a XLAFSDP model in the provided checkpoint state.\"\n                    \" Please provide the model without any XLAFSDP wrapper.\"\n                )\n            if \"model\" not in state or not isinstance(model := state[\"model\"], torch.nn.Module):\n                raise NotImplementedError(\"XLAFSDP only supports a single model instance with 'model' as the key.\")\n            full_ckpt = torch.load(path, weights_only=weights_only)\n            model.load_state_dict(full_ckpt.pop(\"model\"), strict=strict)\n            return full_ckpt\n\n        raise ValueError(f\"Unknown state_dict_type: {self._state_dict_type}\")\n\n    @classmethod\n    @override\n    def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:\n        strategy_registry.register(\"xla_fsdp\", cls, description=cls.__name__)\n\n    def _parse_fsdp_kwargs(self) -> dict:\n        # this needs to be delayed because `self.precision` isn't available at init","sourceCodeStart":588,"sourceCodeEnd":624,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/strategies/xla_fsdp.py#L588-L624","documentation":"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.","triggerScenarios":"Calling fabric.load_checkpoint(path, state={'model': wrapped_model}) with state_dict_type='full' where wrapped_model is the fabric.setup() output (XLAFSDP instance).","commonSituations":"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.","solutions":["Pass the unwrapped module (e.g. the original nn.Module before fabric.setup, or wrapper.module) in the state","Alternatively use state_dict_type='sharded' with the wrapped module if you intend to continue FSDP training"],"exampleFix":"# before\nmodel = fabric.setup(MyModel())\nfabric.load_checkpoint('full.ckpt', state={'model': model})  # wrapped\n\n# after\nraw = MyModel()\nfabric.load_checkpoint('full.ckpt', state={'model': raw})  # unwrapped\nmodel = fabric.setup(raw)","handlingStrategy":"type-guard","validationCode":"from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\nassert not any(isinstance(v, XLAFSDP) for v in state.values()), 'pass the unwrapped module for full-checkpoint loading'","typeGuard":"from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\n\ndef state_is_unwrapped(state: dict) -> bool:\n    return not any(isinstance(v, XLAFSDP) for v in state.values())","tryCatchPattern":null,"preventionTips":["Keep a reference to the raw module (or use wrapper.module) for full-checkpoint eval flows"],"tags":["xla","fsdp","checkpoint","load","full-checkpoint","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"}