| """WSD learning-rate schedule + EMA wrapper. |
| |
| WSD = Warmup → Steady → Decay. |
| - warmup over [0, w_steps]: linear from 0 to peak_lr |
| - steady over [w_steps, decay_start]: constant peak_lr |
| - decay over [decay_start, total_steps]: 1/sqrt(progress), reaching min_lr |
| """ |
| import math |
| from typing import Optional |
|
|
| import torch |
|
|
|
|
| class WSDSchedule: |
| def __init__(self, peak_lr: float, warmup_steps: int, total_steps: int, |
| decay_frac: float = 0.2, min_lr: float = 1e-5, |
| decay_shape: str = "sqrt", decay_steps: Optional[int] = None): |
| self.peak_lr = peak_lr |
| self.warmup_steps = max(1, warmup_steps) |
| self.total_steps = total_steps |
| if decay_steps is not None: |
| self.decay_start = max(self.warmup_steps + 1, total_steps - decay_steps) |
| else: |
| self.decay_start = max(self.warmup_steps + 1, int(total_steps * (1 - decay_frac))) |
| self.min_lr = min_lr |
| assert decay_shape in ("sqrt", "linear"), decay_shape |
| self.decay_shape = decay_shape |
|
|
| def lr_at(self, step: int) -> float: |
| if step < self.warmup_steps: |
| return self.peak_lr * (step / self.warmup_steps) |
| if step < self.decay_start: |
| return self.peak_lr |
| progress = (step - self.decay_start) / max(1, self.total_steps - self.decay_start) |
| progress = min(1.0, progress) |
| if self.decay_shape == "linear": |
| lr = self.peak_lr * (1.0 - progress) + self.min_lr * progress |
| else: |
| lr = self.peak_lr * (1.0 - math.sqrt(progress)) + self.min_lr * math.sqrt(progress) |
| return max(lr, self.min_lr) |
|
|
| def is_in_decay(self, step: int) -> bool: |
| return step >= self.decay_start |
|
|
| def state_dict(self): |
| return {"peak_lr": self.peak_lr, "warmup_steps": self.warmup_steps, |
| "total_steps": self.total_steps, "decay_start": self.decay_start, |
| "min_lr": self.min_lr, "decay_shape": self.decay_shape} |
|
|
| def load_state_dict(self, sd): |
| self.peak_lr = sd["peak_lr"] |
| self.warmup_steps = sd["warmup_steps"] |
| self.total_steps = sd["total_steps"] |
| self.decay_start = sd["decay_start"] |
| self.min_lr = sd["min_lr"] |
| self.decay_shape = sd.get("decay_shape", "sqrt") |
|
|
|
|
| class EMA: |
| """Exponential-moving-average shadow weights. Only updated when activated |
| (plan: late-only). Lives on the same device as the model.""" |
| def __init__(self, model: torch.nn.Module, decay: float = 0.9999): |
| self.decay = decay |
| self.shadow = {} |
| self.activated = False |
| for n, p in model.named_parameters(): |
| if p.requires_grad: |
| self.shadow[n] = p.detach().clone() |
|
|
| @torch.no_grad() |
| def activate(self, model: torch.nn.Module): |
| self.activated = True |
| for n, p in model.named_parameters(): |
| if p.requires_grad: |
| self.shadow[n].copy_(p.detach()) |
|
|
| @torch.no_grad() |
| def update(self, model: torch.nn.Module): |
| if not self.activated: |
| return |
| for n, p in model.named_parameters(): |
| if not p.requires_grad: |
| continue |
| s = self.shadow[n] |
| s.mul_(self.decay).add_(p.detach(), alpha=1 - self.decay) |
|
|
| @torch.no_grad() |
| def swap_into(self, model: torch.nn.Module): |
| """Copy shadow into model, returning the original weights so they can be |
| restored later.""" |
| backup = {} |
| for n, p in model.named_parameters(): |
| if p.requires_grad: |
| backup[n] = p.detach().clone() |
| p.copy_(self.shadow[n]) |
| return backup |
|
|
| @torch.no_grad() |
| def restore(self, model: torch.nn.Module, backup): |
| for n, p in model.named_parameters(): |
| if n in backup: |
| p.copy_(backup[n]) |
|
|
| def state_dict(self): |
| return {"decay": self.decay, "activated": self.activated, |
| "shadow": {k: v.cpu() for k, v in self.shadow.items()}} |
|
|
| def load_state_dict(self, sd, device="cuda"): |
| self.decay = sd["decay"] |
| self.activated = sd["activated"] |
| for k, v in sd["shadow"].items(): |
| if k in self.shadow: |
| self.shadow[k].copy_(v.to(self.shadow[k].device)) |
|
|
