Lightning-AI/pytorch-lightning · error · NotImplementedError

`{empty_init=}` is not a valid choice with `DeepSpeedStrateg

Error message

`{empty_init=}` is not a valid choice with `DeepSpeedStrategy` when ZeRO stage 3 is enabled.

What it means

With DeepSpeed ZeRO stage 3, model parameters are partitioned at creation time, so the module must be materialized inside DeepSpeed's init context — you cannot pre-create a fully-initialized module. module_init_context therefore rejects empty_init=False when zero_stage_3 is enabled.

Source

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

        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.

        import deepspeed

        assert self._config_initialized
        return deepspeed.zero.Init(

View on GitHub (pinned to 9fed5c27d2)

Solutions

1. Create the model inside fabric's init context: with fabric.init_module(): model = MyModel() so ZeRO-3 can partition it at creation
2. Keep zero_stage < 3 if you must pass an already-materialized module
3. For pretrained weights under ZeRO-3, use the deepspeed checkpoint load path or load state after engine init

Example fix

# before
model = MyModel()  # materialized outside
strategy = DeepSpeedStrategy(stage=3)
fabric.setup(model, opt)
# after
strategy = DeepSpeedStrategy(stage=3)
fabric = Fabric(strategy=strategy)
with fabric.init_module(empty_init=True):
    model = MyModel()
model, opt = fabric.setup(model, opt)

from lightning.fabric.strategies.deepspeed import DeepSpeedStrategy
if isinstance(fabric.strategy, DeepSpeedStrategy) and fabric.strategy.zero_stage_3:
    with fabric.init_module(empty_init=True):
        model = MyModel()  # create inside the context
else:
    model = MyModel()

from lightning.fabric.strategies.deepspeed import DeepSpeedStrategy

def must_init_in_context(strategy) -> bool:
    return isinstance(strategy, DeepSpeedStrategy) and strategy.zero_stage_3

try:
    ctx = fabric.strategy.module_init_context(empty_init=False)
except NotImplementedError:
    ctx = fabric.strategy.module_init_context(empty_init=True)

When it happens

Trigger: DeepSpeedStrategy(stage=3) combined with fabric.setup_module(module_created_with_empty_init=False), i.e. passing a module built outside the init context / with meta or normal init while ZeRO-3 partitioning is on; also explicitly passing empty_init=False to module_init_context.

Common situations: Loading a pretrained model then handing it to a ZeRO-3 strategy; code that instantiates the model before fabric.setup and requests real (non-empty) init.

Related errors

• `precision={precision!r})` is not supported in DeepSpeed. `p
• To use the `DeepSpeedStrategy`, you must have DeepSpeed inst
• PyTorch >= 2.6 requires DeepSpeed >= 0.16.0. Detected DeepSp
• Currently only one optimizer is supported with DeepSpeed. Go
• The `{type(self).__name__}` does not support setting up the

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