Upload scripts/mk_lineage.py with huggingface_hub

0da00ec verified 2 months ago

4.9 kB

#!/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)")

	#!/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)")