Add learning rate scheduling support for DeepSpeedStrategy

src/lightning/fabric/fabric.py

@@ -32,6 +32,7 @@
 from lightning_utilities.core.overrides import is_overridden
 from torch import Tensor
 from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
 from torch.utils.data import BatchSampler, DataLoader, DistributedSampler, RandomSampler, SequentialSampler
 
 import lightning.fabric
@@ -206,6 +207,7 @@ def setup(
         self,
         module: nn.Module,
         *optimizers: Optimizer,
+        scheduler: Optional[_LRScheduler] = None,
         move_to_device: bool = True,
         _reapply_compile: bool = True,
     ) -> Any:  # no specific return because the way we want our API to look does not play well with mypy
@@ -214,6 +216,7 @@ def setup(
         Args:
             module: A :class:`torch.nn.Module` to set up
             *optimizers: The optimizer(s) to set up (no optimizers is also possible)
+            scheduler: The learning rate scheduler to set up (no learning rate scheduler is also possible)
             move_to_device: If set ``True`` (default), moves the model to the correct device. Set this to ``False``
                 and alternatively use :meth:`to_device` manually.
             _reapply_compile: If ``True`` (default), and the model was ``torch.compile``d before, the
@@ -236,8 +239,8 @@ def setup(
 
         # Let accelerator/plugin wrap and connect the models and optimizers
         if optimizers:
-            module, optimizers = self._strategy.setup_module_and_optimizers(  # type: ignore[assignment]
-                module, list(optimizers)
+            module, optimizers, scheduler = self._strategy.setup_module_and_optimizers(  # type: ignore[assignment]
+                module, list(optimizers), scheduler
             )
         else:
             module = self._strategy.setup_module(module)
@@ -266,7 +269,7 @@ def setup(
 
         if optimizers:
             # join both types in a tuple for API convenience
-            return (module, *optimizers)
+            return (module, *optimizers, scheduler)
         return module
 
     def setup_module(

src/lightning/fabric/strategies/deepspeed.py

@@ -26,6 +26,7 @@
 from lightning_utilities.core.imports import RequirementCache
 from torch.nn import Module
 from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
 from typing_extensions import override
 
 from lightning.fabric.accelerators import Accelerator, CUDAAccelerator
@@ -312,11 +313,10 @@ def model(self) -> "DeepSpeedEngine":
 
     @override
     def setup_module_and_optimizers(
-        self, module: Module, optimizers: list[Optimizer]
-    ) -> tuple["DeepSpeedEngine", list[Optimizer]]:
-        """Set up a model and multiple optimizers together.
-
-        Currently, only a single optimizer is supported.
+        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
@@ -329,9 +329,9 @@ def setup_module_and_optimizers(
                 f" Got {len(optimizers)} optimizers instead."
             )
 
-        self._deepspeed_engine, optimizer = self._initialize_engine(module, optimizers[0])
+        self._deepspeed_engine, optimizer, scheduler = self._initialize_engine(module, optimizers[0], scheduler)
         self._set_deepspeed_activation_checkpointing()
-        return self._deepspeed_engine, [optimizer]
+        return self._deepspeed_engine, [optimizer], scheduler
 
     @override
     def setup_module(self, module: Module) -> "DeepSpeedEngine":
@@ -340,7 +340,7 @@ def setup_module(self, module: Module) -> "DeepSpeedEngine":
         For training, see :meth:`setup_module_and_optimizers`.
 
         """
-        self._deepspeed_engine, _ = self._initialize_engine(module)
+        self._deepspeed_engine, _, _ = self._initialize_engine(module)
         return self._deepspeed_engine
 
     @override
@@ -593,10 +593,8 @@ def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
         )
 
     def _initialize_engine(
-        self,
-        model: Module,
-        optimizer: Optional[Optimizer] = None,
-    ) -> tuple["DeepSpeedEngine", Optimizer]:
+        self, model: Module, optimizer: Optional[Optimizer] = None, scheduler: Optional[_LRScheduler] = None
+    ) -> tuple["DeepSpeedEngine", Optimizer, Optional[_LRScheduler]]:
         """Initialize one model and one optimizer with an optional learning rate scheduler.
 
         This calls ``deepspeed.initialize`` internally.
@@ -605,15 +603,16 @@ def _initialize_engine(
         import deepspeed
 
         model_parameters = filter(lambda p: p.requires_grad, model.parameters())
-        deepspeed_engine, deepspeed_optimizer, _, _ = deepspeed.initialize(
+        deepspeed_engine, deepspeed_optimizer, _, deepspeed_scheduler = deepspeed.initialize(
             args=argparse.Namespace(device_rank=self.root_device.index),
             config=self.config,
             model=model,
             model_parameters=model_parameters,
             optimizer=optimizer,
+            lr_scheduler=scheduler,
             dist_init_required=False,
         )
-        return deepspeed_engine, deepspeed_optimizer
+        return deepspeed_engine, deepspeed_optimizer, deepspeed_scheduler
 
     @override
     def setup_environment(self) -> None:

src/lightning/fabric/strategies/fsdp.py

@@ -33,6 +33,7 @@
 from torch import Tensor
 from torch.nn import Module
 from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
 from typing_extensions import TypeGuard, override
 
 from lightning.fabric.accelerators import Accelerator
@@ -261,8 +262,8 @@ def setup_environment(self) -> None:
 
     @override
     def setup_module_and_optimizers(
-        self, module: Module, optimizers: list[Optimizer]
-    ) -> tuple[Module, list[Optimizer]]:
+        self, module: Module, optimizers: list[Optimizer], scheduler: Optional[_LRScheduler] = None
+    ) -> tuple[Module, list[Optimizer], Optional[_LRScheduler]]:
         """Wraps the model into a :class:`~torch.distributed.fsdp.fully_sharded_data_parallel.FullyShardedDataParallel`
         module and sets `use_orig_params=True` to keep the reference to the original parameters in the optimizer."""
         use_orig_params = self._fsdp_kwargs.get("use_orig_params")
@@ -274,7 +275,7 @@ def setup_module_and_optimizers(
                 " call `setup_optimizer`."
             )
         module = self.setup_module(module)
-        return module, optimizers
+        return module, optimizers, scheduler
 
     @override
     def setup_module(self, module: Module) -> Module:

src/lightning/fabric/strategies/strategy.py

@@ -21,6 +21,7 @@
 from torch import Tensor
 from torch.nn import Module
 from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
 from torch.utils.data import DataLoader
 
 from lightning.fabric.accelerators import Accelerator
@@ -145,8 +146,8 @@ def module_init_context(self, empty_init: Optional[bool] = None) -> AbstractCont
         return stack
 
     def setup_module_and_optimizers(
-        self, module: Module, optimizers: list[Optimizer]
-    ) -> tuple[Module, list[Optimizer]]:
+        self, module: Module, optimizers: list[Optimizer], scheduler: Optional[_LRScheduler] = None
+    ) -> tuple[Module, list[Optimizer], Optional[_LRScheduler]]:
         """Set up a model and multiple optimizers together.
 
         The returned objects are expected to be in the same order they were passed in. The default implementation will
@@ -155,7 +156,7 @@ def setup_module_and_optimizers(
         """
         module = self.setup_module(module)
         optimizers = [self.setup_optimizer(optimizer) for optimizer in optimizers]
-        return module, optimizers
+        return module, optimizers, scheduler
 
     def setup_module(self, module: Module) -> Module:
         """Performs setup for the model, e.g., by wrapping it by another class."""

src/lightning/fabric/strategies/xla_fsdp.py

@@ -21,6 +21,7 @@
 from torch import Tensor
 from torch.nn import Module
 from torch.optim import Optimizer
+from torch.optim.lr_scheduler import _LRScheduler
 from torch.utils.data import DataLoader
 from typing_extensions import override
 
@@ -196,8 +197,8 @@ def setup_environment(self) -> None:
 
     @override
     def setup_module_and_optimizers(
-        self, module: Module, optimizers: list[Optimizer]
-    ) -> tuple[Module, list[Optimizer]]:
+        self, module: Module, optimizers: list[Optimizer], scheduler: Optional[_LRScheduler] = None
+    ) -> tuple[Module, list[Optimizer], Optional[_LRScheduler]]:
         """Returns NotImplementedError since for XLAFSDP optimizer setup must happen after module setup."""
         raise NotImplementedError(
             f"The `{type(self).__name__}` does not support the joint setup of module and optimizer(s)."

src/lightning/fabric/utilities/seed.py

@@ -104,7 +104,10 @@ def pl_worker_init_function(worker_id: int, rank: Optional[int] = None) -> None:
     if _NUMPY_AVAILABLE:
         import numpy as np
 
-        np.random.seed(seed_sequence[3] & 0xFFFFFFFF)  # numpy takes 32-bit seed only
+        ss = np.random.SeedSequence([base_seed, worker_id, global_rank])
+        np_rng_seed = ss.generate_state(4)
+
+        np.random.seed(np_rng_seed)
 
 
 def _generate_seed_sequence(base_seed: int, worker_id: int, global_rank: int, count: int) -> list[int]: