Lightning-AI/pytorch-lightning · error · NotImplementedError

The `{type(self).__name__}` does not support setting up the

Error message

The `{type(self).__name__}` does not support setting up the module and optimizer(s) independently. Please call `setup_module_and_optimizers(model, [optimizer, ...])` to jointly set them up.

What it means

DeepSpeed must initialize the model and optimizer together inside deepspeed.initialize (the engine creation allocates optimizer states per ZeRO stage). Therefore setup_optimizer alone cannot work and always raises NotImplementedError with a pointer to setup_module_and_optimizers.

Source

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

    @override
    def setup_module(self, module: Module) -> "DeepSpeedEngine":
        """Set up a module for inference (no optimizers).

        For training, see :meth:`setup_module_and_optimizers`.

        """
        self._deepspeed_engine, _, _ = self._initialize_engine(module)
        return self._deepspeed_engine

    @override
    def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
        """Optimizers can only be set up jointly with the model in this strategy.

        Please use :meth:`setup_module_and_optimizers` to set up both module and optimizer together.

        """
        raise NotImplementedError(self._err_msg_joint_setup_required())

    @override
    def module_init_context(self, empty_init: Optional[bool] = None) -> AbstractContextManager:
        if self.zero_stage_3 and empty_init is False:
            raise NotImplementedError(
                f"`{empty_init=}` is not a valid choice with `DeepSpeedStrategy` when ZeRO stage 3 is enabled."
            )
        module_sharded_ctx = self.module_sharded_context()
        stack = ExitStack()
        if not self.zero_stage_3:
            stack.enter_context(super().module_init_context(empty_init=empty_init))
        stack.enter_context(module_sharded_ctx)
        return stack

    @override
    def module_sharded_context(self) -> AbstractContextManager:
        # Current limitation in Fabric: The config needs to be fully determined at the time of calling the context
        # manager. Later modifications through e.g. `Fabric.setup()` won't have an effect here.

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Replace separate calls with fabric.setup(module, optimizer) / strategy.setup_module_and_optimizers(module, [optimizer])
  2. If you have multiple optimizers, see error 110 — DeepSpeed supports only one
  3. Guard strategy-agnostic code with a check for strategy.requires_joint_setup semantics (isinstance(strategy, DeepSpeedStrategy))

Example fix

# before
module = fabric.setup_module(module)
opt = fabric.setup_optimizer(opt)  # raises under deepspeed
# after
module, opt = fabric.setup(module, opt)
Defensive patterns

Strategy: validation

Validate before calling

from lightning.fabric.strategies.deepspeed import DeepSpeedStrategy
if isinstance(fabric.strategy, DeepSpeedStrategy):
    module, optimizer = fabric.setup(module, optimizer)  # joint setup
else:
    module = fabric.setup_module(module)
    optimizer = fabric.setup_optimizer(optimizer)

Type guard

from lightning.fabric.strategies.deepspeed import DeepSpeedStrategy

def needs_joint_setup(strategy) -> bool:
    return isinstance(strategy, DeepSpeedStrategy)

Try / catch

try:
    optimizer = fabric.setup_optimizer(optimizer)
except NotImplementedError:
    raise RuntimeError("use fabric.setup(module, optimizer) for deepspeed")

Prevention

When it happens

Trigger: Calling fabric.setup_optimizer(optimizer) or strategy.setup_optimizer(optimizer) while using DeepSpeedStrategy; e.g. code that sets up the model first and then calls setup for each optimizer separately.

Common situations: Porting strategy-agnostic Fabric code that assumes setup() and setup_optimizer() can be called independently; looping over optimizers calling setup_optimizer.

Related errors

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