Lightning-AI/pytorch-lightning · error · ValueError

Could not find a DeepSpeed model in the provided checkpoint

Error message

Could not find a DeepSpeed 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.

What it means

Symmetric to saving: DeepSpeed loads checkpoints through the DeepSpeedEngine, so the state passed to load_checkpoint must contain at least one module set up via the strategy. If no engine is found in state, this ValueError is raised.

Source

Thrown at src/lightning/fabric/strategies/deepspeed.py:516

        """
        if isinstance(state, (Module, Optimizer)) or self.load_full_weights and self.zero_stage_3:
            # This code path to enables loading a checkpoint from a non-deepspeed checkpoint or from
            # a consolidated checkpoint
            path = self.broadcast(path)
            return super().load_checkpoint(path=path, state=state, strict=strict, weights_only=weights_only)

        if not state:
            raise ValueError(
                f"Got DeepSpeedStrategy.load_checkpoint(..., state={state!r}) but a state with at least "
                f" a model instance to reload is required. Pass it in like so:"
                " DeepSpeedStrategy.load_checkpoint(..., state={'model': model, ...})"
            )
        _validate_checkpoint_directory(path)

        engines = _get_deepspeed_engines_from_state(state)
        if len(engines) == 0:
            raise ValueError(
                "Could not find a DeepSpeed 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(engines) > 1:
            raise ValueError(
                "Found multiple DeepSpeed engine modules in the given state. Saving and loading checkpoints"
                " with DeepSpeed is currently limited to a single model per checkpoint. To load multiple model"
                " states, call the load method for each model checkpoint separately."
            )
        engine = engines[0]

        from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer

        optimzer_state_requested = any(isinstance(item, (Optimizer, DeepSpeedOptimizer)) for item in state.values())

        torch.cuda.empty_cache()
        _, client_state = engine.load_checkpoint(

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Call fabric.setup(model, optimizer) before fabric.load_checkpoint and include the model in state
  2. Use the same Fabric/strategy instance for setup and loading
  3. For raw weights, load with torch.load + model.load_state_dict instead of the strategy checkpoint API

Example fix

# before
model = MyModel()
fabric.load_checkpoint(path, state={'weights': sd})
# after
model, optimizer = fabric.setup(MyModel(), make_optimizer())
fabric.load_checkpoint(path, state={'model': model, 'optimizer': optimizer})
Defensive patterns

Strategy: validation

Validate before calling

from deepspeed import DeepSpeedEngine
def has_engine(state) -> bool:
    return any(isinstance(v, DeepSpeedEngine) for v in state.values())
model, optimizer = fabric.setup(model, optimizer)
assert has_engine({'model': model})

Type guard

from deepspeed import DeepSpeedEngine

def load_state_valid(state: dict) -> bool:
    return any(isinstance(v, DeepSpeedEngine) for v in state.values())

Prevention

When it happens

Trigger: fabric.load_checkpoint(path, state={'something': x}) where none of the values is a DeepSpeedEngine-backed module — e.g. model not passed through fabric.setup, or state holds only raw tensors/dicts.

Common situations: Resuming with a freshly constructed model that was never set up; calling load_checkpoint before setup; model set up under a different Fabric instance than the one loading.

Related errors


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