Add hyperparameter scan

2026-06-22 07:27:02 +02:00
parent aef0a588ce
commit 46068f356f
6 changed files with 255 additions and 5 deletions
@@ -43,9 +43,10 @@ def run_train_job(
    )
    events_arr = np.array(sorted(train_events))
    n_train_steps = int(np.isin(all_event_ids, events_arr).sum())
+    total_train_batches = n_train_steps // t["batch_size"]
    echo(
        f"  {len(all_event_ids):,} steps  |  "
-        f"{len(train_events)} train events (~{n_train_steps:,} steps)  |  "
+        f"{len(train_events)} train events (~{n_train_steps:,} steps, ~{total_train_batches:,} batches)  |  "
        f"{len(val_events)} val events"
    )

@@ -154,4 +155,5 @@ def run_train_job(
        resume_path=resume,
        validate_every=t["validate_every"],
        validate_steps=t["validate_steps"],
+        total_train_batches=total_train_batches,
    )
@@ -10,6 +10,7 @@ from typing import Callable
 import torch
 import torch.optim as optim
 from torch.utils.data import DataLoader
+from tqdm import tqdm

 from giant.model.schedule import CosineSchedule, flow_matching_loss
 from giant.validate import validate_marginals
@@ -73,6 +74,7 @@ def train(
    resume_path: str | Path | None = None,
    validate_every: int = 0,
    validate_steps: int = 10,
+    total_train_batches: int = 0,
 ) -> None:
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
@@ -108,6 +110,7 @@ def train(
    if write_header:
        metrics_writer.writeheader()

+    epoch_w = len(str(epochs))
    last_completed_epoch = start_epoch - 1
    with _GracefulShutdown() as shutdown:
        for epoch in range(start_epoch, epochs + 1):
@@ -116,7 +119,16 @@ def train(
            model.train()
            train_loss_sum = 0.0
            train_n = 0
-            for cond_cont, cond_cat, x1 in train_loader:
+            ema_loss = 0.0
+            bar = tqdm(
+                train_loader,
+                desc=f"  epoch {epoch:{epoch_w}d}/{epochs}",
+                total=total_train_batches or None,
+                leave=False,
+                unit="batch",
+                dynamic_ncols=True,
+            )
+            for cond_cont, cond_cat, x1 in bar:
                cond_cont = cond_cont.to(device)
                cond_cat = cond_cat.to(device)
                x1 = x1.to(device)
@@ -131,11 +143,19 @@ def train(
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
-                train_loss_sum += loss.item() * x1.size(0)
+                batch_loss = loss.item()
+                train_loss_sum += batch_loss * x1.size(0)
                train_n += x1.size(0)
+                ema_loss = (
+                    batch_loss
+                    if train_n == x1.size(0)
+                    else 0.95 * ema_loss + 0.05 * batch_loss
+                )
+                bar.set_postfix_str(f"loss={ema_loss:.4f}", refresh=False)

                if shutdown.requested:
                    break
+            bar.close()

            if shutdown.requested:
                # Mid-epoch: discard the partial epoch rather than persist an
@@ -163,9 +183,12 @@ def train(
            val_loss = val_loss_sum / max(val_n, 1)
            epoch_time = time.monotonic() - epoch_start

+            is_best = val_loss < best_val_loss
+            marker = "  [best]" if is_best else ""
            print(
-                f"epoch {epoch:4d}  train {train_loss:.4f}  val {val_loss:.4f}  "
-                f"lr {current_lr:.2e}  {epoch_time:.1f}s"
+                f"epoch {epoch:{epoch_w}d}/{epochs}"
+                f"  train {train_loss:.4f}  val {val_loss:.4f}"
+                f"  lr {current_lr:.2e}  {epoch_time:.1f}s{marker}"
            )
            metrics_writer.writerow(
                {
@@ -8,6 +8,7 @@ dependencies = [
    "numpy>=1.26,<3",
    "pandas>=2.2,<4",
    "pyarrow>=16,<25",
+    "tqdm>=4.60,<5",
    "typer>=0.12,<1",
 ]

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+"""Hyperparameter scan over dropout x n_blocks x hidden_dim.
+
+Runs `giant train` sequentially (this machine has a single GPU) for every
+combination, plus one extra run at the default architecture with a higher
+learning rate. Runs are shuffled so the parameter space gets coarse coverage
+early rather than exhausting one corner of the grid first.
+
+Usage:
+    uv run python scripts/hparam_scan.py
+    uv run python scripts/hparam_scan.py --dry-run
+    uv run python scripts/hparam_scan.py --seed 1 --data /path/to/parquet
+"""
+
+import argparse
+import csv
+import itertools
+import os
+import random
+import subprocess
+import time
+from pathlib import Path
+
+DATA_DEFAULT = "/home/lars/geant_steps/train"
+SCAN_DIR_DEFAULT = "checkpoints/scan"
+EPOCHS = 50
+DEFAULT_LR = 3e-4
+
+DROPOUTS = [0.0, 0.1]
+N_BLOCKS = [5, 6, 7, 8]
+HIDDEN_DIMS = [256, 512, 1024]
+
+EXTRA_LR_RUN = {"hidden_dim": 256, "n_blocks": 6, "dropout": 0.1, "lr": 1e-3}
+
+_SUMMARY_FIELDS = [
+    "name",
+    "hidden_dim",
+    "n_blocks",
+    "dropout",
+    "lr",
+    "epochs_completed",
+    "final_val_loss",
+    "best_val_loss",
+    "wall_time_s",
+]
+
+
+def build_runs(seed: int) -> list[dict]:
+    runs = [
+        {"hidden_dim": h, "n_blocks": n, "dropout": d, "lr": DEFAULT_LR}
+        for d, n, h in itertools.product(DROPOUTS, N_BLOCKS, HIDDEN_DIMS)
+    ]
+    runs.append(dict(EXTRA_LR_RUN))
+    random.Random(seed).shuffle(runs)
+    return runs
+
+
+def run_name(run: dict) -> str:
+    return f"h{run['hidden_dim']}_n{run['n_blocks']}_d{run['dropout']}_lr{run['lr']}"
+
+
+def last_completed_epoch(metrics_path: Path) -> int:
+    if not metrics_path.exists():
+        return 0
+    with open(metrics_path, newline="") as f:
+        rows = list(csv.DictReader(f))
+    if not rows:
+        return 0
+    return int(rows[-1]["epoch"])
+
+
+def final_metrics(metrics_path: Path) -> tuple[int, float, float]:
+    with open(metrics_path, newline="") as f:
+        rows = list(csv.DictReader(f))
+    epochs_completed = int(rows[-1]["epoch"])
+    final_val_loss = float(rows[-1]["val_loss"])
+    best_val_loss = min(float(r["val_loss"]) for r in rows)
+    return epochs_completed, final_val_loss, best_val_loss
+
+
+def append_summary(summary_path: Path, row: dict) -> None:
+    write_header = not summary_path.exists()
+    with open(summary_path, "a", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=_SUMMARY_FIELDS)
+        if write_header:
+            writer.writeheader()
+        writer.writerow(row)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--data", default=DATA_DEFAULT)
+    parser.add_argument("--scan-dir", default=SCAN_DIR_DEFAULT)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--dry-run", action="store_true")
+    args = parser.parse_args()
+
+    runs = build_runs(args.seed)
+    scan_dir = Path(args.scan_dir)
+
+    if args.dry_run:
+        for i, run in enumerate(runs, 1):
+            print(f"[{i}/{len(runs)}] {run_name(run)}")
+        return
+
+    scan_dir.mkdir(parents=True, exist_ok=True)
+    summary_path = scan_dir / "scan_summary.csv"
+
+    env = os.environ.copy()
+    env["TQDM_DISABLE"] = "1"
+
+    for i, run in enumerate(runs, 1):
+        name = run_name(run)
+        out_dir = scan_dir / name
+        metrics_path = out_dir / "metrics.csv"
+        last_ckpt = out_dir / "last.pt"
+
+        completed = last_completed_epoch(metrics_path)
+        if completed >= EPOCHS:
+            print(f"[{i}/{len(runs)}] {name} — already complete, skipping")
+            continue
+
+        out_dir.mkdir(parents=True, exist_ok=True)
+        cmd = [
+            "giant",
+            "train",
+            args.data,
+            "--mode",
+            "flow",
+            "--epochs",
+            str(EPOCHS),
+            "--batch-size",
+            "auto",
+            "--hidden-dim",
+            str(run["hidden_dim"]),
+            "--n-blocks",
+            str(run["n_blocks"]),
+            "--dropout",
+            str(run["dropout"]),
+            "--lr",
+            str(run["lr"]),
+            "--out",
+            str(out_dir),
+        ]
+        if last_ckpt.exists():
+            cmd += ["--resume", str(last_ckpt)]
+            print(f"[{i}/{len(runs)}] {name} — resuming from epoch {completed}")
+        else:
+            print(f"[{i}/{len(runs)}] {name} — starting")
+
+        start = time.monotonic()
+        try:
+            with open(out_dir / "train.log", "a") as log:
+                subprocess.run(cmd, env=env, stdout=log, stderr=subprocess.STDOUT)
+        except KeyboardInterrupt:
+            print(
+                f"\ninterrupted during {name} — re-run this script to resume "
+                f"(checkpoint/resume is handled by `giant train` itself)"
+            )
+            return
+        wall_time_s = time.monotonic() - start
+
+        if metrics_path.exists():
+            epochs_completed, final_val_loss, best_val_loss = final_metrics(
+                metrics_path
+            )
+            append_summary(
+                summary_path,
+                {
+                    "name": name,
+                    "hidden_dim": run["hidden_dim"],
+                    "n_blocks": run["n_blocks"],
+                    "dropout": run["dropout"],
+                    "lr": run["lr"],
+                    "epochs_completed": epochs_completed,
+                    "final_val_loss": final_val_loss,
+                    "best_val_loss": best_val_loss,
+                    "wall_time_s": round(wall_time_s, 1),
+                },
+            )
+            print(
+                f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} "
+                f"({wall_time_s:.1f}s)"
+            )
+        else:
+            print(
+                f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log"
+            )
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,18 @@
+import torch
+
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available:  {torch.cuda.is_available()}")
+
+if torch.cuda.is_available():
+    print(f"CUDA version:    {torch.version.cuda}")
+    print(f"Device count:    {torch.cuda.device_count()}")
+    print(f"Device name:     {torch.cuda.get_device_name(0)}")
+
+    # Run a small tensor op on the GPU
+    a = torch.randn(1000, 1000, device="cuda")
+    b = torch.randn(1000, 1000, device="cuda")
+    c = a @ b
+    torch.cuda.synchronize()
+    print(f"Matrix multiply:  OK  (result shape {c.shape}, device {c.device})")
+else:
+    print("No CUDA device found — check driver/CUDA installation.")
@@ -317,6 +317,7 @@ dependencies = [
    { name = "numpy" },
    { name = "pandas" },
    { name = "pyarrow" },
+    { name = "tqdm" },
    { name = "typer" },
 ]

@@ -356,6 +357,7 @@ requires-dist = [
    { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.15,<1" },
    { name = "torch", marker = "extra == 'cpu'", specifier = ">=2.3,<2.4", index = "https://download.pytorch.org/whl/cpu", conflict = { package = "giant", extra = "cpu" } },
    { name = "torch", marker = "extra == 'cuda'", specifier = ">=2.3,<2.4", index = "https://download.pytorch.org/whl/cu118", conflict = { package = "giant", extra = "cuda" } },
+    { name = "tqdm", specifier = ">=4.60,<5" },
    { name = "ty", marker = "extra == 'dev'", specifier = ">=0.0.50,<0.1" },
    { name = "typer", specifier = ">=0.12,<1" },
    { name = "uproot", marker = "extra == 'convert'", specifier = ">=5.3,<6" },
@@ -1253,6 +1255,18 @@ wheels = [
    { url = "https://download-r2.pytorch.org/whl/cu118/torch-2.3.1%2Bcu118-cp312-cp312-win_amd64.whl", hash = "sha256:f44c7b64d990a6b1a382d1cd63c359806153974e7db8d16f6780645a8a9c9fe0", upload-time = "2024-06-05T21:18:59Z" },
 ]

+[[package]]
+name = "tqdm"
+version = "4.68.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/87/d7/0535a28b1f5f24f6612fb3ff1e89fb1a8d160fee0f976e0aa6803862134b/tqdm-4.68.3.tar.gz", hash = "sha256:00dfa48452b6b6cfae3dd9885636c23d3422d1ec97c66d96818cbd5e0821d482", size = 170596, upload-time = "2026-06-17T07:36:52.105Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d8/8e/bb97bb0c71802080bfc8952937d174e49cfc50de5c951dd47b2496f0dcdb/tqdm-4.68.3-py3-none-any.whl", hash = "sha256:39832cc2def2789a6f29df83f172db7416cea70052c0907a57801c5f2fdccb03", size = 78337, upload-time = "2026-06-17T07:36:50.132Z" },
+]
+
 [[package]]
 name = "ty"
 version = "0.0.50"