Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/strategies/deepspeed.py (Python)

Code before:
        self.initial_scale_power = initial_scale_power
        self.loss_scale_window = loss_scale_window
        self.hysteresis = hysteresis
        self.min_loss_scale = min_loss_scale

        self._deepspeed_engine: Optional[DeepSpeedEngine] = None

    @property
    def zero_stage_3(self) -> bool:
        assert isinstance(self.config, dict)
        zero_optimization = self.config.get("zero_optimization")
        return zero_optimization is not None and zero_optimization.get("stage") == 3

    @property
    @override
    def distributed_sampler_kwargs(self) -> dict[str, int]:
        return {"num_replicas": self.world_size, "rank": self.global_rank}

    @property
    def model(self) -> "DeepSpeedEngine":
        return self._deepspeed_engine

    @override
    def setup_module_and_optimizers(
        self, module: Module, optimizers: list[Optimizer], scheduler: Optional[_LRScheduler] = None
    ) -> tuple["DeepSpeedEngine", list[Optimizer], Optional[_LRScheduler]]:
        """Set up a model and multiple optimizers together, along with an optional learning rate scheduler. Currently,
        only a single optimizer is supported.

        Return:
            The model wrapped into a :class:`deepspeed.DeepSpeedEngine` and a list with a single
            deepspeed optimizer.

        """
        if len(optimizers) != 1:
            raise ValueError(
                f"Currently only one optimizer is supported with DeepSpeed."
                f" Got {len(optimizers)} optimizers instead."
            )

        self._deepspeed_engine, optimizer, scheduler = self._initialize_engine(module, optimizers[0], scheduler)
        self._set_deepspeed_activation_checkpointing()
        return self._deepspeed_engine, [optimizer], scheduler

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

        For training, see :meth:`setup_module_and_optimizers`.

        """

Review comment (suggestion):
This will return `None`, we need to return `Any` here so we can ignore the `scheduler` if it is not provided in input.

Code after:
        self.load_full_weights = load_full_weights

        # default FP16 parameters.
        self.loss_scale = loss_scale
        self.initial_scale_power = initial_scale_power
        self.loss_scale_window = loss_scale_window
        self.hysteresis = hysteresis
        self.min_loss_scale = min_loss_scale

        self._deepspeed_engine: Optional[DeepSpeedEngine] = None

    @property
    def zero_stage_3(self) -> bool:
        assert isinstance(self.config, dict)
        zero_optimization = self.config.get("zero_optimization")
        return zero_optimization is not None and zero_optimization.get("stage") == 3

    @property
    @override
    def distributed_sampler_kwargs(self) -> dict[str, int]:
        return {"num_replicas": self.world_size, "rank": self.global_rank}

    @property
    def model(self) -> "DeepSpeedEngine":
        return self._deepspeed_engine

    @override
    def setup_module_and_optimizers(
        self, module: Module, optimizers: list[Optimizer], scheduler: Optional["_LRScheduler"] = None
    ) -> tuple["DeepSpeedEngine", list[Optimizer], Any]:
        """Set up a model and multiple optimizers together, along with an optional learning rate scheduler. Currently,
        only a single optimizer is supported.

        Return:
            The model wrapped into a :class:`deepspeed.DeepSpeedEngine`, a list with a single
            deepspeed optimizer, and an optional learning rate scheduler.

        """
        if len(optimizers) != 1:
            raise ValueError(
                f"Currently only one optimizer is supported with DeepSpeed. Got {len(optimizers)} optimizers instead."
            )

        self._deepspeed_engine, optimizer, scheduler = self._initialize_engine(module, optimizers[0], scheduler)
        self._set_deepspeed_activation_checkpointing()
        return self._deepspeed_engine, [optimizer], scheduler

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