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

Under ZeRO stage 3, model parameters are partitioned across ranks at creation time, so tensors cannot be materialized with real values on every rank during `__init__`/`configure_model`. `empty_init=False` (which requests non-empty, real initialization) therefore contradicts ZeRO-3 and raises NotImplementedError in tensor_init_context.

Source

Thrown at src/lightning/pytorch/strategies/deepspeed.py:533

        self.optimizers = [deepspeed_optimizer]

        deepspeed_scheduler = model.lr_scheduler
        if deepspeed_scheduler is not None:
            # disable deepspeed lr scheduling as lightning manages scheduling
            model.lr_scheduler = None
            if lr_scheduler is None:
                lr_scheduler = LRSchedulerConfig(deepspeed_scheduler, interval="step")
            else:
                lr_scheduler.scheduler = deepspeed_scheduler
            self.lr_scheduler_configs = [lr_scheduler]
        self.model = model

    @contextmanager
    @override
    def tensor_init_context(self, empty_init: Optional[bool] = None) -> Generator[None, None, None]:
        if self.zero_stage_3:
            if empty_init is False:
                raise NotImplementedError(
                    f"`{empty_init=}` is not a valid choice with `DeepSpeedStrategy` when ZeRO stage 3 is enabled."
                )
            yield
            return
        with super().tensor_init_context(empty_init=empty_init):
            yield

    @contextmanager
    @override
    def model_sharded_context(self) -> Generator[None, None, None]:
        import deepspeed

        self._init_config_if_needed()
        with deepspeed.zero.Init(
            enabled=self.zero_stage_3,
            remote_device=self.remote_device,
            config_dict_or_path=self.config,
        ):

View on GitHub (pinned to 9fed5c27d2)

Solutions

  1. Use the default `empty_init=None`/`True` with ZeRO stage 3 and initialize weights in `configure_model()` or after setup
  2. If real initialization is required, lower the ZeRO stage (e.g. stage 2)
  3. Seed-based determinism: keep empty init and set PL_GLOBAL_SEED so post-init re-initialization is reproducible

Example fix

# before
strategy = DeepSpeedStrategy(config={"zero_optimization": {"stage": 3}})
model = MyModel(..., empty_init=False)

# after
strategy = DeepSpeedStrategy(config={"zero_optimization": {"stage": 3}})
model = MyModel(...)  # init under empty context; real init in configure_model()
Defensive patterns

Strategy: validation

Validate before calling

cfg = {"zero_optimization": {"stage": 3}}
empty_init = None  # never pass False under ZeRO-3
model = L.LightningModule(...)  # rely on default empty init

Prevention

When it happens

Trigger: `DeepSpeedStrategy(config with zero_optimization stage=3)` combined with Lightning's `TensorInitContext` requesting `empty_init=False` — typically set via `Trainer(..., plugins=...)` or Fabric's `init_tensor_workers`/model init with empty_init=False.

Common situations: Users passing `empty_init=False` to get deterministic real init for debugging while enabling ZeRO-3; migrating from FSDP-style configs where empty_init is tunable.

Related errors


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