Upload scripts/mk_lineage.py with huggingface_hub

scripts/mk_lineage.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+"""Standalone lineage loss plot: every training stage this 1.06B model has been
+through, forking at ext64k step 695 into the two 64k runs.
+
+  100% b5  -> collapsed (cliff at step 1199, ended 45/17/2/1)
+  MIX 50/30/20 -> this run
+
+Loss curves come from each stage's own training log; nothing is transcribed.
+"""
+import json, math, os, re, statistics as st
+import matplotlib; matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+S={"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"}
+PURE="#b3243b"; MIX="#12508f"; PARENT="#eb6834"
+R="mix64k_ref"
+klab=lambda v:(f"{round(v/1024)}k" if v>=1024 else str(int(v)))
+
+def rows_of(p):
+    try: t=open(p,errors="ignore").read()
+    except Exception: return []
+    return [(int(a),float(b)) for a,b in re.findall(r"^step +(\d+) \| loss: ([0-9.]+)",t,re.M)]
+
+def roll(rows,k=50):
+    d=dict(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
+
+stages=json.load(open(f"{R}/lineage/stages.json"))
+ORDER=[("base pretrain",4096,"#8e8e88"),("l2 16k bucket",16384,"#9dc3ea"),
+       ("ext32k",32768,"#4a8fd4"),("b4 full epoch",32768,"#12508f"),("ext64k",65536,PARENT)]
+pr=rows_of(f"{R}/cont64k_purec5.log"); mr=rows_of("mix64k_logs/train.log")
+
+fig,(axA,axB)=plt.subplots(2,1,figsize=(11.5,9.2),facecolor=S["surface"],
+                           gridspec_kw={"height_ratios":[1.2,1],"hspace":0.26},sharex=True)
+off=0; spans=[]
+for name,ctx,col in ORDER:
+    stg=stages.get(name)
+    if not stg: continue
+    rw=stg["rows"]; b=rw[0][0]
+    xs=[off+(r[0]-b) for r in rw]; ys=[r[1] for r in rw]
+    for ax in (axA,axB):
+        ax.plot(xs,ys,color=col,lw=.8,alpha=.85,zorder=3,
+                label=(f"{name}  ({klab(ctx)})" if ax is axA else None))
+    spans.append((off,xs[-1],name,col)); off=xs[-1]
+fork=off
+for rw,col,lab,lw in ((pr,PURE,"100% b5  ->  COLLAPSED",1.5),(mr,MIX,"MIX 50/30/20  ->  this run",1.9)):
+    if not rw: continue
+    xs=[fork+s for s,_ in rw]; ys=[l for _,l in rw]
+    for ax in (axA,axB): ax.plot(xs,ys,color=col,lw=.45,alpha=.22,zorder=2)
+    rx,ry=roll(rw)
+    for ax in (axA,axB):
+        ax.plot([fork+x for x in rx],ry,color=col,lw=lw,zorder=5,
+                label=(lab+"  (50-step mean)" if ax is axA else None))
+    spans.append((fork,fork+rw[-1][0],lab.split("  ->")[0],col))
+
+ends=[stages[n]["rows"][-1][1] for n,_,_ in ORDER if n in stages]
+lo,hi=min(ends),max(ends)
+for ax in (axA,axB):
+    ax.axvline(fork,color=S["ink2"],lw=1.1,ls="--",alpha=.75,zorder=4)
+    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_yscale("log"); axA.set_ylim(0.9,13)
+axA.set_yticks([1,2,3,5,10]); axA.set_yticklabels(["1","2","3","5","10"])
+axA.set_ylabel("loss (nats, log)",color=S["ink2"],fontsize=10)
+# axA.set_title("Every stage this 1.06B model has been through -- and where the two 64k runs diverge",
+            #   color=S["ink"],fontsize=12.5,loc="left",pad=12)
+# axA.annotate("both 64k runs fork here\n(ext64k step 695, 67/65/52/29)",xy=(fork,8.2),
+#              xytext=(fork*0.50,8.6),fontsize=8.8,color=S["ink2"],
+            #  arrowprops=dict(arrowstyle="->",color=S["ink2"],lw=.9))
+axA.legend(frameon=False,fontsize=8.2,labelcolor=S["ink2"],loc="lower left",ncol=2)
+axB.axhspan(lo,hi,color="#7a9a3a",alpha=.13,zorder=1)
+# axB.annotate(f"every stage before the fork ends in {lo:.2f}-{hi:.2f}\nthe MIX stays in the band; 100% b5 leaves it and never returns",
+#              xy=(fork+1300,1.30),xytext=(off*0.06,1.18),fontsize=8.8,color="#5d7a22",
+#              arrowprops=dict(arrowstyle="->",color="#7a9a3a",lw=1.0,alpha=.85))
+axB.set_ylim(1.0,3.5); axB.set_ylabel("loss (nats)",color=S["ink2"],fontsize=10)
+axB.set_xlabel("cumulative training step across the lineage",color=S["ink2"],fontsize=10)
+for s,e,name,col in spans:
+    axB.text((s+e)/2,3.38,name,ha="center",fontsize=7.4,color=col)
+fig.savefig("mix64k_lineage.png",dpi=140,facecolor=S["surface"],bbox_inches="tight")
+print(f"  wrote mix64k_lineage.png ({os.path.getsize('mix64k_lineage.png')//1024} KB)")
+print(f"\n  {'stage':<26}{'window':>8}{'steps':>7}{'first':>9}{'last':>8}")
+for n,c,_ in ORDER:
+    if n in stages:
+        rw=stages[n]["rows"]; print(f"  {n:<26}{klab(c):>8}{len(rw):>7}{rw[0][1]:>9.3f}{rw[-1][1]:>8.3f}")
+for rw,n in ((pr,"100% b5 (collapsed)"),(mr,"MIX (this run)")):
+    if rw: print(f"  {n:<26}{'64k':>8}{len(rw):>7}{rw[0][1]:>9.3f}{rw[-1][1]:>8.3f}")
+print(f"\n  pre-fork stages end in {lo:.3f}-{hi:.3f}")
+if mr:
+    d=dict(mr); mx=max(d); w=[d[k] for k in range(mx-49,mx+1) if k in d]
+    print(f"  MIX rolling-50 at step {mx}: {st.mean(w):.3f}  (pure-b5 was 1.80 by here)")