Lightning-AI/pytorch-lightning · error · ValueError

Found multiple XLAFSDP modules in the given state. Saving ch

Error message

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.

What it means

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.

Source

Thrown at src/lightning/fabric/strategies/xla_fsdp.py:438

        optimizer shards stored per file. If the user specifies full checkpointing, the directory will contain a
        consolidated checkpoint combining all of the sharded checkpoints.

        """
        # broadcast the path from rank 0 to ensure all the states are saved in a common path
        path = Path(self.broadcast(path))
        if path.is_dir() and any(path.iterdir()):
            raise FileExistsError(f"The checkpoint directory already exists and is not empty: {path}")
        from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP

        modules = [module for module in state.values() if isinstance(module, XLAFSDP)]
        if len(modules) == 0:
            raise ValueError(
                "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."
            )
        if len(modules) > 1:
            raise ValueError(
                "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."
            )
        import torch_xla.core.xla_model as xm

        # ensure model parameters are updated
        xm.mark_step()

        parallel_devices = self.parallel_devices
        assert parallel_devices is not None
        if self._sequential_save:
            # each host runs this in parallel, but the ranks in the host run it sequentially
            for rank in range(len(parallel_devices)):
                if rank == self.local_rank:
                    self._save_checkpoint_shard(path, state, storage_options, filter)
                self.barrier(f"wait-for-{rank}-save")
        else:

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. 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})
2. Keep only one XLAFSDP-wrapped model per checkpoint and store other models unwrapped via a separate mechanism if they don't need sharding

Example fix

// before
fabric.save_checkpoint(path, state={'gen': gen, 'disc': disc})

// after
fabric.save_checkpoint(path / 'gen.ckpt', state={'model': gen})
fabric.save_checkpoint(path / 'disc.ckpt', state={'model': disc})

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP
wrapped = {k: v for k, v in state.items() if isinstance(v, XLAFSDP)}
assert len(wrapped) <= 1, f'one model per XLAFSDP checkpoint, got {list(wrapped)}'

from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP

def xlafsdp_count(state: dict) -> int:
    return sum(isinstance(v, XLAFSDP) for v in state.values())

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• Found multiple XLAFSDP modules in the given state. Loading c
• Could not find a XLAFSDP model in the provided checkpoint st
• Multihost setups do not have a shared filesystem, so the che
• Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})`
• Loading a single module or optimizer object from a checkpoin

AI-assisted analysis of Lightning-AI/pytorch-lightning@9fed5c27d2 (2026-08-28). Data as JSON: /api/errors/b1f51027ba86d644. Report an issue: GitHub.