{"record":{"id":"b1f51027ba86d644","repo":"Lightning-AI/pytorch-lightning","slug":"found-multiple-xlafsdp-modules-in-the-given-state","errorCode":null,"errorMessage":"Found multiple XLAFSDP modules in the given state. Saving checkpoints with FSDP is currently limited to a single model per checkpoint. To save multiple models, call the save method for each model separately with a different path.","messagePattern":"Found multiple XLAFSDP modules in the given state\\. Saving checkpoints with FSDP is currently limited to a single model per checkpoint\\. To save multiple models, call the save 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":438,"sourceCode":"        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:\n            # each host runs this in parallel, but the ranks in the host run it sequentially\n            for rank in range(len(parallel_devices)):\n                if rank == self.local_rank:\n                    self._save_checkpoint_shard(path, state, storage_options, filter)\n                self.barrier(f\"wait-for-{rank}-save\")\n        else:","sourceCodeStart":420,"sourceCodeEnd":456,"githubUrl":"https://github.com/Lightning-AI/pytorch-lightning/blob/9fed5c27d2a62ff0efd6c3573599921d6ff67c14/src/lightning/fabric/strategies/xla_fsdp.py#L420-L456","documentation":"XLAFSDPStrategy.save_checkpoint refuses to save a checkpoint whose state dict contains more than one XlaFullyShardedDataParallel module. torch_xla's FSDP checkpoint format and consolidation tooling assume exactly one sharded model per checkpoint shard set, so multiple wrapped models in one call are unsupported. Each model must be saved in its own separate checkpoint.","triggerScenarios":"Calling fabric.save_checkpoint(path, state={'model1': m1, 'model2': m2}) where both m1 and m2 are XLAFSDP-wrapped (e.g. GAN generator+discriminator or ensembles on TPU with XLA FSDP).","commonSituations":"Training GANs or multi-stage models where several models were each passed through fabric.setup(); migrating multi-model training code from single-device to XLA FSDP strategy.","solutions":["Split into one save_checkpoint call per model, each with its own path: save_checkpoint(path_a, {'model': m1}); save_checkpoint(path_b, {'model': m2})","Keep only one XLAFSDP-wrapped model per checkpoint and store other models unwrapped via a separate mechanism if they don't need sharding"],"exampleFix":"// before\nfabric.save_checkpoint(path, state={'gen': gen, 'disc': disc})\n\n// after\nfabric.save_checkpoint(path / 'gen.ckpt', state={'model': gen})\nfabric.save_checkpoint(path / 'disc.ckpt', state={'model': disc})","handlingStrategy":"validation","validationCode":"from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\nwrapped = {k: v for k, v in state.items() if isinstance(v, XLAFSDP)}\nassert len(wrapped) <= 1, f'one model per XLAFSDP checkpoint, got {list(wrapped)}'","typeGuard":"from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP\n\ndef xlafsdp_count(state: dict) -> int:\n    return sum(isinstance(v, XLAFSDP) for v in state.values())","tryCatchPattern":null,"preventionTips":["Adopt one-checkpoint-per-model conventions in multi-model training from the start","Name checkpoint paths after the model they contain"],"tags":["xla","fsdp","checkpoint","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"}