| """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) |
|
|
