Upload scripts/mk_lineage_680m.py with huggingface_hub

scripts/mk_lineage_680m.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python3
+"""680M phase-1 lineage: one parent (varybatch), four forks at 4k -> 16k.
+
+All four arms read the same mix (b3 50 / b2 30 / b1 20, b3=320011377 tokens)
+and all four start from loss 2.731203, so the curves are directly comparable.
+Loss rows are parsed from each arm's own training log; nothing is transcribed.
+"""
+import re, sys, os
+import matplotlib; matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+S = {"surface": "#fcfcfb", "ink": "#0b0b0b", "ink2": "#52514e", "grid": "#e4e4e1"}
+D = os.path.dirname(os.path.abspath(__file__))
+
+ARMS = [
+    ("p1final  Muon  WSD peak 1.8e-4", f"{D}/lineage_src/p1_train.log",          "#eb6834", 1.7),
+    ("const    Muon  8e-5 flat",       f"{D}/lineage_src/abl_p1_const_train.log", "#12508f", 1.7),
+    ("cosine   Muon  8e-5 -> 8e-6",    f"{D}/lineage_src/abl_p1_cosine_train.log","#4a8fd4", 1.7),
+    ("const    AdamW 8e-5 flat",       f"{D}/p1_const_adamw_logs/train.log",      "#b3243b", 2.2),
+]
+
+def rows_of(p):
+    try: t = open(p, errors="ignore").read()
+    except Exception: return []
+    return [(int(a), float(b), float(c))
+            for a, b, c in re.findall(r"^step +(\d+) \| loss: ([0-9.]+) \| lr ([0-9.e+-]+)", t, re.M)]
+
+def roll(rows, k=50):
+    d = {s: l for s, l, _ in rows}; xs = []; ys = []
+    for s, _, _ in rows:
+        w = [d[j] for j in range(s - k + 1, s + 1) if j in d]
+        if len(w) == k: xs.append(s); ys.append(sum(w) / k)
+    return xs, ys
+
+fig, (axA, axB) = plt.subplots(2, 1, figsize=(11.5, 8.4), facecolor=S["surface"],
+                               gridspec_kw={"height_ratios": [1.45, 1], "hspace": 0.22}, sharex=True)
+parent_loss = None; cut = None; summary = []
+for name, path, col, lw in ARMS:
+    rw = rows_of(path)
+    if not rw:
+        print(f"  MISSING: {path}"); continue
+    parent_loss = rw[0][1] if parent_loss is None else parent_loss
+    axA.plot([s for s, _, _ in rw], [l for _, l, _ in rw], color=col, lw=.4, alpha=.16, zorder=2)
+    rx, ry = roll(rw)
+    axA.plot(rx, ry, color=col, lw=lw, zorder=5, label=f"{name}   (50-step mean)")
+    axB.plot([s for s, _, _ in rw], [r[2] for r in rw], color=col, lw=lw, zorder=4, label=name)
+    if "AdamW" in name:
+        cut = rw[-1][0]
+        axA.scatter([rx[-1]], [ry[-1]], s=46, color=col, zorder=7, edgecolor=S["surface"], lw=1.2)
+    summary.append((name, rw[0][0], rw[-1][0], rw[-1][1], ry[-1] if ry else float("nan")))
+
+axA.axhline(parent_loss, color=S["ink2"], lw=1.0, ls=":", alpha=.8, zorder=3)
+axA.annotate(f"parent varybatch   loss {parent_loss:.4f}   (all four arms start here)",
+             xy=(6, parent_loss), xytext=(6, parent_loss + 0.013), color=S["ink2"], fontsize=8.6)
+if cut:
+    for ax in (axA, axB):
+        ax.axvline(cut, color="#b3243b", lw=1.0, ls="--", alpha=.55, zorder=4)
+    axA.annotate(f"AdamW arm cut at {cut}", xy=(cut, parent_loss + 0.055),
+                 xytext=(cut + 18, parent_loss + 0.055), color="#b3243b", fontsize=8.6)
+
+for ax in (axA, axB):
+    ax.set_facecolor(S["surface"]); ax.grid(True, color=S["grid"], lw=.8, zorder=0); ax.set_axisbelow(True)
+    for k, sp in ax.spines.items():
+        sp.set_visible(k in ("left", "bottom")); sp.set_color(S["grid"])
+    ax.tick_params(colors=S["ink2"], labelsize=9)
+axA.set_ylabel("training loss (nats)", color=S["ink2"], fontsize=10)
+axA.set_title("680M phase 1: 4k -> 16k extension, one parent, four arms",
+              color=S["ink"], fontsize=12.5, loc="left", pad=11)
+axA.legend(frameon=False, fontsize=9, labelcolor=S["ink2"], loc="upper right")
+axB.set_yscale("log"); axB.set_ylabel("learning rate", color=S["ink2"], fontsize=10)
+axB.set_xlabel("step  (window 16,384 x grad-accum 32)", color=S["ink2"], fontsize=10)
+axB.legend(frameon=False, fontsize=8.5, labelcolor=S["ink2"], loc="lower left", ncol=2)
+
+out = f"{D}/plots_abl/lineage_680m_phase1.png"
+os.makedirs(os.path.dirname(out), exist_ok=True)
+fig.savefig(out, dpi=170, bbox_inches="tight", facecolor=S["surface"])
+print(f"wrote {out}")
+print(f"{'arm':34s} {'steps':>12s} {'final':>9s} {'50-step mean':>13s}")
+for n, a, b, f, m in summary:
+    print(f"{n:34s} {a:5d}..{b:<5d} {f:9.4f} {m:13.4f}")