Lightning-AI/pytorch-lightning · error · ValueError

Found multiple FSDP models in the given state. Saving checkp

Error message

Found multiple FSDP models 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

The FSDP strategy's save_checkpoint implementation only supports checkpointing one FSDP-wrapped model at a time. When it scans the state dict for FSDP modules it finds more than one, it refuses to save because sharded checkpoint metadata cannot unambiguously represent multiple independent FSDP root modules.

Source

Thrown at src/lightning/fabric/strategies/fsdp.py:466

                "FSDP doesn't support loading sharded filtered checkpoints, so saving them is disabled."
            )

        # broadcast the path from rank 0 to ensure all the states are saved in a common path
        path = _resolve_path(self.broadcast(path))
        if self._state_dict_type == "full" and _is_checkpoint_dir(path) and not _is_sharded_checkpoint(path):
            raise IsADirectoryError(f"The checkpoint path exists and is a directory: {path}")

        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

        modules = [module for module in state.values() if _has_fsdp_modules(module)]
        if len(modules) == 0:
            raise ValueError(
                "Could not find a FSDP 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 FSDP models 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."
            )
        module = modules[0]

        if self._state_dict_type == "sharded":
            _prepare_directory_checkpoint(path)

            state_dict_ctx = _get_sharded_state_dict_context(module)

            # replace the modules and optimizer objects in the state with their local state dict
            # and separate the user's metadata
            converted_state: dict[str, Any] = {}
            metadata: dict[str, Any] = {}
            with state_dict_ctx:
                for key, obj in state.items():
                    converted: Any

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Split the save into separate calls: save_checkpoint(path1, {'model': model1}) and save_checkpoint(path2, {'model': model2})
2. Keep only one FSDP model in the state and pass the other unwrapped or handle it manually
3. Upgrade to torch.distributed.checkpoint usage directly if you truly need one file for multiple sharded models

Example fix

# before
fabric.save('ckpt', {'unet': unet, 'vae': vae})
# after
fabric.save('ckpt/unet', {'unet': unet})
fabric.save('ckpt/vae', {'vae': vae})

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

def count_fsdp_models(state: dict) -> int:
    def has_fsdp(m):
        return isinstance(m, FSDP) or any(isinstance(mod, FSDP) for mod in m.modules()) if hasattr(m, 'modules') else False
    return sum(1 for v in state.values() if has_fsdp(v))

if count_fsdp_models(state) > 1:
    for key, val in state.items():
        fabric.save(f'{out_dir}/{key}', {key: val})

def is_fsdp_model(obj) -> bool:
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
    return isinstance(obj, FSDP) or (hasattr(obj, 'modules') and any(isinstance(m, FSDP) for m in obj.modules()))

When it happens

Trigger: Calling fabric.save(...) / strategy.save_checkpoint(path, state) where state is a dict containing two or more modules that contain FullyShardedDataParallel-wrapped submodules, e.g. {'model': fsdp_model, 'vae': fsdp_vae}.

Common situations: Diffusion-style training with a UNet + text encoder + VAE all wrapped in FSDP in one Fabric setup; stacking a student and teacher model in a single state dict; migrating from DDP where multi-model checkpoints worked fine.

Related errors

• Found multiple FSDP models in the given state. Loading check
• `FSDPStrategy.save_checkpoint(..., storage_options=...)` is
• Could not find a FSDP model in the provided checkpoint state
• Got FSDPStrategy.load_checkpoint(..., state={state!r}) but a
• Loading a single optimizer object from a checkpoint is not s

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