diff --git a/src/speculators/train/cli.py b/src/speculators/train/cli.py index 2f20306e5..8e0ee11bb 100644 --- a/src/speculators/train/cli.py +++ b/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, diff --git a/src/speculators/train/config/schema.py b/src/speculators/train/config/schema.py index 3dc678c5b..9046a40e3 100644 --- a/src/speculators/train/config/schema.py +++ b/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): diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py index ee34500c5..ee4e85960 100644 --- a/src/speculators/train/trainer.py +++ b/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, diff --git a/tests/unit/train/test_trainer_scheduler.py b/tests/unit/train/test_trainer_scheduler.py index b324983a7..fbeeddf5a 100644 --- a/tests/unit/train/test_trainer_scheduler.py +++ b/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()