diff --git a/giant/pipeline.py b/giant/pipeline.py index 5976ae4..c52419f 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -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, ) diff --git a/giant/train.py b/giant/train.py index 539691c..894c47b 100644 --- a/giant/train.py +++ b/giant/train.py @@ -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( { diff --git a/pyproject.toml b/pyproject.toml index ad124be..6128bb8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,7 @@ dependencies = [ "numpy>=1.26,<3", "pandas>=2.2,<4", "pyarrow>=16,<25", + "tqdm>=4.60,<5", "typer>=0.12,<1", ] diff --git a/scripts/hparam_scan.py b/scripts/hparam_scan.py new file mode 100644 index 0000000..ac0cb1f --- /dev/null +++ b/scripts/hparam_scan.py @@ -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() diff --git a/test-cuda.py b/test-cuda.py new file mode 100644 index 0000000..2230890 --- /dev/null +++ b/test-cuda.py @@ -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.") diff --git a/uv.lock b/uv.lock index 3b81c20..646571e 100644 --- a/uv.lock +++ b/uv.lock @@ -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"