{"record":{"id":"661e6d37068fd986","repo":"Lightning-AI/pytorch-lightning","slug":"found-multiple-xlafsdp-modules-in-the-given-state-661e6d","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"src/lightning/fabric/strategies/xla_fsdp.py","lineNumber":561,"sourceCode":"        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\"\n                    \" part of the state like so: `load_checkpoint(..., state={'model': model, ...})`. Make sure\"\n                    \" you set up the model (and optimizers if any) through the strategy before loading the checkpoint.\"\n                )\n            if len(modules) > 1:\n                raise ValueError(\n                    \"Found multiple XLAFSDP modules in the given state. Loading checkpoints with FSDP is\"\n                    \" currently limited to a single model per checkpoint. To load multiple models, call the\"\n                    \" load method for each model separately with a different path.\"\n                )\n\n            _, module = list(modules.items())[0]\n            sharded_ckpt = torch.load(file)\n\n            module.load_state_dict(sharded_ckpt[\"model\"], strict=strict)\n            for opt_key, opt in optimizers.items():\n                opt.load_state_dict(sharded_ckpt[opt_key])\n\n            # Load anything leftover from sharded_ckpt\n            loaded_metadata_keys = sharded_ckpt.keys() - modules.keys() - optimizers.keys()\n            requested_metadata_keys = state.keys() - modules.keys() - optimizers.keys()\n            _validate_keys_for_strict_loading(requested_metadata_keys, loaded_metadata_keys, strict=strict)\n            for key in requested_metadata_keys:\n                if key in loaded_metadata_keys:","sourceCodeStart":543,"sourceCodeEnd":579,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/strategies/xla_fsdp.py#L543-L579","documentation":"The load-side counterpart to the multiple-models save restriction: loading a sharded XLA FSDP checkpoint with more than one XLAFSDP module in the state is unsupported because each checkpoint shard set corresponds to exactly one model. Call load_checkpoint separately per model with each model's own checkpoint path.","triggerScenarios":"fabric.load_checkpoint(path, state={'gen': gen, 'disc': disc}) with both modules XLAFSDP-wrapped (GANs, ensembles, multi-stage pipelines on TPU FSDP).","commonSituations":"Multi-model TPU training resuming from checkpoints; adapting single-model example code to a GAN setup.","solutions":["Issue one load_checkpoint call per model, each with its own path and single-model state","Keep per-model checkpoint directories aligned with per-model save calls"],"exampleFix":"# before\nfabric.load_checkpoint(path, state={'gen': gen, 'disc': disc})\n\n# after\nfabric.load_checkpoint(path / 'gen', state={'model': gen})\nfabric.load_checkpoint(path / 'disc', state={'model': disc})","handlingStrategy":"validation","validationCode":"wrapped = [k for k, v in state.items() if isinstance(v, XLAFSDP)]\nassert len(wrapped) <= 1, 'load each XLAFSDP model with its own load_checkpoint call'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Maintain a dict of model->checkpoint-path mapping in multi-model training"],"tags":["xla","fsdp","checkpoint","load","multiple-models","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"}