train/optimizer.py

"""Optimizer and scheduler factories."""

from __future__ import annotations

import math
from dataclasses import dataclass

import torch


@dataclass(frozen=True)
class ScheduleConfig:
    """Training schedule settings."""

    peak_learning_rate: float = 3.0e-4
    min_learning_rate: float = 3.0e-5
    warmup_steps: int = 2000
    weight_decay: float = 0.1
    betas: tuple[float, float] = (0.9, 0.95)
    adam_eps: float = 1.0e-8
    total_steps: int = 25_000


def create_optimizer(model: torch.nn.Module, config: ScheduleConfig) -> torch.optim.Optimizer:
    """Create an AdamW optimizer with correct weight-decay exclusions."""
    decay: list[torch.nn.Parameter] = []
    no_decay: list[torch.nn.Parameter] = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if param.ndim == 1 or "norm" in name:
            no_decay.append(param)
        else:
            decay.append(param)
    return torch.optim.AdamW(
        [
            {"params": decay, "weight_decay": config.weight_decay},
            {"params": no_decay, "weight_decay": 0.0},
        ],
        lr=config.peak_learning_rate,
        betas=config.betas,
        eps=config.adam_eps,
    )


def lr_lambda(current_step: int, config: ScheduleConfig) -> float:
    """Warm up linearly and then decay with cosine."""
    if current_step < config.warmup_steps:
        return float(current_step + 1) / float(max(1, config.warmup_steps))
    progress = (current_step - config.warmup_steps) / float(max(1, config.total_steps - config.warmup_steps))
    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
    floor = config.min_learning_rate / config.peak_learning_rate
    return floor + (1.0 - floor) * cosine


def create_scheduler(optimizer: torch.optim.Optimizer, config: ScheduleConfig) -> torch.optim.lr_scheduler.LambdaLR:
    """Create the training LR scheduler."""
    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step: lr_lambda(step, config))