Repository Public
1"""Resumable trainer: AdamW, warmup + cosine schedule, gradient accumulation and clipping, optional bf16 autocast2(CUDA; MPS stays fp32 by default), early stopping on validation token loss, safetensors checkpoints with full3metadata, JSONL logging, FLOPs/accelerator-time accounting, one-batch overfit test."""4from __future__ import annotations56import json7import math8import time9from dataclasses import asdict, dataclass, field10from pathlib import Path1112import torch13from safetensors.torch import load_file, save_file1415from htr.models.config import HTRConfig16from htr.models.losses import LossWeights, compute_loss17from htr.models.transformer import HTRModel18from htr.utils.io import append_jsonl, ensure_dir, write_json19from htr.utils.provenance import run_header202122@dataclass23class TrainConfig:24 lr: float = 3e-425 weight_decay: float = 0.0126 betas: tuple[float, float] = (0.9, 0.98)27 warmup_steps: int = 50028 max_steps: int = 2000029 accum: int = 130 clip: float = 1.031 eval_every: int = 50032 log_every: int = 2533 patience: int = 6 # evals without improvement34 amp: str = "auto" # auto | bf16 | none35 seed: int = 136 max_tokens: int = 819237 max_batch: int = 4838 loss: LossWeights = field(default_factory=LossWeights)39 device: str = "auto"40 keep_last: int = 2414243def _device(name: str) -> torch.device:44 if name != "auto":45 return torch.device(name)46 if torch.cuda.is_available():47 return torch.device("cuda")48 if torch.backends.mps.is_available():49 return torch.device("mps")50 return torch.device("cpu")515253def _to(batch: dict, dev: torch.device) -> dict:54 out = {}55 for k, v in batch.items():56 if torch.is_tensor(v):57 out[k] = v.to(dev, non_blocking=True)58 elif isinstance(v, dict):59 out[k] = {kk: vv.to(dev) for kk, vv in v.items()}60 else:61 out[k] = v62 return out636465def lr_at(step: int, cfg: TrainConfig) -> float:66 if step < cfg.warmup_steps:67 return cfg.lr * (step + 1) / cfg.warmup_steps68 prog = min(1.0, (step - cfg.warmup_steps) / max(1, cfg.max_steps - cfg.warmup_steps))69 return cfg.lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * prog)))707172def save_checkpoint(model: HTRModel, path: Path, extra: dict) -> None:73 ensure_dir(path)74 state = {k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()}75 if model.cfg.tie_embeddings and "out.weight" in state:76 del state["out.weight"] # tied; restored on load77 save_file(state, str(path / "model.safetensors"), metadata={"format": "pt", "arch": model.cfg.arch, "htr_version": "0.1.0"})78 write_json(path / "config.json", {"model_type": "htr", **model.cfg.to_dict(), "n_params": model.n_params})79 write_json(path / "checkpoint_meta.json", {**run_header(), **extra})808182def load_checkpoint(path: Path, device: str | torch.device = "cpu") -> HTRModel:83 cfg = HTRConfig.from_dict(json.loads((Path(path) / "config.json").read_text()))84 model = HTRModel(cfg)85 state = load_file(str(Path(path) / "model.safetensors"))86 if cfg.tie_embeddings and "out.weight" not in state:87 state["out.weight"] = state["emb.weight"]88 missing, unexpected = model.load_state_dict(state, strict=False)89 missing = [m for m in missing if m != "out.weight"]90 if missing or unexpected:91 raise RuntimeError(f"checkpoint mismatch: missing={missing[:5]} unexpected={unexpected[:5]}")92 return model.to(device).eval()939495class Trainer:96 def __init__(self, model: HTRModel, cfg: TrainConfig, run_dir: Path, train_loader, val_loader):97 self.model, self.cfg, self.run_dir = model, cfg, ensure_dir(run_dir)98 self.train_loader, self.val_loader = train_loader, val_loader99 self.dev = _device(cfg.device)100 self.model.to(self.dev)101 decay, no_decay = [], []102 for n, p in model.named_parameters():103 (no_decay if p.ndim < 2 or "ln" in n or "norm" in n else decay).append(p)104 self.opt = torch.optim.AdamW([{"params": decay, "weight_decay": cfg.weight_decay}, {"params": no_decay, "weight_decay": 0.0}], lr=cfg.lr, betas=cfg.betas, eps=1e-8)105 self.step = 0106 self.best = float("inf")107 self.bad_evals = 0108 self.tokens_seen = 0109 self.accel_seconds = 0.0110 use_bf16 = (cfg.amp == "bf16") or (cfg.amp == "auto" and self.dev.type == "cuda")111 self.autocast = torch.autocast(device_type=self.dev.type, dtype=torch.bfloat16) if use_bf16 else None112 if self.dev.type == "cuda" and hasattr(torch.backends.cuda, "enable_cudnn_sdp"):113 # The cuDNN SDPA backend builds a new execution plan (~3.5 s on an H100, torch 2.11/cu128) for every new114 # (batch, length) shape; with dynamic token-budget batches that dominated the step time (3.3 s vs 0.2 s).115 torch.backends.cuda.enable_cudnn_sdp(False)116 self.log_path = self.run_dir / "train_log.jsonl"117 torch.manual_seed(cfg.seed)118119 # ---- resume120 def resume(self) -> bool:121 ck = self.run_dir / "last"122 st = self.run_dir / "trainer_state.pt"123 if not (ck.exists() and st.exists()):124 return False125 model = load_checkpoint(ck, self.dev)126 self.model.load_state_dict(model.state_dict())127 s = torch.load(st, map_location="cpu", weights_only=False)128 self.opt.load_state_dict(s["opt"])129 self.step, self.best, self.bad_evals, self.tokens_seen, self.accel_seconds = s["step"], s["best"], s["bad_evals"], s["tokens_seen"], s["accel_seconds"]130 return True131132 extra_files: list[Path] = [] # e.g. tokenizer.model, copied into every checkpoint so partial runs stay loadable133134 def _save(self, tag: str, extra: dict) -> None:135 save_checkpoint(self.model, self.run_dir / tag, {"step": self.step, "tokens_seen": self.tokens_seen, "accel_seconds": round(self.accel_seconds, 1), "train_config": asdict(self.cfg), **extra})136 import shutil137138 for f in self.extra_files:139 if Path(f).exists():140 shutil.copy(f, self.run_dir / tag / Path(f).name)141 if tag == "last":142 torch.save({"opt": self.opt.state_dict(), "step": self.step, "best": self.best, "bad_evals": self.bad_evals, "tokens_seen": self.tokens_seen, "accel_seconds": self.accel_seconds}, self.run_dir / "trainer_state.pt")143144 @torch.no_grad()145 def evaluate(self, max_batches: int | None = None) -> dict:146 self.model.eval()147 tot, n, extra = 0.0, 0, {}148 for i, batch in enumerate(self.val_loader):149 if max_batches and i >= max_batches:150 break151 b = _to(batch, self.dev)152 w = LossWeights(ce=1.0, label_smoothing=0.0)153 with (self.autocast or _nullctx()):154 loss, logs = compute_loss(self.model, b, w)155 tot += logs["ce"] * logs["n_tok"]156 n += logs["n_tok"]157 self.model.train()158 ce = tot / max(1, n)159 return {"val_ce": ce, "val_ppl": math.exp(min(20, ce)), "val_tokens": n, **extra}160161 def flops_estimate(self) -> float:162 return 6.0 * self.model.n_params * self.tokens_seen # standard 6ND approximation (training)163164 def train(self) -> dict:165 self.model.train()166 epoch = 0167 t_start = time.time()168 done = False169 while not done:170 if hasattr(self.train_loader.batch_sampler, "set_epoch"):171 self.train_loader.batch_sampler.set_epoch(epoch)172 for i, batch in enumerate(self.train_loader):173 t0 = time.perf_counter()174 b = _to(batch, self.dev)175 with (self.autocast or _nullctx()):176 loss, logs = compute_loss(self.model, b, self.cfg.loss)177 (loss / self.cfg.accum).backward()178 self.tokens_seen += logs["n_tok"] + int(b["src"].ne(self.model.cfg.pad_id).sum())179 if (i + 1) % self.cfg.accum != 0:180 continue181 gn = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.clip)182 for g in self.opt.param_groups:183 g["lr"] = lr_at(self.step, self.cfg)184 self.opt.step()185 self.opt.zero_grad(set_to_none=True)186 if self.dev.type == "mps":187 torch.mps.synchronize()188 elif self.dev.type == "cuda":189 torch.cuda.synchronize()190 self.accel_seconds += time.perf_counter() - t0191 self.step += 1192 if self.dev.type == "mps" and self.step % 20 == 0:193 torch.mps.empty_cache() # MPS keeps per-shape allocations; release them regularly194 if self.step % self.cfg.log_every == 0:195 mem = None196 if self.dev.type == "mps":197 mem = round(torch.mps.driver_allocated_memory() / 2**30, 2)198 elif self.dev.type == "cuda":199 mem = round(torch.cuda.max_memory_allocated() / 2**30, 2)200 rec = {"step": self.step, "epoch": epoch, "lr": lr_at(self.step, self.cfg), "grad_norm": float(gn), "tokens_seen": self.tokens_seen, "elapsed_s": round(time.time() - t_start, 1), "mem_gb": mem, **logs}201 append_jsonl(self.log_path, rec)202 if not math.isfinite(logs["total"]):203 raise RuntimeError(f"non-finite loss at step {self.step}: {logs}")204 if self.step % self.cfg.eval_every == 0 or self.step >= self.cfg.max_steps:205 ev = self.evaluate()206 append_jsonl(self.log_path, {"step": self.step, "eval": ev})207 improved = ev["val_ce"] < self.best - 1e-4208 if improved:209 self.best = ev["val_ce"]210 self.bad_evals = 0211 self._save("best", {"eval": ev})212 else:213 self.bad_evals += 1214 self._save("last", {"eval": ev})215 if self.bad_evals >= self.cfg.patience or self.step >= self.cfg.max_steps:216 done = True217 break218 epoch += 1219 summary = {"steps": self.step, "epochs": epoch, "best_val_ce": self.best, "best_val_ppl": math.exp(min(20, self.best)), "tokens_seen": self.tokens_seen, "accel_seconds": round(self.accel_seconds, 1), "wall_seconds": round(time.time() - t_start, 1), "flops_6ND": self.flops_estimate(), "n_params": self.model.n_params, "device": str(self.dev)}220 write_json(self.run_dir / "train_summary.json", {**run_header(), **summary})221 return summary222223224class _nullctx:225 def __enter__(self):226 return self227228 def __exit__(self, *a):229 return False230231232def overfit_one_batch(model: HTRModel, batch: dict, steps: int = 200, lr: float = 1e-3, device: str = "auto", loss_w: LossWeights | None = None) -> dict:233 """Sanity check: a healthy model drives the token loss on a single batch near zero."""234 dev = _device(device)235 model.to(dev).train()236 b = _to(batch, dev)237 opt = torch.optim.AdamW(model.parameters(), lr=lr)238 w = loss_w or LossWeights(label_smoothing=0.0)239 first = None240 for s in range(steps):241 loss, logs = compute_loss(model, b, w)242 opt.zero_grad()243 loss.backward()244 torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)245 opt.step()246 if first is None:247 first = logs["ce"]248 return {"first_ce": first, "last_ce": logs["ce"], "steps": steps, "device": str(dev)}249