src/ares/training/optimizer.py

"""
Optimizers: AdamW with cosine decay, gradient clipping, flaw checks.
Flaws to catch:
- No weight decay on norm/bias (should exclude)
- No warmup -> loss spike
- No grad clipping -> explosion
- No LR decay -> no convergence
"""
import torch
import math

def get_optimizer(model, lr=3e-4, weight_decay=0.1, betas=(0.9,0.95)):
    # Separate decay / no decay
    decay_params = []
    no_decay_params = []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if "norm" in name or "bias" in name or "embed" in name:
            no_decay_params.append(p)
        else:
            decay_params.append(p)
    optim_groups = [
        {"params": decay_params, "weight_decay": weight_decay},
        {"params": no_decay_params, "weight_decay": 0.0},
    ]
    optimizer = torch.optim.AdamW(optim_groups, lr=lr, betas=betas, eps=1e-8)
    return optimizer

def cosine_schedule_with_warmup(current_step, total_steps, warmup_steps, max_lr, min_lr=3e-5):
    if current_step < warmup_steps:
        return max_lr * (current_step / max(warmup_steps,1))
    # cosine
    progress = (current_step - warmup_steps) / max(1, total_steps - warmup_steps)
    cosine_decay = 0.5 * (1 + math.cos(math.pi * progress))
    return min_lr + (max_lr - min_lr) * cosine_decay

def get_lr_scheduler(optimizer, total_steps, warmup_steps=2000, max_lr=3e-4, min_lr=3e-5):
    from torch.optim.lr_scheduler import LambdaLR
    def lr_lambda(step):
        lr = cosine_schedule_with_warmup(step, total_steps, warmup_steps, max_lr, min_lr)
        return lr / max_lr
    return LambdaLR(optimizer, lr_lambda)