| |
| """Chart for reports/d58_scaling.md and the model card: KL against training tokens, one panel per codebook format |
| (small multiples: the formats' KLs differ 6x, so each panel has its own scale, starting at zero), from |
| reports/d58_curve.csv. Dashed line: llama.cpp's own file of that format (the start). Filled markers: finished files |
| (with an LR decay); hollow: constant-LR checkpoints; line: the finished-point fit (IQ3_S has none: it saturates by |
| ~20M). Needs matplotlib (gen's system python3 has it). |
| |
| python3 scripts/d58_scaling_chart.py --out reports/d58_scaling.png |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| from pathlib import Path |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| REPO = Path(__file__).resolve().parents[1] |
| SURFACE, INK, INK2, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#e6e5e1" |
| SERIES = [("IQ2_XXS", "2-bit IQ2_XXS", "#2a78d6", (0.0775, 0.1354, 0.637)), |
| ("IQ2_S", "2-bit IQ2_S", "#eb6834", (0.0506, 0.0591, 0.622)), |
| ("IQ3_S", "3-bit IQ3_S", "#1baf7a", None)] |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--csv", type=Path, default=REPO / "reports" / "d58_curve.csv") |
| ap.add_argument("--out", type=Path, required=True) |
| a = ap.parse_args() |
| rows = list(csv.DictReader(a.csv.open())) |
| fig, axes = plt.subplots(1, 3, figsize=(12, 4.6), dpi=150) |
| fig.patch.set_facecolor(SURFACE) |
| for ax, (fmt, title, color, fit) in zip(axes, SERIES): |
| ax.set_facecolor(SURFACE) |
| start = next(float(r["kl"]) for r in rows if r["format"] == fmt and r["phase"] == "start") |
| fin = sorted((int(r["tokens"]) / 1e6, float(r["kl"])) for r in rows |
| if r["format"] == fmt and r["phase"] in ("branch", "final")) |
| con = sorted((int(r["tokens"]) / 1e6, float(r["kl"])) for r in rows |
| if r["format"] == fmt and r["phase"] == "constant") |
| ax.axhline(start, color=INK2, lw=1.2, ls=(0, (4, 3)), zorder=1) |
| ax.text(52, start, f"llama.cpp's own file: {start:.3f}", color=INK2, fontsize=8.5, va="bottom", ha="right") |
| if fit: |
| c, k, b = fit |
| xs = [x / 2 for x in range(10, 101)] |
| ax.plot(xs, [c + k * x ** -b for x in xs], color=color, lw=2, zorder=2) |
| else: |
| ax.plot([x for x, _ in fin], [y for _, y in fin], color=color, lw=2, zorder=2) |
| ax.plot([x for x, _ in con], [y for _, y in con], marker="o", ms=8, mfc=SURFACE, mec=color, mew=2, |
| ls="none", zorder=3) |
| ax.plot([x for x, _ in fin], [y for _, y in fin], marker="o", ms=8, color=color, mec=SURFACE, mew=2, |
| ls="none", zorder=4) |
| best = min(fin, key=lambda t: t[1]) |
| ax.annotate(f"best {best[1]:.4f}\n({best[0]:.0f}M tokens, {100 * (1 - best[1] / start):.0f}% below llama.cpp)", |
| best, xytext=(6 if best[0] > 30 else 0, -36), textcoords="offset points", |
| ha="right" if best[0] > 30 else "center", fontsize=8.5, color=INK) |
| ax.set_ylim(0, start * 1.12) |
| ax.set_xlim(0, 53) |
| ax.set_title(title, color=INK, loc="left", fontsize=11) |
| ax.set_xlabel("Training tokens (millions)", color=INK2, fontsize=9) |
| ax.grid(True, color=GRID, lw=0.8) |
| ax.set_axisbelow(True) |
| for s in ("top", "right"): |
| ax.spines[s].set_visible(False) |
| for s in ("left", "bottom"): |
| ax.spines[s].set_color(GRID) |
| ax.tick_params(colors=INK2, labelsize=8.5) |
| axes[0].set_ylabel("KL to the original (lower is better)", color=INK2, fontsize=9) |
| fig.suptitle("Training the codebook scales: KL against training tokens", color=INK, x=0.01, y=0.985, ha="left", fontsize=12) |
| fig.text(0.01, 0.855, "Training curves, measured before attention and embeddings were converted to Q8_0 " |
| "(the shipped files are about 0.001 higher). 40 × 4,096 held-out tokens.\n" |
| "Filled: runs ending with an LR decay. Hollow: constant-LR checkpoints. Line: fit. Each panel has its own " |
| "scale.", color=INK2, fontsize=8.5, va="bottom") |
| fig.tight_layout(rect=(0, 0, 1, 0.83)) |
| a.out.parent.mkdir(parents=True, exist_ok=True) |
| fig.savefig(a.out, facecolor=SURFACE) |
| print(f"wrote {a.out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
