File size: 5,802 Bytes
026f63e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """Pre-flight training time / cost estimate from a *measured* throughput.
Runs the real training loop (same data pipeline, model, micro-batch, AMP,
compile) for `bench_seconds` on the current machine, takes the steady-state
samples/s (first logging window dropped: compile / warm-up), then extrapolates.
Run it on the rented instance itself for GPU numbers - measuring locally and
dividing by a guessed speed-up is exactly the kind of invented number to avoid.
"""
from __future__ import annotations
import csv
import copy
import tempfile
import time
from pathlib import Path
from minijev.model import human_params
def measure_throughput(cfg: dict, seconds: float = 60.0) -> dict:
from minijev.training.supervised import train
cfg = copy.deepcopy(cfg)
data_dir = Path(cfg["data"]["dir"])
tmp = Path(tempfile.mkdtemp(prefix="minijev-est-"))
if not (data_dir / "train.npz").exists():
# same generator settings, small size: per-row encoding cost is identical
gen = dict(cfg["data"].get("generate", {}))
gen["samples"] = 30_000
cfg["data"] = {"dir": str(tmp / "data"), "generate": gen}
tcfg = cfg["training"]
tcfg["log_every"] = 10
tcfg["eval_every"] = 10**9
tcfg["checkpoint_every"] = 10**9
s = train(cfg, run_dir=tmp / "run", max_minutes=seconds / 60, quiet=True, final_eval=False)
rows = list(csv.DictReader(open(tmp / "run" / "metrics.csv")))
rates = [float(r["samples_per_sec"]) for r in rows if r.get("samples_per_sec")]
steady = rates[1:] or rates
sps = sorted(steady)[len(steady) // 2] if steady else s["throughput_samples_per_sec"]
return {"samples_per_sec": sps, "runtime": s["runtime"], "params": s["params"], "train_rows": s["train_rows"],
"windows": len(rates)}
def measure_generation(cfg: dict, n: int = 5000) -> float | None:
"""Rows/s of oracle data generation on this machine (all cores)."""
gen = cfg["data"].get("generate")
if not gen:
return None
import importlib
mod = importlib.import_module(f"environments.{cfg['task']['name']}.dataset")
g = {k: v for k, v in gen.items() if k not in ("samples", "seed")}
t = time.time()
mod.generate(n, **g)
return n / (time.time() - t)
def estimate(cfg: dict, hourly_cost: float, throughput: float | None = None, bench_seconds: float = 60,
gpu_name: str | None = None, budget: float | None = None, overhead_minutes: float = 3.0) -> dict:
from minijev.device import resolve_runtime
from minijev.tasks import get_task
from minijev.training.supervised import make_model
task = get_task(cfg["task"]["name"], cfg["task"])
params = make_model(cfg, task).num_parameters()
data_dir = Path(cfg["data"]["dir"])
if (data_dir / "train.npz").exists():
import numpy as np
rows = int(len(np.load(data_dir / "train.npz")["target"]))
gen_rows = 0
else:
total = int(cfg["data"].get("generate", {}).get("samples", 0))
rows, gen_rows = int(total * 0.9), total
if cfg["data"].get("train_limit"):
rows = min(rows, int(cfg["data"]["train_limit"]))
epochs = float(cfg["training"].get("epochs", 1))
prof = resolve_runtime(cfg["training"], params)
measured = throughput is None
if measured:
print(f"[estimate] measuring real throughput for {bench_seconds:.0f}s on {prof.device} ({prof.device_name}) ...")
m = measure_throughput(cfg, bench_seconds)
throughput = m["samples_per_sec"]
gen_rate = measure_generation(cfg) if gen_rows else None
train_s = rows * epochs / throughput
evals = max(1, int(rows * epochs / cfg["training"].get("batch_size", 256) / cfg["training"].get("eval_every", 1000)))
eval_rows = min(cfg["training"].get("eval_max_rows", 20000), int(rows / 18))
eval_s = evals * eval_rows / (3 * throughput) # inference ~3x cheaper than a training step
gen_s = gen_rows / gen_rate if gen_rate else 0.0
total_s = train_s + eval_s + gen_s + overhead_minutes * 60
hours = total_s / 3600
cost = hours * hourly_cost
out = {"model": cfg.get("name"), "params": params, "train_rows": rows, "epochs": epochs,
"device": gpu_name or prof.device_name, "throughput_samples_per_sec": round(throughput, 1),
"throughput_measured": measured, "train_minutes": round(train_s / 60, 1),
"data_generation_minutes": round(gen_s / 60, 1), "eval_minutes": round(eval_s / 60, 1),
"overhead_minutes": overhead_minutes, "total_minutes": round(total_s / 60, 1),
"gpu_hours": round(hours, 3), "hourly_cost": hourly_cost, "estimated_cost": round(cost, 3)}
print()
print(f"Model: MiniJEV {cfg.get('name')}")
print(f"Parameters: {human_params(params)}")
print(f"Dataset: {rows:,} training examples" + (f" (+ generating {gen_rows:,})" if gen_rows else ""))
print(f"Epochs: {epochs:g}")
print()
print(f"Device: {out['device']} (micro-batch {prof.micro_batch} x accum {prof.grad_accum}, amp={prof.amp_dtype})")
print(f"{'Measured' if measured else 'Given'} throughput: {throughput:,.0f} samples/sec")
print()
print(f"Estimated duration: {_fmt_min(total_s / 60)} (train {_fmt_min(train_s / 60)}, "
f"data gen {_fmt_min(gen_s / 60)}, eval {_fmt_min(eval_s / 60)}, overhead {overhead_minutes:g} min)")
print(f"Estimated GPU hours: {hours:.3f}")
print(f"Estimated cost: €{cost:.2f} (at €{hourly_cost:.3f}/h)")
if budget is not None and cost > budget:
scale = budget / cost
print(f"\n!! Over budget (€{budget:.2f}). Reduce e.g. epochs to {epochs * scale:.2f} "
f"or train rows to {int(rows * scale):,}.")
return out
def _fmt_min(m: float) -> str:
return f"{m:.1f} min" if m < 120 else f"{m / 60:.1f} h"
|