feat: added LinearWarmupCosineAnnealingLRScheduler

Author: le1nux

src/modalities/config/config.py

@@ -227,6 +227,22 @@ class CosineAnnealingLRSchedulerConfig(BaseModel):
     last_epoch: Annotated[int, Field(strict=True, ge=-1)] = -1
 
 
+class LinearWarmupCosineAnnealingLRSchedulerConfig(BaseModel):
+    optimizer: PydanticOptimizerIFType
+    warmup_steps: Annotated[int, Field(strict=True, gt=0)]
+    total_steps: Annotated[int, Field(strict=True, gt=0)]
+    initial_lr: Annotated[float, Field(strict=True, ge=0.0)]
+    final_lr: Annotated[float, Field(strict=True, ge=0.0)]
+    max_lr: Annotated[float, Field(strict=True, ge=0.0)]
+    last_epoch: Annotated[int, Field(strict=True, ge=-1)] = -1
+
+    @model_validator(mode="after")
+    def check_total_steps_greater_than_warmup_steps(self) -> "LinearWarmupCosineAnnealingLRSchedulerConfig":
+        if self.total_steps <= self.warmup_steps:
+            raise ValueError("total_steps must be greater than warmup_steps.")
+        return self
+
+
 class FSDP1CheckpointedOptimizerConfig(BaseModel):
     checkpoint_loading: PydanticFSDP1CheckpointLoadingIFType
     checkpoint_path: Path

src/modalities/optimizers/lr_schedulers.py

@@ -1,21 +1,64 @@
 import warnings
-from typing import Optional
 
+from torch import Tensor
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import LRScheduler
+from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, LRScheduler, SequentialLR
 
 
 class DummyLRScheduler(LRScheduler):
-    def __init__(self, optimizer: Optimizer, last_epoch: Optional[int] = -1):
+    def __init__(self, optimizer: Optimizer, last_epoch: int = -1):
         super().__init__(optimizer, last_epoch)
 
-    def get_lr(self) -> list[float]:
+    def get_lr(self) -> list[float | Tensor]:
         if not self._get_lr_called_within_step:  # type error expected due to internal pytorch implementation
             warnings.warn(
-                "To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.", UserWarning
+                "To get the last learning rate computed by the scheduler, please use `get_last_lr()`.",
+                UserWarning,
             )
 
         return [group["lr"] for group in self.optimizer.param_groups]
 
-    def _get_closed_form_lr(self) -> list[float]:
+    def _get_closed_form_lr(self) -> list[float | Tensor]:
         return self.base_lrs
+
+
+class LRSchedulerFactory:
+    @staticmethod
+    def get_linear_warmup_cosine_annealing_lr_scheduler(
+        optimizer: Optimizer,
+        warmup_steps: int,
+        total_steps: int,
+        initial_lr: float,
+        final_lr: float,
+        max_lr: float,
+        last_epoch: int = -1,
+    ) -> SequentialLR:
+        if warmup_steps <= 0:
+            raise ValueError("warmup_steps must be greater than 0.")
+        if total_steps <= warmup_steps:
+            raise ValueError("total_steps must be greater than warmup_steps.")
+
+        if not all(base_lr == max_lr for base_lr in [group["lr"] for group in optimizer.param_groups]):
+            raise ValueError(
+                "All parameter groups must have the same initial_lr."
+                "and it must be equal to the initial_lr passed to the LR scheduler factory."
+            )
+
+        warmup_scheduler = LinearLR(
+            optimizer=optimizer,
+            start_factor=initial_lr / max_lr,
+            end_factor=1,
+            total_iters=warmup_steps,
+        )
+        cosine_scheduler = CosineAnnealingLR(
+            optimizer=optimizer,
+            T_max=total_steps - warmup_steps,
+            eta_min=final_lr,
+        )
+
+        return SequentialLR(
+            optimizer=optimizer,
+            schedulers=[warmup_scheduler, cosine_scheduler],
+            milestones=[warmup_steps],
+            last_epoch=last_epoch,
+        )

src/modalities/registry/components.py

@@ -48,6 +48,7 @@
     GPT2MFUCalculatorConfig,
     GPT2ModelTPConfig,
     LinearLRSchedulerConfig,
+    LinearWarmupCosineAnnealingLRSchedulerConfig,
     LLMDataLoaderConfig,
     MemMapDatasetConfig,
     OneCycleLRSchedulerConfig,
@@ -108,7 +109,7 @@
     ComposedInitializationRoutines,
     ComposedModelInitializationConfig,
 )
-from modalities.optimizers.lr_schedulers import DummyLRScheduler
+from modalities.optimizers.lr_schedulers import DummyLRScheduler, LRSchedulerFactory
 from modalities.optimizers.optimizer_factory import OptimizerFactory
 from modalities.optimizers.optimizer_list import OptimizersList
 from modalities.optimizers.scheduler_list import SchedulerList
@@ -285,6 +286,12 @@ class ComponentEntity:
         maybe_optimizer_list(torch.optim.lr_scheduler.CosineAnnealingLR),
         CosineAnnealingLRSchedulerConfig,
     ),
+    ComponentEntity(
+        "scheduler",
+        "linear_warmup_cosine_annealing_lr",
+        maybe_optimizer_list(LRSchedulerFactory.get_linear_warmup_cosine_annealing_lr_scheduler),
+        LinearWarmupCosineAnnealingLRSchedulerConfig,
+    ),
     # tokenizers
     ComponentEntity("tokenizer", "pretrained_hf_tokenizer", PreTrainedHFTokenizer, PreTrainedHFTokenizerConfig),
     ComponentEntity("tokenizer", "pretrained_sp_tokenizer", PreTrainedSPTokenizer, PreTrainedSPTokenizerConfig),

tests/test_lr_scheduler.py

@@ -1,14 +1,15 @@
 from unittest.mock import MagicMock, call
 
 import numpy as np
+import torch
 
 from modalities.checkpointing.checkpoint_saving import CheckpointSaving
 from modalities.checkpointing.stateful.app_state import AppState
 from modalities.dataloader.dataloader import LLMDataLoader
 from modalities.evaluator import Evaluator
 from modalities.gym import Gym
 from modalities.loss_functions import Loss
-from modalities.optimizers.lr_schedulers import DummyLRScheduler
+from modalities.optimizers.lr_schedulers import DummyLRScheduler, LRSchedulerFactory
 from modalities.trainer import Trainer
 from tests.utility import configure_dataloader_mock
 
@@ -76,3 +77,27 @@ def test_dummy_lr_scheduler(optimizer_with_param_groups_mock: MagicMock):
     assert np.allclose(scheduler.get_lr(), [0.08, 0.18, 0.28], atol=1e-6)
     assert scheduler._get_closed_form_lr() == [0.1, 0.2, 0.3]
     assert np.allclose(scheduler.get_last_lr(), [0.08, 0.18, 0.28], atol=1e-6)
+
+
+def test_linear_warmup_cosine_annealing_lr_scheduler():
+    parameter = torch.nn.Parameter(torch.tensor([1.0]))
+    optimizer = torch.optim.SGD([parameter], lr=1.0)
+    scheduler = LRSchedulerFactory.get_linear_warmup_cosine_annealing_lr_scheduler(
+        optimizer=optimizer,
+        warmup_steps=2,
+        total_steps=6,
+        initial_lr=0.1,
+        final_lr=0.2,
+        max_lr=1.0,
+    )
+
+    learning_rates = [scheduler.get_last_lr()[0]]
+    for _ in range(6):
+        optimizer.step()
+        scheduler.step()
+        learning_rates.append(scheduler.get_last_lr()[0])
+
+    assert learning_rates[0] < learning_rates[1] < learning_rates[2]
+    assert np.isclose(learning_rates[2], 1.0, atol=1e-6)
+    assert learning_rates[2] > learning_rates[3] > learning_rates[4] > learning_rates[5] > learning_rates[6]
+    assert np.isclose(learning_rates[6], 0.2, atol=1e-6)