"""Single training entrypoint; ``--config`` selects the stage. Faithful to Transolver ``exp_elas.py``: - AdamW (lr, weight_decay), CosineAnnealingLR(T_max=epochs) - batch_size 1, gradient clipping at max_grad_norm (0.1) - loss = relative-L2 in physical units: predictions are de-normalized before the loss; targets are physical (decode(encode(s)) == s, so storing physical targets is equivalent). Writes a run-log JSON to ``results/`` per master plan ยง0.2. """ from __future__ import annotations import argparse import os import time from typing import Any, Dict, Optional import torch import yaml from .data.dataset import build_splits from .losses.relative_l2 import relative_l2 from .models.transolver import build_model, count_parameters from .seeds import set_seed from .utils.logging import MODAL_RATES_PER_SEC, write_run_log def load_config(path: str) -> Dict[str, Any]: with open(path) as f: return yaml.safe_load(f) def run_training( config: Dict[str, Any], seed: int, data_dir: str, device: Optional[str] = None, gpu_name: str = "CPU", results_path: Optional[str] = None, ckpt_path: Optional[str] = None, log_every: int = 50, max_epochs: Optional[int] = None, ntrain_override: Optional[int] = None, splits=None, ) -> Dict[str, Any]: """Train one model for one seed; return final metrics and write a run-log JSON. If ``ckpt_path`` is given, also save ``{state_dict, normalizer{mean,std}, config, seed, metrics}`` (the normalizer stats are required to de-normalize predictions at inference). If ``splits`` (a ``Splits`` from ``build_splits_from_indices``) is given, it overrides the default first-1000/last-200 split (used by the OOD evaluation). """ device = device or ("cuda" if torch.cuda.is_available() else "cpu") set_seed(seed) data_cfg = config["data"] train_cfg = config["train"] model_cfg = config["model"] if splits is None: ntrain = ntrain_override or data_cfg.get("ntrain", 1000) ntest = data_cfg.get("ntest", 200) splits = build_splits(data_dir, ntrain=ntrain, ntest=ntest) ntest = splits.test_coords.shape[0] normalizer = splits.normalizer.to(device) # GPU-resident dataset: the whole thing is tiny (~10 MB), so we keep it on-device and # batch by index. This removes DataLoader + per-iteration host->device + per-iteration # .item() sync overhead, which dominates wall-clock at batch_size 1. The math is identical # to the DataLoader path (same batch_size, same loss, same seeded shuffle order). batch_size = train_cfg.get("batch_size", 1) eval_every = int(train_cfg.get("eval_every", 1)) def _3d(t): return (t if t.dim() == 3 else t.unsqueeze(-1)).to(device) train_coords = splits.train_coords.to(device) # (ntrain, 972, 2) train_sigma = _3d(splits.train_sigma) # (ntrain, 972, 1) physical test_coords = splits.test_coords.to(device) test_sigma = _3d(splits.test_sigma) ntrain_eff = train_coords.shape[0] base_model = build_model(model_cfg).to(device) n_params = count_parameters(base_model) # Optional torch.compile (CUDA graphs) to cut per-iteration kernel-launch overhead, which # dominates wall-clock at batch_size 1. Same math, static input shape (1, 972, 2). The # checkpoint is saved from base_model so its state_dict keys stay clean (no _orig_mod prefix). model = base_model if bool(train_cfg.get("compile", False)) and device == "cuda": try: model = torch.compile(base_model, mode="reduce-overhead") print(f"[seed {seed}] torch.compile enabled (reduce-overhead)", flush=True) except Exception as e: # pragma: no cover print(f"[seed {seed}] torch.compile failed ({e}); falling back to eager", flush=True) model = base_model lr = float(train_cfg.get("lr", 1e-3)) wd = float(train_cfg.get("weight_decay", 1e-5)) betas = tuple(train_cfg.get("betas", (0.9, 0.999))) epochs = max_epochs or int(train_cfg.get("epochs", 500)) max_grad_norm = train_cfg.get("max_grad_norm", None) optimizer = torch.optim.AdamW(base_model.parameters(), lr=lr, weight_decay=wd, betas=betas) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) shuffle_gen = torch.Generator().manual_seed(seed) # reproducible per-epoch shuffle @torch.no_grad() def eval_test() -> float: model.eval() total = 0.0 for i in range(0, ntest, batch_size): out = normalizer.decode(model(test_coords[i:i + batch_size], None)) total += relative_l2(out, test_sigma[i:i + batch_size], reduction="sum").item() return total / ntest t0 = time.time() best_rel = float("inf") test_rel = float("nan") history = [] for ep in range(epochs): model.train() perm = torch.randperm(ntrain_eff, generator=shuffle_gen).to(device) running = torch.zeros((), device=device) for s in range(0, ntrain_eff, batch_size): idx = perm[s:s + batch_size] optimizer.zero_grad() out = normalizer.decode(model(train_coords[idx], None)) # -> physical loss = relative_l2(out, train_sigma[idx], reduction="sum") loss.backward() if max_grad_norm is not None: torch.nn.utils.clip_grad_norm_(base_model.parameters(), max_grad_norm) optimizer.step() running += loss.detach() scheduler.step() train_rel = (running / ntrain_eff).item() if (ep % eval_every == 0) or (ep >= epochs - 5): test_rel = eval_test() best_rel = min(best_rel, test_rel) history.append({"epoch": ep, "train_rel": train_rel, "test_rel": test_rel}) if ep % log_every == 0 or ep == epochs - 1: print( f"[seed {seed}] epoch {ep:4d} train_rel={train_rel:.5f} test_rel={test_rel:.5f}", flush=True, ) wall = time.time() - t0 rate = MODAL_RATES_PER_SEC.get(gpu_name, 0.0) est_cost = wall * rate final_metrics = { "test_rel_l2": round(test_rel, 6), "best_test_rel_l2": round(best_rel, 6), "train_rel_l2": round(train_rel, 6), "n_params": n_params, "epochs": epochs, } if ckpt_path is not None: os.makedirs(os.path.dirname(ckpt_path) or ".", exist_ok=True) torch.save( { "state_dict": base_model.state_dict(), "normalizer": { "mean": normalizer.mean.detach().cpu(), "std": normalizer.std.detach().cpu(), }, "config": config, "seed": seed, "metrics": final_metrics, }, ckpt_path, ) print(f"[seed {seed}] saved checkpoint -> {ckpt_path}") if results_path is None: os.makedirs("results", exist_ok=True) results_path = os.path.join("results", f"{config.get('name','run')}_seed{seed}.json") write_run_log( path=results_path, config=config, seed=seed, final_metrics=final_metrics, wall_clock_sec=wall, gpu=gpu_name, est_cost_usd=est_cost, extra={"history_tail": history[-5:]}, ) print( f"[seed {seed}] DONE test_rel_l2={test_rel:.6f} best={best_rel:.6f} " f"params={n_params} wall={wall:.0f}s gpu={gpu_name} est_cost=${est_cost:.4f}" ) return final_metrics def main() -> int: ap = argparse.ArgumentParser(description="Train the stress operator (one seed).") ap.add_argument("--config", required=True) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--data-dir", default="data") ap.add_argument("--device", default=None) ap.add_argument("--gpu-name", default="CPU", help="for cost accounting (A10 / A100-40GB / CPU)") ap.add_argument("--results-path", default=None) ap.add_argument("--max-epochs", type=int, default=None, help="override epochs (local smoke runs)") ap.add_argument("--ntrain", type=int, default=None, help="override ntrain (local smoke runs)") args = ap.parse_args() config = load_config(args.config) run_training( config=config, seed=args.seed, data_dir=args.data_dir, device=args.device, gpu_name=args.gpu_name, results_path=args.results_path, max_epochs=args.max_epochs, ntrain_override=args.ntrain, ) return 0 if __name__ == "__main__": raise SystemExit(main())