File size: 4,832 Bytes
13caab9 | 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 | """Loss curve for GPT-U-20M: training loss (logged value + 100-step mean) and validation loss per domain."""
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
SURFACE, INK, INK_2, MUTED, GRID, AXIS = "#fcfcfb", "#0b0b0b", "#52514e", "#898781", "#e1e0d9", "#c3c2b7"
DOMAIN_COLORS = {"dclm": "#2a78d6", "edu": "#eb6834", "code": "#1baf7a"} # categorical slots 1-3, fixed order
PHASES = ((500, "warmup ends"), (15_872, "decay starts"))
def read_log(path: Path) -> tuple[list, list]:
"""Last record per step wins (a resumed run re-logs the steps after its checkpoint)."""
train, evals = {}, {}
for line in path.read_text(encoding="utf-8").splitlines():
rec = json.loads(line)
if rec.get("event") == "eval":
evals[rec["step"]] = rec
elif "event" not in rec:
train[rec["step"]] = rec["loss"]
return sorted(train.items()), sorted(evals.items())
def _style(ax, title: str, ylabel: str) -> None:
ax.set_facecolor(SURFACE)
ax.set_title(title, loc="left", color=INK, fontsize=11)
ax.set_xlabel("step (131,072 tokens each)", color=INK_2, fontsize=9)
ax.set_ylabel(ylabel, color=INK_2, fontsize=9)
ax.grid(axis="y", color=GRID, linewidth=0.5)
ax.set_axisbelow(True)
for side in ("top", "right"):
ax.spines[side].set_visible(False)
for side in ("left", "bottom"):
ax.spines[side].set_color(AXIS)
ax.tick_params(colors=MUTED, labelsize=8)
def _phase_lines(ax, x_max: int) -> None:
for step, label in PHASES:
if step <= x_max:
ax.axvline(step, color=AXIS, linewidth=0.6, linestyle=(0, (3, 3)))
ax.annotate(label, (step, 1), xycoords=("data", "axes fraction"), xytext=(3, -10),
textcoords="offset points", color=MUTED, fontsize=7)
def plot_loss(log_path: Path, out_path: Path, total_steps: int) -> None:
train, evals = read_log(Path(log_path))
fig, (ax_t, ax_v) = plt.subplots(1, 2, figsize=(12, 4.6), dpi=150, facecolor=SURFACE)
_style(ax_t, "Training loss", "cross-entropy (nats/token)")
_style(ax_v, "Validation loss by domain", "cross-entropy (nats/token)")
x_max = max([s for s, _ in train] + [s for s, _ in evals] + [1])
if train:
steps, loss = zip(*train)
window = 10 # log points are 10 steps apart -> 100-step mean
mean = [sum(loss[max(0, i - window + 1):i + 1]) / len(loss[max(0, i - window + 1):i + 1]) for i in range(len(loss))]
ax_t.plot(steps, loss, color=AXIS, linewidth=0.6, label="logged loss")
ax_t.plot(steps, mean, color=INK, linewidth=1.2, label="100-step mean")
settled = [l for s, l in train if s >= min(300, steps[-1] // 5)] or list(loss)
ax_t.set_ylim(min(loss) - 0.15, max(settled) + 0.3)
ax_t.annotate(f"{mean[-1]:.3f}", (steps[-1], mean[-1]), xytext=(4, 0), textcoords="offset points",
color=INK_2, fontsize=8, va="center")
if steps[0] <= 10:
ax_t.text(0.01, 0.02, f"step {steps[0]} loss {loss[0]:.2f} (off scale)", transform=ax_t.transAxes,
ha="left", va="bottom", color=MUTED, fontsize=7)
ax_t.legend(loc="upper right", bbox_to_anchor=(1, 0.92), frameon=False, fontsize=8, labelcolor=INK_2)
if evals:
ends = []
for domain, color in DOMAIN_COLORS.items():
pts = [(s, r[domain]["loss"]) for s, r in evals if domain in r]
if pts:
xs, ys = zip(*pts)
ax_v.plot(xs, ys, color=color, linewidth=1.2, marker="o", markersize=4.5,
markeredgecolor=SURFACE, markeredgewidth=1, label=domain)
ends.append([ys[-1], domain, xs[-1], ys[-1]])
lo, hi = ax_v.get_ylim()
gap = (hi - lo) * 0.06 # dodge direct labels that would collide
ends.sort()
for i in range(1, len(ends)):
ends[i][0] = max(ends[i][0], ends[i - 1][0] + gap)
for label_y, domain, x, last in ends:
ax_v.annotate(f"{domain} {last:.3f}", (x, label_y), xytext=(6, 0), textcoords="offset points",
color=INK_2, fontsize=8, va="center")
ax_v.legend(loc="upper center", frameon=False, fontsize=8, labelcolor=INK_2)
else:
ax_v.text(0.5, 0.5, "no evaluation yet", transform=ax_v.transAxes, ha="center", color=MUTED)
for ax in (ax_t, ax_v):
ax.set_xlim(0, max(x_max * 1.08, 10))
_phase_lines(ax, x_max)
fig.suptitle(f"GPT-U-20M · step {x_max:,} / {total_steps:,}", x=0.01, ha="left", color=INK_2, fontsize=9)
fig.tight_layout()
out_path.parent.mkdir(parents=True, exist_ok=True)
tmp = out_path.with_suffix(".tmp.png")
fig.savefig(tmp, facecolor=SURFACE)
plt.close(fig)
tmp.replace(out_path)
|