mapgs/train/optim.py

"""Optimizer + LR schedule (§3.2): AdamW (no weight decay on biases / norms /
LayerScale / embeddings), cosine decay with linear warmup."""

from __future__ import annotations

import math

import torch


def build_optimizer(model, lr: float, weight_decay: float):
    decay, no_decay = [], []
    skip_substr = ("temporal_embed", "free_tokens", "template", "type_embed", "ls", "gamma", "norm")
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if p.ndim <= 1 or any(s in name for s in skip_substr):
            no_decay.append(p)
        else:
            decay.append(p)
    groups = [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(groups, lr=lr, betas=(0.9, 0.95))


def cosine_warmup(optimizer, warmup: int, total: int, min_ratio: float = 0.01):
    def fn(step: int) -> float:
        if step < warmup:
            return (step + 1) / max(warmup, 1)
        prog = (step - warmup) / max(total - warmup, 1)
        return min_ratio + (1 - min_ratio) * 0.5 * (1 + math.cos(math.pi * min(prog, 1.0)))

    return torch.optim.lr_scheduler.LambdaLR(optimizer, fn)