Lightning-AI/pytorch-lightning · error · ValueError

Found multiple DeepSpeed engine modules in the given state.

Error message

Found multiple DeepSpeed engine modules in the given state. Saving checkpoints with DeepSpeed 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

A DeepSpeed checkpoint maps to exactly one DeepSpeedEngine. If the state passed to save_checkpoint contains multiple set-up models (multiple engines), saving is ambiguous and this ValueError is raised, suggesting separate save calls per model.

Source

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

            raise TypeError(
                "`DeepSpeedStrategy.save_checkpoint(..., storage_options=...)` is not supported because"
                " `DeepSpeedStrategy` does not use the `CheckpointIO`."
            )
        if filter is not None:
            raise TypeError(
                "`DeepSpeedStrategy.save_checkpoint(..., filter=...)` is not supported because"
                " `DeepSpeedStrategy` manages the state serialization internally."
            )

        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: `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(engines) > 1:
            raise ValueError(
                "Found multiple DeepSpeed engine modules in the given state. Saving checkpoints with DeepSpeed 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."
            )
        engine = engines[0]

        # broadcast the path from rank 0 to ensure all the states are saved in a common path
        path = self.broadcast(path)

        # split the checkpoint into two parts:
        # 1) the deepspeed engine encapsulating both the model and optionally the optimizer(s)
        # 2) the rest of the user's state, which in deepspeed is called `client state`
        excluded_objects = (engine, engine.optimizer) if engine.optimizer is not None else (engine,)
        state = {k: v for k, v in state.items() if v not in excluded_objects}
        _validate_state_keys(state)
        # there might be other stateful objects unrelated to the deepspeed engine - convert them to a state_dict
        state = self._convert_stateful_objects_in_state(state, filter={})
        # use deepspeed's internal checkpointing function to handle partitioned weights across processes

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Call fabric.save_checkpoint separately for each model with different paths
  2. Or use separate Fabric/strategy instances per model from the start
  3. Consolidate into a single model (e.g. one module containing submodules) if a single checkpoint is required

Example fix

# before
fabric.save_checkpoint(path, state={'enc': enc, 'dec': dec})
# after
fabric.save_checkpoint(path_enc, state={'enc': enc})
fabric.save_checkpoint(path_dec, state={'dec': dec})
Defensive patterns

Strategy: validation

Validate before calling

from deepspeed import DeepSpeedEngine
engines = [v for v in state.values() if isinstance(v, DeepSpeedEngine)]
if len(engines) > 1:
    for name, model in state.items():
        if isinstance(model, DeepSpeedEngine):
            fabric.save_checkpoint(f"{path}_{name}", {name: model})

Type guard

from deepspeed import DeepSpeedEngine

def single_engine_state(state: dict) -> bool:
    return sum(isinstance(v, DeepSpeedEngine) for v in state.values()) == 1

Prevention

When it happens

Trigger: fabric.save_checkpoint(path, {'model_a': m1, 'model_b': m2}) where both m1 and m2 were wrapped by DeepSpeedStrategy engines.

Common situations: Multi-model pipelines (e.g. encoder+decoder, GANs) under one DeepSpeed strategy instance; Mixture-of-Experts setups with several engines.

Related errors


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