File size: 4,391 Bytes
e5235fb | 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 | #!/usr/bin/env python3
"""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 # noqa: E402
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)), # validated slots 1-3 (dataviz palette)
("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()
|