Lightning-AI/pytorch-lightning · error · ValueError

Could not find a FSDP model in the provided checkpoint state

Error message

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.

What it means

FSDP save_checkpoint requires at least one FSDP-wrapped module in the state dict so it knows what to shard/gather. If no value in state contains FSDP modules (nothing was set up through the strategy), ValueError is raised with guidance on the expected state format.

Source

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

                "`FSDPStrategy.save_checkpoint(..., storage_options=...)` is not supported because"
                " `FSDPStrategy` does not use the `CheckpointIO`."
            )
        if filter is not None and self._state_dict_type == "sharded":
            # https://github.com/pytorch/pytorch/issues/105379
            raise NotImplementedError(
                "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

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Include the set-up model: state = {'model': model, ...} using the module returned by fabric.setup/setup_module
  2. Call fabric.setup(model, optimizer) before saving so the model is FSDP-wrapped
  3. Use fabric.save_checkpoint(path, state) which routes through the same requirement

Example fix

# before
strategy.save_checkpoint(path, {"step": step, "raw_model": raw_model})

# after
model, optimizer = fabric.setup(model, optimizer)
strategy.save_checkpoint(path, {"model": model, "optimizer": optimizer, "step": step})
Defensive patterns

Strategy: validation

Validate before calling

from torch.distributed.fsdp import FullyShardedDataParallel
assert any(
    isinstance(v, FullyShardedDataParallel) or any(isinstance(m, FullyShardedDataParallel) for m in getattr(v, "modules", lambda: [])())
    for v in state.values()
), "state must contain the FSDP-wrapped model under some key"

Type guard

from torch.distributed.fsdp import FullyShardedDataParallel
def state_has_fsdp(state: dict) -> bool:
    for v in state.values():
        if isinstance(v, FullyShardedDataParallel):
            return True
        if hasattr(v, "modules"):
            if any(isinstance(m, FullyShardedDataParallel) for m in v.modules()):
                return True
    return False

Prevention

When it happens

Trigger: strategy.save_checkpoint(path, state={'step': 10, 'loss': 0.5}) with no model entry; or saving the raw model that was never passed through fabric.setup/setup_module (so it is not FSDP-wrapped).

Common situations: Saving metadata-only checkpoints; keeping a reference to the unwrapped model in state instead of the wrapped one returned by setup; calling strategy.save_checkpoint instead of fabric.save_checkpoint before setup.

Related errors


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