1"""Training loop: AdamW, linear warm-up + cosine decay, gradient clipping, early stopping on the validation2probabilistic score, fixed seeds, crash-safe checkpoints with resume, JSONL tracking, throughput/memory measurement."""3from __future__ import annotations45import math6import os7import random8import time9from pathlib import Path1011import numpy as np12import torch13from torch.utils.data import DataLoader1415from ..models.model import MRPT, MRPTConfig16from .losses import forecast_loss17from .sampling import OriginDataset, collate18from .tracking import Tracker19from . import checkpoint as ckpt202122def seed_all(seed: int) -> None:23 random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)242526def pick_device(name: str | None = None) -> torch.device:27 if name:28 return torch.device(name)29 if torch.backends.mps.is_available():30 return torch.device("mps")31 if torch.cuda.is_available():32 return torch.device("cuda")33 return torch.device("cpu")343536def to_device(batch: dict[str, torch.Tensor], dev: torch.device) -> dict[str, torch.Tensor]:37 return {k: v.to(dev, non_blocking=True) for k, v in batch.items()}383940def lr_lambda(warmup: int, total: int):41 def f(step):42 if step < warmup:43 return (step + 1) / max(1, warmup)44 p = (step - warmup) / max(1, total - warmup)45 return 0.5 * (1 + math.cos(math.pi * min(1.0, p))) * 0.9 + 0.146 return f474849def mem_mb(dev: torch.device) -> float:50 if dev.type == "mps":51 return torch.mps.current_allocated_memory() / 1e652 if dev.type == "cuda":53 return torch.cuda.max_memory_allocated() / 1e654 return 0.0555657@torch.no_grad()58def evaluate_nll(model: MRPT, loader: DataLoader, dev: torch.device, head: str, levels, max_batches: int | None = None):59 model.eval()60 tot, n = 0.0, 061 for i, b in enumerate(loader):62 if max_batches is not None and i >= max_batches:63 break64 b = to_device(b, dev)65 out = model(b)66 loss, _ = forecast_loss(out, b, head, levels)67 k = int(b["y_mask"].sum())68 tot += float(loss) * k; n += k69 model.train()70 return tot / max(1, n)717273def train(model: MRPT, train_ds: OriginDataset, val_ds: OriginDataset, out_dir: Path, *, batch_size: int = 64,74 max_steps: int = 1000, lr: float = 3e-4, warmup_steps: int = 100, weight_decay: float = 0.01,75 grad_clip: float = 1.0, eval_every: int = 200, patience: int = 5, seed: int = 1, num_workers: int = 4,76 device: str | None = None, aux_vol_weight: float = 0.0, resume: bool = True, val_max_batches: int = 50,77 copula_only: bool = False) -> dict:78 seed_all(seed)79 dev = pick_device(device)80 model.to(dev)81 head = model.cfg.head82 params = [p for n, p in model.named_parameters() if (not copula_only) or n.startswith("copula.")]83 opt = torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay, betas=(0.9, 0.98))84 sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda(warmup_steps, max_steps))85 tracker = Tracker(out_dir / "train_log.jsonl")86 ck_path = out_dir / "last.pt"; best_path = out_dir / "best.pt"87 step, best, bad = 0, float("inf"), 088 if resume and ck_path.exists():89 ck = ckpt.load(ck_path, model, opt, sched, map_location=dev)90 step, best = ck["step"], ck["best"]91 tracker.log(event="resume", step=step, best=best)92 g = torch.Generator(); g.manual_seed(seed + step)93 train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate,94 drop_last=True, generator=g, persistent_workers=num_workers > 0, prefetch_factor=4 if num_workers > 0 else None)95 val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate)96 model.train()97 t0 = time.time(); n_ex = 0; t_data = 0.0; t_last = time.time()98 done = step >= max_steps99 while not done:100 for b in train_loader:101 t_data += time.time() - t_last102 b = to_device(b, dev)103 out = model(b)104 if copula_only:105 from ..models.joint_path import pit_student_t, z_from_pit106 with torch.no_grad():107 u = pit_student_t(b["y"].cpu().numpy(), out["mu"].detach().cpu().numpy(),108 out["scale"].detach().cpu().numpy(), out["nu"].detach().cpu().numpy())109 z = torch.from_numpy(z_from_pit(u).astype(np.float32)).to(dev)110 loss = model.copula.nll(out["hidden"], z, b["y_mask"].bool())111 logs = {"copula_nll": float(loss.detach())}112 else:113 loss, logs = forecast_loss(out, b, head, model.cfg.quantiles, aux_vol_weight)114 if not torch.isfinite(loss):115 tracker.log(event="nonfinite_loss", step=step)116 opt.zero_grad(set_to_none=True); t_last = time.time()117 continue118 opt.zero_grad(set_to_none=True)119 loss.backward()120 gn = torch.nn.utils.clip_grad_norm_(params, grad_clip)121 opt.step(); sched.step(); step += 1; n_ex += batch_size122 if step % 20 == 0:123 el = time.time() - t0124 tracker.log(event="train", step=step, loss=float(loss.detach()), grad_norm=float(gn), lr=sched.get_last_lr()[0],125 ex_per_s=n_ex / el, data_frac=t_data / el, mem_mb=mem_mb(dev), **logs)126 if step % eval_every == 0 or step >= max_steps:127 if copula_only:128 vl = float(loss)129 else:130 vl = evaluate_nll(model, val_loader, dev, head, model.cfg.quantiles, val_max_batches)131 improved = vl < best - 1e-5132 if improved:133 best, bad = vl, 0134 ckpt.save(best_path, model, None, None, step, best, {"val": vl})135 else:136 bad += 1137 ckpt.save(ck_path, model, opt, sched, step, best, {"val": vl})138 tracker.log(event="val", step=step, val_nll=vl, best=best, bad=bad, mem_mb=mem_mb(dev),139 elapsed_s=time.time() - t0)140 if bad >= patience:141 tracker.log(event="early_stop", step=step)142 done = True143 if step >= max_steps:144 done = True145 t_last = time.time()146 if done:147 break148 el = time.time() - t0149 summary = {"steps": step, "best_val": best, "wall_s": el, "ex_per_s": n_ex / max(el, 1e-9),150 "peak_mem_mb": mem_mb(dev), "n_params": model.n_params(), "device": str(dev),151 "data_time_frac": t_data / max(el, 1e-9)}152 tracker.log(event="done", **summary)153 return summary154