File size: 4,858 Bytes
93cd35a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """Training-time utilities: train state, optimizer/schedule helpers.
Extracted from train.py for reuse and readability.
"""
import math
import queue
import threading
from dataclasses import dataclass, field
from typing import Any, Dict, Optional
import torch
import torch.nn as nn
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR
from utils.logging_utils import log_for_0
# ============================================
# Train State with EMA
# ============================================
@dataclass
class TrainState:
"""Lightweight container for the trainable model and its EMA shadow."""
model: nn.Module
optimizer: Optimizer
lr_scheduler: Any = None
ema_params1: Dict[str, torch.Tensor] = field(default_factory=dict)
step: int = 0
epoch: int = 0
dropout_generator: Optional[torch.Generator] = None
def replace(self, **kwargs) -> "TrainState":
new = TrainState(
model=self.model, optimizer=self.optimizer, lr_scheduler=self.lr_scheduler,
ema_params1=self.ema_params1, step=self.step, epoch=self.epoch,
dropout_generator=self.dropout_generator,
)
for k, v in kwargs.items():
setattr(new, k, v)
return new
@staticmethod
def init_ema(model: nn.Module) -> Dict[str, torch.Tensor]:
return {k: v.detach().clone() for k, v in model.named_parameters()}
def prefetch_to_device(iterator, size: int = 2):
"""Prefetch batches asynchronously via a background thread."""
q = queue.Queue(maxsize=size)
def enqueue():
for item in iterator:
q.put(item)
q.put(None)
threading.Thread(target=enqueue, daemon=True).start()
while True:
item = q.get()
if item is None:
break
yield item
# ============================================
# Optimizer
# ============================================
def get_optimizer(model: nn.Module, config, lr: float, grad_accum_steps: int = 1):
"""Build optimizer (AdamW or Muon). Gradient clipping is applied in train.py."""
if config.optimizer == "muon":
from utils.muon_utils import muon_with_aux_adam
opt = muon_with_aux_adam(model, lr=lr)
log_for_0("Using Muon optimizer")
elif config.optimizer == "adamw":
params = [p for p in model.parameters() if p.requires_grad]
opt = torch.optim.AdamW(
params, lr=lr, weight_decay=config.weight_decay,
betas=(config.adam_b1, config.adam_b2),
)
log_for_0("Using AdamW optimizer")
else:
raise ValueError(f"Unknown optimizer: {config.optimizer}. Choose 'adamw' or 'muon'.")
return opt
# ============================================
# Learning Rate Schedule
# ============================================
def create_learning_rate_fn(
num_train_steps: int,
num_warmup_steps: int,
learning_rate: float,
schedule: str = "constant",
min_lr: float = 0.0,
):
"""Create learning rate schedule with linear warmup."""
alpha = (min_lr / learning_rate) if learning_rate > 0 else 0.0
def fn(step: int) -> float:
step = int(step)
if num_warmup_steps > 0 and step < num_warmup_steps:
return learning_rate * step / max(1, num_warmup_steps)
if schedule == "cosine":
progress = (step - num_warmup_steps) / max(1, num_train_steps - num_warmup_steps)
progress = min(max(progress, 0.0), 1.0)
cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
return learning_rate * (alpha + (1.0 - alpha) * cosine)
return learning_rate
return fn
def attach_lr_scheduler(optimizer: Optimizer, lr_fn) -> LambdaLR:
"""Wrap `lr_fn` in a LambdaLR (lambda returns multiplier on base lr)."""
base_lr = optimizer.param_groups[0]["lr"]
return LambdaLR(optimizer, lr_lambda=lambda step: lr_fn(step) / max(base_lr, 1e-12))
# ============================================
# EMA update
# ============================================
def unwrap_model(model: nn.Module) -> nn.Module:
"""Strip DDP (`.module`) and `torch.compile` (`._orig_mod`) wrappers."""
seen = set()
while id(model) not in seen:
seen.add(id(model))
if hasattr(model, "_orig_mod"):
model = model._orig_mod
elif hasattr(model, "module") and isinstance(model.module, nn.Module):
model = model.module
else:
break
return model
@torch.no_grad()
def ema_update(ema_state: Dict[str, torch.Tensor], model: nn.Module, decay: float) -> None:
"""In-place EMA over trainable params: `ema = decay*ema + (1-decay)*param`."""
inner = unwrap_model(model)
for name, param in inner.named_parameters():
if name in ema_state:
ema_state[name].lerp_(param.detach(), 1.0 - decay)
|