Lightning-AI/pytorch-lightning · error · ValueError

Found multiple XLAFSDP modules in the given state. Loading c

Error message

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.

What it means

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.

Source

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

        from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as XLAFSDP

        modules = {key: module for key, module in state.items() if isinstance(module, XLAFSDP)}
        optimizers = {key: optim for key, optim in state.items() if isinstance(optim, Optimizer)}
        if self._state_dict_type == "sharded":
            file = path / f"checkpoint_rank-{self.global_rank:08d}-of-{self.world_size:08d}.pth"
            if not file.is_file():
                raise ValueError(
                    f"The path {str(file)!r} does not point to valid sharded checkpoints. Make sure the path points to"
                    " a directory with XLAFSDP checkpoint shards."
                )
            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: `load_checkpoint(..., state={'model': model, ...})`. Make sure"
                    " you set up the model (and optimizers if any) through the strategy before loading the checkpoint."
                )
            if len(modules) > 1:
                raise ValueError(
                    "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."
                )

            _, module = list(modules.items())[0]
            sharded_ckpt = torch.load(file)

            module.load_state_dict(sharded_ckpt["model"], strict=strict)
            for opt_key, opt in optimizers.items():
                opt.load_state_dict(sharded_ckpt[opt_key])

            # Load anything leftover from sharded_ckpt
            loaded_metadata_keys = sharded_ckpt.keys() - modules.keys() - optimizers.keys()
            requested_metadata_keys = state.keys() - modules.keys() - optimizers.keys()
            _validate_keys_for_strict_loading(requested_metadata_keys, loaded_metadata_keys, strict=strict)
            for key in requested_metadata_keys:
                if key in loaded_metadata_keys:

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Issue one load_checkpoint call per model, each with its own path and single-model state
2. Keep per-model checkpoint directories aligned with per-model save calls

Example fix

# before
fabric.load_checkpoint(path, state={'gen': gen, 'disc': disc})

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

wrapped = [k for k, v in state.items() if isinstance(v, XLAFSDP)]
assert len(wrapped) <= 1, 'load each XLAFSDP model with its own load_checkpoint call'

When it happens

Trigger: fabric.load_checkpoint(path, state={'gen': gen, 'disc': disc}) with both modules XLAFSDP-wrapped (GANs, ensembles, multi-stage pipelines on TPU FSDP).

Common situations: Multi-model TPU training resuming from checkpoints; adapting single-model example code to a GAN setup.

Related errors

• Found multiple XLAFSDP modules in the given state. Saving ch
• Got `XLAFSDPStrategy.load_checkpoint(..., state={state!r})`
• Loading a single module or optimizer object from a checkpoin
• The path {str(file)!r} does not point to valid sharded check
• Could not find a XLAFSDP model in the provided checkpoint st

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