[Train] Add warmup-stable-decay learning rate scheduler

src/speculators/train/cli.py

@@ -701,6 +701,10 @@ def main(cfg: TrainConfig):  # noqa: C901
         scheduler_warmup_ratio=args.scheduler_warmup_ratio,
         scheduler_total_steps=args.scheduler_total_steps,
         scheduler_num_cosine_cycles=args.scheduler_num_cosine_cycles,
+        scheduler_warmup_init_lr_ratio=args.scheduler_warmup_init_lr_ratio,
+        scheduler_min_lr_ratio=args.scheduler_min_lr_ratio,
+        scheduler_wsd_decay_ratio=args.scheduler_wsd_decay_ratio,
+        scheduler_wsd_decay_style=args.scheduler_wsd_decay_style,
         checkpoint_freq=args.checkpoint_freq,
         save_best=args.save_best,
         hidden_states_dtype=hidden_states_dtype,

src/speculators/train/config/schema.py

@@ -370,7 +370,7 @@ class OptimizerArgs(_Group):
 
 
 class SchedulerArgs(_Group):
-    scheduler_type: Literal["linear", "cosine", "none"] = Field(
+    scheduler_type: Literal["linear", "cosine", "wsd", "none"] = Field(
         default="linear", description="LR scheduler type."
     )
     scheduler_warmup_steps: int | None = Field(
@@ -387,6 +387,27 @@ class SchedulerArgs(_Group):
     scheduler_num_cosine_cycles: float = Field(
         default=0.5, description="Number of cosine cycles for the cosine scheduler."
     )
+    scheduler_warmup_init_lr_ratio: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="Initial LR as a fraction of peak LR during WSD warmup.",
+    )
+    scheduler_min_lr_ratio: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="Final LR as a fraction of peak LR after WSD decay.",
+    )
+    scheduler_wsd_decay_ratio: float = Field(
+        default=0.2,
+        gt=0.0,
+        le=1.0,
+        description="Fraction of total scheduler steps used by WSD final decay.",
+    )
+    scheduler_wsd_decay_style: Literal[
+        "linear", "cosine", "exponential", "minus_sqrt"
+    ] = Field(default="cosine", description="Decay curve for the final WSD phase.")
 
 
 class TrainerArgs(_Group):

src/speculators/train/trainer.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import math
 import time
 import warnings
 from pathlib import Path
@@ -12,6 +13,7 @@
     set_model_state_dict,
 )
 from torch.nn.parallel import DistributedDataParallel
+from torch.optim.lr_scheduler import LambdaLR
 from torch.utils.data import DataLoader
 from tqdm import TqdmExperimentalWarning
 from tqdm.rich import tqdm
@@ -99,6 +101,8 @@ def profile(self, num_tokens: int) -> dict[str, float] | None:
 
 # Bound rank skew before the validation metrics reduction.
 _VAL_SYNC_INTERVAL = 50
+WSDDecayStyle = Literal["linear", "cosine", "exponential", "minus_sqrt"]
+_WSD_DECAY_STYLES = {"linear", "cosine", "exponential", "minus_sqrt"}
 
 
 def _should_sync_recovery(
@@ -128,11 +132,15 @@ class TrainerConfig(NamedTuple):
     muon_weight_decay: float = 0.1
     muon_ns_steps: int = 5
     muon_adjust_lr_fn: str = "match_rms_adamw"
-    scheduler_type: Literal["linear", "cosine", "none"] = "linear"
+    scheduler_type: Literal["linear", "cosine", "wsd", "none"] = "linear"
     scheduler_warmup_steps: int | None = None
     scheduler_warmup_ratio: float | None = None
     scheduler_total_steps: int | None = None
     scheduler_num_cosine_cycles: float = 0.5
+    scheduler_warmup_init_lr_ratio: float = 0.0
+    scheduler_min_lr_ratio: float = 0.0
+    scheduler_wsd_decay_ratio: float = 0.2
+    scheduler_wsd_decay_style: WSDDecayStyle = "cosine"
     checkpoint_freq: float = 1
     save_best: bool = False
     hidden_states_dtype: torch.dtype = torch.bfloat16
@@ -179,6 +187,93 @@ def _resolve_scheduler_steps(
     return scheduler_warmup_steps, scheduler_total_steps
 
 
+def _validate_wsd_schedule(
+    *,
+    num_warmup_steps: int,
+    num_training_steps: int,
+    warmup_init_lr_ratio: float,
+    min_lr_ratio: float,
+    decay_ratio: float,
+    decay_style: WSDDecayStyle,
+) -> tuple[int, int]:
+    if num_training_steps <= 0:
+        raise ValueError("num_training_steps must be greater than zero.")
+    if not 0 <= num_warmup_steps < num_training_steps:
+        raise ValueError(
+            "num_warmup_steps must be non-negative and smaller than num_training_steps."
+        )
+    if not 0.0 <= warmup_init_lr_ratio <= 1.0:
+        raise ValueError("warmup_init_lr_ratio must be between zero and one.")
+    if not 0.0 <= min_lr_ratio <= 1.0:
+        raise ValueError("min_lr_ratio must be between zero and one.")
+    if not 0.0 < decay_ratio <= 1.0:
+        raise ValueError("decay_ratio must be greater than zero and at most one.")
+    if decay_style not in _WSD_DECAY_STYLES:
+        raise ValueError(f"Unknown WSD decay_style: {decay_style!r}.")
+
+    decay_steps = max(1, int(decay_ratio * num_training_steps))
+    decay_start = num_training_steps - decay_steps
+    if decay_start < num_warmup_steps:
+        raise ValueError(
+            "WSD warmup and final decay phases overlap; reduce num_warmup_steps "
+            "or decay_ratio."
+        )
+    return decay_start, decay_steps
+
+
+def _get_wsd_decay_coefficient(
+    decay_progress: float, decay_style: WSDDecayStyle
+) -> float:
+    if decay_style == "linear":
+        return 1.0 - decay_progress
+    if decay_style == "cosine":
+        return 0.5 * (math.cos(math.pi * decay_progress) + 1.0)
+    if decay_style == "exponential":
+        return 2.0 * math.pow(0.5, decay_progress) - 1.0
+    if decay_style == "minus_sqrt":
+        return 1.0 - math.sqrt(decay_progress)
+    raise AssertionError(f"Unhandled WSD decay style: {decay_style}")
+
+
+def _get_wsd_schedule_with_warmup(
+    optimizer: torch.optim.Optimizer,
+    *,
+    num_warmup_steps: int,
+    num_training_steps: int,
+    warmup_init_lr_ratio: float = 0.0,
+    min_lr_ratio: float = 0.0,
+    decay_ratio: float = 0.2,
+    decay_style: WSDDecayStyle = "cosine",
+) -> LambdaLR:
+    """Create a warmup-stable-decay schedule with a configurable final anneal.
+
+    The decay coefficients follow ``lightseekorg/TorchSpec``'s
+    ``torchspec/training/lr_scheduler.py`` at commit ``4f447655``.
+    """
+    decay_start, decay_steps = _validate_wsd_schedule(
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps,
+        warmup_init_lr_ratio=warmup_init_lr_ratio,
+        min_lr_ratio=min_lr_ratio,
+        decay_ratio=decay_ratio,
+        decay_style=decay_style,
+    )
+
+    def lr_lambda(current_step: int) -> float:
+        if num_warmup_steps > 0 and current_step <= num_warmup_steps:
+            warmup_ratio = current_step / num_warmup_steps
+            return warmup_init_lr_ratio + (1.0 - warmup_init_lr_ratio) * warmup_ratio
+        if current_step <= decay_start:
+            return 1.0
+        if current_step >= num_training_steps:
+            return min_lr_ratio
+        decay_progress = (current_step - decay_start) / decay_steps
+        coefficient = _get_wsd_decay_coefficient(decay_progress, decay_style)
+        return min_lr_ratio + coefficient * (1.0 - min_lr_ratio)
+
+    return LambdaLR(optimizer, lr_lambda)
+
+
 class Trainer:
     def __init__(
         self,
@@ -378,6 +473,16 @@ def make_scheduler(opt: torch.optim.Optimizer):
                     num_training_steps=scheduler_total_steps,
                     last_epoch=last_epoch,
                 )
+            if self.config.scheduler_type == "wsd":
+                return _get_wsd_schedule_with_warmup(
+                    opt,
+                    num_warmup_steps=scheduler_warmup_steps,
+                    num_training_steps=scheduler_total_steps,
+                    warmup_init_lr_ratio=self.config.scheduler_warmup_init_lr_ratio,
+                    min_lr_ratio=self.config.scheduler_min_lr_ratio,
+                    decay_ratio=self.config.scheduler_wsd_decay_ratio,
+                    decay_style=self.config.scheduler_wsd_decay_style,
+                )
             return get_cosine_schedule_with_warmup(
                 opt,
                 num_warmup_steps=scheduler_warmup_steps,

tests/unit/train/test_trainer_scheduler.py

@@ -8,6 +8,9 @@
 from speculators.train.config import TrainConfig
 from speculators.train.trainer import (
     TrainerConfig,
+    WSDDecayStyle,
+    _get_wsd_decay_coefficient,
+    _get_wsd_schedule_with_warmup,
     _resolve_scheduler_steps,
 )
 
@@ -75,6 +78,61 @@ def test_scheduler_type_rejects_unsupported_values():
         )
 
 
+def test_wsd_scheduler():
+    parameter = torch.nn.Parameter(torch.zeros(()))
+    optimizer = torch.optim.AdamW([parameter], lr=1.0)
+    scheduler = _get_wsd_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=2,
+        num_training_steps=10,
+        warmup_init_lr_ratio=0.25,
+        min_lr_ratio=0.1,
+        decay_ratio=0.4,
+        decay_style="linear",
+    )
+
+    observed = [scheduler.get_last_lr()[0]]
+    for _ in range(10):
+        optimizer.step()
+        scheduler.step()
+        observed.append(scheduler.get_last_lr()[0])
+
+    assert observed == pytest.approx(
+        [0.25, 0.625, 1.0, 1.0, 1.0, 1.0, 1.0, 0.775, 0.55, 0.325, 0.1]
+    )
+
+
+@pytest.mark.parametrize(
+    ("decay_style", "expected_midpoint"),
+    [
+        ("linear", 0.5),
+        ("cosine", 0.5),
+        ("exponential", 2.0 * 0.5**0.5 - 1.0),
+        ("minus_sqrt", 1.0 - 0.5**0.5),
+    ],
+)
+def test_wsd_decay_styles(
+    decay_style: WSDDecayStyle,
+    expected_midpoint: float,
+):
+    assert _get_wsd_decay_coefficient(0.5, decay_style) == pytest.approx(
+        expected_midpoint
+    )
+
+
+def test_wsd_rejects_overlapping_phases():
+    parameter = torch.nn.Parameter(torch.zeros(()))
+    optimizer = torch.optim.AdamW([parameter], lr=1.0)
+
+    with pytest.raises(ValueError, match="overlap"):
+        _get_wsd_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=9,
+            num_training_steps=10,
+            decay_ratio=0.2,
+        )
+
+
 def test_scheduler_resume_restores_optimizer_learning_rate(tmp_path: Path):
     checkpoint_dir = tmp_path / "0"
     checkpoint_dir.mkdir()