| """Optimizer, learning-rate schedule and EMA of weights.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
| from collections.abc import Collection, Iterator |
| from contextlib import contextmanager |
|
|
| import torch |
| from torch import Tensor, nn |
| from torch.optim.lr_scheduler import LambdaLR |
|
|
| from nida.train.config import TrainConfig |
|
|
| __all__ = ["EMA", "build_optimizer", "build_scheduler", "lr_factor"] |
|
|
|
|
| def lr_factor(step: int, *, warmup_steps: int, total_steps: int, min_lr_ratio: float) -> float: |
| """Multiplier of the peak learning rate for the optimizer step number ``step`` (0-based). |
| |
| Linear warmup ``(step + 1) / warmup_steps`` for ``step < warmup_steps`` (so the very |
| first step already has a non-zero rate and step ``warmup_steps - 1`` reaches 1.0), then |
| cosine decay from 1.0 at ``step == warmup_steps`` to ``min_lr_ratio`` at |
| ``step == total_steps``; constant ``min_lr_ratio`` afterwards. |
| """ |
| if step < 0: |
| raise ValueError("step must be >= 0") |
| if step < warmup_steps: |
| return (step + 1) / warmup_steps |
| decay_steps = total_steps - warmup_steps |
| progress = 1.0 if decay_steps <= 0 else min(1.0, (step - warmup_steps) / decay_steps) |
| cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) |
| return min_lr_ratio + (1.0 - min_lr_ratio) * cosine |
|
|
|
|
| def build_optimizer(model: nn.Module, cfg: TrainConfig) -> torch.optim.AdamW: |
| """AdamW with weight decay on weight matrices (``ndim >= 2``) only. |
| |
| Biases, norm scales and embedding tables are exempt (decaying a lookup table shrinks |
| the rows of rare tokens/classes towards zero between their rare updates). |
| """ |
| embedding_ids = { |
| id(module.weight) for module in model.modules() if isinstance(module, nn.Embedding) |
| } |
| decay: list[nn.Parameter] = [] |
| no_decay: list[nn.Parameter] = [] |
| for _, param in model.named_parameters(): |
| if param.requires_grad: |
| exempt = param.ndim < 2 or id(param) in embedding_ids |
| (no_decay if exempt else decay).append(param) |
| o = cfg.optim |
| groups = [ |
| {"params": decay, "weight_decay": o.weight_decay}, |
| {"params": no_decay, "weight_decay": 0.0}, |
| ] |
| return torch.optim.AdamW(groups, lr=o.lr, betas=o.betas, eps=o.eps) |
|
|
|
|
| def build_scheduler(optimizer: torch.optim.Optimizer, cfg: TrainConfig) -> LambdaLR: |
| """``LambdaLR`` over :func:`lr_factor`; call ``scheduler.step()`` once per optimizer step.""" |
| warmup, total, ratio = cfg.schedule.warmup_steps, cfg.steps, cfg.schedule.min_lr_ratio |
|
|
| def factor(step: int) -> float: |
| return lr_factor(step, warmup_steps=warmup, total_steps=total, min_lr_ratio=ratio) |
|
|
| return LambdaLR(optimizer, factor) |
|
|
|
|
| class EMA: |
| """Exponential moving average of a model's weights. |
| |
| ``shadow = decay_n * shadow + (1 - decay_n) * param`` after every optimizer step, where |
| ``decay_n = decay`` or, with ``warmup``, ``min(decay, (1 + n) / (10 + n))`` for the |
| ``n``-th update (0-based) -- a short memory early on so the average is useful before |
| ``1 / (1 - decay)`` steps have passed. Buffers are copied, not averaged. The shadow is |
| stored in float32 on the model's device. |
| """ |
|
|
| def __init__(self, model: nn.Module, *, decay: float, warmup: bool) -> None: |
| if not 0.0 <= decay < 1.0: |
| raise ValueError("decay must be in [0, 1)") |
| self.decay = decay |
| self.warmup = warmup |
| self.num_updates = 0 |
| self.shadow: dict[str, Tensor] = { |
| name: tensor.detach().clone().float() for name, tensor in _state(model).items() |
| } |
|
|
| def current_decay(self) -> float: |
| """Decay that the *next* :meth:`update` will use.""" |
| if not self.warmup: |
| return self.decay |
| n = self.num_updates |
| return min(self.decay, (1.0 + n) / (10.0 + n)) |
|
|
| @torch.no_grad() |
| def update(self, model: nn.Module) -> None: |
| decay = self.current_decay() |
| params = dict(model.named_parameters()) |
| for name, tensor in _state(model).items(): |
| shadow = self.shadow[name] |
| if name in params: |
| shadow.mul_(decay).add_(tensor.detach().float(), alpha=1.0 - decay) |
| else: |
| shadow.copy_(tensor.detach().float()) |
| self.num_updates += 1 |
|
|
| @torch.no_grad() |
| def copy_to(self, model: nn.Module) -> None: |
| """Overwrite ``model``'s weights with the averaged ones.""" |
| for name, tensor in _state(model).items(): |
| tensor.copy_(self.shadow[name].to(tensor.dtype)) |
|
|
| @contextmanager |
| def swapped_in(self, model: nn.Module) -> Iterator[None]: |
| """Temporarily load the EMA weights into ``model`` (restored on exit, also on error).""" |
| backup = {name: t.detach().clone() for name, t in _state(model).items()} |
| self.copy_to(model) |
| try: |
| yield |
| finally: |
| with torch.no_grad(): |
| for name, tensor in _state(model).items(): |
| tensor.copy_(backup[name]) |
|
|
| def state_dict(self) -> dict[str, object]: |
| return { |
| "decay": self.decay, |
| "warmup": self.warmup, |
| "num_updates": self.num_updates, |
| "shadow": {k: v.clone() for k, v in self.shadow.items()}, |
| } |
|
|
| def load_state_dict( |
| self, state: dict[str, object], *, allow_missing: Collection[str] = () |
| ) -> None: |
| """Load ``state``; its names must be exactly this EMA's names. |
| |
| ``allow_missing`` names tensors that ``state`` must NOT contain (a fine-tune that |
| adds parameters, :meth:`nida.train.loop.Trainer.init_from`): they keep their current |
| shadow value. The names missing from ``state`` have to be exactly this set. |
| |
| Raises: |
| ValueError: names, shapes or ``num_updates`` do not match. |
| """ |
| shadow = state["shadow"] |
| allowed = set(allow_missing) |
| if not allowed <= set(self.shadow): |
| raise ValueError( |
| f"allow_missing names unknown to the EMA: {sorted(allowed - set(self.shadow))}" |
| ) |
| if not isinstance(shadow, dict) or set(shadow) != set(self.shadow) - allowed: |
| raise ValueError("EMA state does not match the model's parameter names") |
| for name, tensor in shadow.items(): |
| if tensor.shape != self.shadow[name].shape: |
| raise ValueError(f"EMA tensor {name!r} has shape {tuple(tensor.shape)}") |
| self.shadow[name].copy_(tensor) |
| num_updates = state["num_updates"] |
| if not isinstance(num_updates, int): |
| raise ValueError("EMA num_updates must be an int") |
| self.num_updates = num_updates |
|
|
|
|
| def _state(model: nn.Module) -> dict[str, Tensor]: |
| """Parameters and buffers by name, in a deterministic order.""" |
| out: dict[str, Tensor] = dict(model.named_parameters()) |
| out.update(model.named_buffers()) |
| return out |
|
|
