| """Figures + summary numbers for the report from pulled metrics (results/{ice,dm}/<run>/metrics.jsonl). |
| |
| Run names: [<prefix>-]<init>-<freeze>-lr<lr>-<budget>[-s<seed>] |
| """ |
|
|
| import argparse |
| import json |
| import re |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
|
|
| COLORS = {"scratch": "#2a78d6", "ivgpt": "#eb6834", "shuffled": "#1baf7a", "rescaled": "#eda100"} |
| STYLES = {"scratch": "-", "ivgpt": "-", "shuffled": "--", "rescaled": "-."} |
| LABELS = {"scratch": "scratch", "ivgpt": "iVideoGPT init", "shuffled": "shuffled iVideoGPT", "rescaled": "rescaled iVideoGPT"} |
| NAME = re.compile(r"(?:(?P<prefix>[^-]+)-)?(?P<init>scratch|ivgpt|shuffled|rescaled|sp[0-9.]+(?:attn|mlp|lower|upper)?|spshuf[0-9.]+|nobody)-(?P<freeze>none|body|lora|top4)-lr(?P<lr>[0-9.e-]+)-(?P<budget>\d+M)(?:-s(?P<seed>\d+))?$") |
|
|
| plt.rcParams.update({ |
| "font.size": 8, "axes.titlesize": 8.5, "axes.labelsize": 8, "legend.fontsize": 7, |
| "axes.spines.top": False, "axes.spines.right": False, "axes.edgecolor": "#52514e", |
| "axes.labelcolor": "#0b0b0b", "xtick.color": "#52514e", "ytick.color": "#52514e", |
| "axes.grid": True, "grid.color": "#e4e3df", "grid.linewidth": 0.5, "lines.linewidth": 1.4, |
| "figure.dpi": 200, "savefig.bbox": "tight", "pdf.fonttype": 42, |
| }) |
|
|
|
|
| def load_runs(roots): |
| runs = [] |
| for root in roots: |
| for d in sorted(Path(root).glob("*/")): |
| m = NAME.match(d.name) |
| f = d / "metrics.jsonl" |
| if not m or not f.exists(): |
| continue |
| recs = [json.loads(l) for l in f.read_text().splitlines() if l.strip()] |
| val = np.array([(r["tokens"], r["val_loss"]) for r in recs if "val_loss" in r]) |
| train = [r for r in recs if "tok_per_s" in r] |
| if len(val) < 2: |
| continue |
| g = m.groupdict() |
| runs.append({ |
| "name": d.name, "src": Path(root).name, "init": g["init"], "freeze": g["freeze"], |
| "lr": float(g["lr"]), "budget": g["budget"], "seed": int(g["seed"] or 0), |
| "done": (d / "DONE").exists(), "tokens": val[:, 0], "val": val[:, 1], |
| "tok_per_s": float(np.median([r["tok_per_s"] for r in train[2:]])) if len(train) > 3 else None, |
| "hardware": (d / "hardware.txt").read_text().strip() if (d / "hardware.txt").exists() else d.parent.name, |
| "evals": load_evals(d / "eval.json"), |
| }) |
| return runs |
|
|
|
|
| EVAL_TASKS = {"piqa": "acc_norm", "hellaswag": "acc_norm", "arc_easy": "acc_norm", "sciq": "acc_norm", |
| "winogrande": "acc", "lambada_openai": "acc"} |
|
|
|
|
| def load_evals(path): |
| if not path.exists(): |
| return None |
| res = json.loads(path.read_text())["results"] |
| return {t: {"v": res[t][f"{m},none"], "se": res[t].get(f"{m}_stderr,none")} for t, m in EVAL_TASKS.items() if t in res} |
|
|
|
|
| def tokens_to_reach(run, target): |
| """Tokens at which the (log-linearly interpolated) val curve first hits target, else None.""" |
| t, v = run["tokens"], run["val"] |
| idx = np.where(v <= target)[0] |
| if len(idx) == 0: |
| return None |
| i = idx[0] |
| if i == 0: |
| return float(t[0]) |
| lt0, lt1 = np.log(max(t[i - 1], 1.0)), np.log(t[i]) |
| frac = (v[i - 1] - target) / (v[i - 1] - v[i]) |
| return float(np.exp(lt0 + frac * (lt1 - lt0))) |
|
|
|
|
| def best_by_arm(runs, budget, freeze): |
| """Best-LR finished run per init (seed 0) at a given budget/freeze.""" |
| best = {} |
| for r in runs: |
| if r["budget"] != budget or r["freeze"] != freeze or r["seed"] != 0 or not r["done"]: |
| continue |
| if r["init"] not in best or r["val"][-1] < best[r["init"]]["val"][-1]: |
| best[r["init"]] = r |
| return best |
|
|
|
|
| def direct_label(ax, runs_by_init, xmax=None): |
| ends = sorted(((r["val"][-1], init) for init, r in runs_by_init.items()), reverse=True) |
| lo, hi = ax.get_ylim() |
| gap = 0.045 * (hi - lo) |
| last_y = None |
| for y, init in ends: |
| ly = y if last_y is None else min(y, last_y - gap) |
| r = runs_by_init[init] |
| ax.annotate(LABELS[init], (r["tokens"][-1], y), xytext=(4, (ly - y) / (hi - lo) * 100), textcoords="offset points", |
| fontsize=6.5, color="#0b0b0b", va="center") |
| last_y = ly |
|
|
|
|
| def fig_curves(runs, budget, freeze, out, title, ylim=None, logx=True): |
| best = best_by_arm(runs, budget, freeze) |
| if not best: |
| return None |
| fig, ax = plt.subplots(figsize=(3.3, 2.3)) |
| for init in ("scratch", "shuffled", "rescaled", "ivgpt"): |
| if init in best: |
| r = best[init] |
| m = r["tokens"] > 0 |
| ax.plot(r["tokens"][m] / 1e6, r["val"][m], STYLES[init], color=COLORS[init], |
| label=f"{LABELS[init]} (lr {r['lr']:g})") |
| if logx: |
| ax.set_xscale("log") |
| if ylim: |
| ax.set_ylim(*ylim) |
| ax.set_xlabel("text tokens (M)") |
| ax.set_ylabel("validation loss (nats/token)") |
| ax.set_title(title) |
| ax.legend(frameon=False, loc="upper right") |
| fig.savefig(out) |
| plt.close(fig) |
| return {k: {"lr": v["lr"], "final": float(v["val"][-1]), "name": v["name"]} for k, v in best.items()} |
|
|
|
|
| def fig_lr(runs, budget, freeze, out): |
| """Final val loss vs peak LR, one line per arm (seed 0).""" |
| table = defaultdict(dict) |
| fig, ax = plt.subplots(figsize=(3.3, 2.3)) |
| for init in ("scratch", "shuffled", "rescaled", "ivgpt"): |
| pts = sorted((r["lr"], float(r["val"][-1])) for r in runs |
| if r["budget"] == budget and r["freeze"] == freeze and r["init"] == init and r["done"] and r["seed"] == 0) |
| for lr, v in pts: |
| table[init][lr] = v |
| if pts: |
| x, y = zip(*pts) |
| ax.plot(x, y, STYLES[init], marker="o", markersize=3.5, color=COLORS[init], label=LABELS[init]) |
| if not table: |
| plt.close(fig) |
| return None |
| ax.set_xscale("log") |
| ax.set_xlabel("peak learning rate") |
| ax.set_ylabel("final validation loss") |
| ax.set_title(f"{'full fine-tune' if freeze == 'none' else 'frozen body'}, {budget} tokens") |
| ax.legend(frameon=False) |
| fig.savefig(out) |
| plt.close(fig) |
| return dict(table) |
|
|
|
|
| def fig_multiplier(runs, budget, freeze, out, seeds=False): |
| """Data multiplier: tokens scratch needs / tokens arm needs to reach the same loss.""" |
| groups = defaultdict(list) |
| best = best_by_arm(runs, budget, freeze) |
| if "scratch" not in best: |
| return None |
| for r in runs: |
| if r["budget"] == budget and r["freeze"] == freeze and r["done"] and r["lr"] == best.get(r["init"], {}).get("lr"): |
| groups[r["init"]].append(r) |
| scratch_runs = groups["scratch"] |
| lo = max(min(r["val"].min() for rs in groups.values() for r in rs), 0) |
| hi = min(r["val"][1] for rs in groups.values() for r in rs) |
| targets = np.linspace(hi, lo + 0.02, 40) |
| fig, ax = plt.subplots(figsize=(3.3, 2.3)) |
| summary = {} |
| for init in ("shuffled", "rescaled", "ivgpt"): |
| if init not in groups: |
| continue |
| ratios = [] |
| for a in groups[init]: |
| for s in scratch_runs: |
| row = [] |
| for tg in targets: |
| ta, ts = tokens_to_reach(a, tg), tokens_to_reach(s, tg) |
| row.append(ts / ta if ta and ts else np.nan) |
| ratios.append(row) |
| ratios = np.array(ratios) |
| med = np.nanmedian(ratios, axis=0) |
| ax.plot(targets, med, STYLES[init], color=COLORS[init], label=LABELS[init]) |
| if len(ratios) > 1: |
| ax.fill_between(targets, np.nanmin(ratios, 0), np.nanmax(ratios, 0), color=COLORS[init], alpha=0.15, lw=0) |
| summary[init] = {f"{t:.2f}": float(m) for t, m in zip(targets[::8], med[::8])} |
| ax.axhline(1.0, color="#52514e", lw=0.8) |
| ax.invert_xaxis() |
| ax.set_xlabel("target validation loss (lower = later in training)") |
| ax.set_ylabel("data multiplier vs scratch") |
| ax.set_title(f"tokens(scratch) / tokens(arm), {budget}") |
| ax.legend(frameon=False) |
| fig.savefig(out) |
| plt.close(fig) |
| return summary |
|
|
|
|
| def fig_delta(runs, budget, freeze, out): |
| """Val loss minus best-LR scratch (seed 0) at matched eval steps; seeds at each arm's best LR, mean +/- range.""" |
| best = best_by_arm(runs, budget, freeze) |
| if "scratch" not in best: |
| return None |
| ref = dict(zip(best["scratch"]["tokens"], best["scratch"]["val"])) |
| fig, ax = plt.subplots(figsize=(3.3, 2.3)) |
| summary = {} |
| for init in ("scratch", "shuffled", "rescaled", "ivgpt"): |
| if init not in best: |
| continue |
| reps = [r for r in runs if r["init"] == init and r["budget"] == budget and r["freeze"] == freeze |
| and r["lr"] == best[init]["lr"] and r["done"]] |
| toks = sorted(set.intersection(*[set(r["tokens"]) for r in reps]) & set(ref)) |
| toks = [t for t in toks if t > 0] |
| d = np.array([[dict(zip(r["tokens"], r["val"]))[t] - ref[t] for t in toks] for r in reps]) |
| x = np.array(toks) / 1e6 |
| ax.plot(x, d.mean(0), STYLES[init], color=COLORS[init], label=f"{LABELS[init]} (n={len(reps)})") |
| if len(reps) > 1: |
| ax.fill_between(x, d.min(0), d.max(0), color=COLORS[init], alpha=0.18, lw=0) |
| summary[init] = {"n": len(reps), "final_mean": float(d[:, -1].mean()), "final_min": float(d[:, -1].min()), |
| "final_max": float(d[:, -1].max()), "lr": best[init]["lr"]} |
| ax.axhline(0.0, color="#52514e", lw=0.8) |
| ax.set_xscale("log") |
| ax.set_xlabel("text tokens (M)") |
| ax.set_ylabel("val loss minus scratch (nats/token)") |
| ax.set_title(f"relative to best-LR scratch, {budget}") |
| ax.legend(frameon=False, fontsize=6.5) |
| fig.savefig(out) |
| plt.close(fig) |
| return summary |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--roots", nargs="+", default=["results/ice", "results/dm_backup"]) |
| p.add_argument("--out", default="report/figs") |
| args = p.parse_args() |
| out = Path(args.out) |
| out.mkdir(parents=True, exist_ok=True) |
| runs = load_runs(args.roots) |
| s = {"runs": [{k: (v if not isinstance(v, np.ndarray) else None) for k, v in r.items()} |
| | {"final_val": float(r["val"][-1]), "final_tokens": float(r["tokens"][-1])} for r in runs]} |
| budgets = sorted({r["budget"] for r in runs}, key=lambda b: int(b[:-1])) |
| for b in budgets: |
| s[f"best_full_{b}"] = fig_curves(runs, b, "none", out / f"curves_full_{b}.pdf", f"full fine-tune, best LR per arm ({b})") |
| s[f"best_frozen_{b}"] = fig_curves(runs, b, "body", out / f"curves_frozen_{b}.pdf", f"frozen body, best LR per arm ({b})") |
| s[f"lr_full_{b}"] = fig_lr(runs, b, "none", out / f"lr_full_{b}.pdf") |
| s[f"lr_frozen_{b}"] = fig_lr(runs, b, "body", out / f"lr_frozen_{b}.pdf") |
| s[f"mult_full_{b}"] = fig_multiplier(runs, b, "none", out / f"multiplier_full_{b}.pdf") |
| s[f"mult_frozen_{b}"] = fig_multiplier(runs, b, "body", out / f"multiplier_frozen_{b}.pdf") |
| s[f"delta_full_{b}"] = fig_delta(runs, b, "none", out / f"delta_full_{b}.pdf") |
| s[f"delta_frozen_{b}"] = fig_delta(runs, b, "body", out / f"delta_frozen_{b}.pdf") |
| (out.parent / "summary.json").write_text(json.dumps(s, indent=1, default=str)) |
| for r in sorted(runs, key=lambda r: r["name"]): |
| print(f"{r['name']:45s} {'done' if r['done'] else 'run '} {r['tokens'][-1] / 1e6:>7.0f}M " |
| f"val={r['val'][-1]:.3f} tok/s={r['tok_per_s']}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
