File size: 4,838 Bytes
e0cce9b | 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 | #!/usr/bin/env python3
"""The scale sweep: retrieval vs YaRN scale, and the two-factor grid.
Every number read from the summary.csv files this session produced."""
import csv, os, glob
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
S={"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"}
def style(ax):
ax.set_facecolor(S["surface"])
for s in ax.spines.values(): s.set_color(S["grid"])
ax.grid(True,color=S["grid"],lw=.7,alpha=.9); ax.tick_params(colors=S["ink2"],labelsize=9)
def read(path,model):
if not os.path.exists(path): return None
r=list(csv.reader(open(path)))
if len(r)<2: return None
hdr=r[0]
for row in r[1:]:
if row and row[0]=="niah_single" and model in hdr:
try: return float(row[hdr.index(model)])
except Exception: return None
return None
# (scale, prompt) -> dir
SRC={(4,1024):"eshort_cache16k_len1024",(8,1024):"eshort_cache32k_len1024",
(4,4096):"eshort_cache16k_len4096",(8,4096):"eshort_cache32k_len4096",
(16,4096):"e64k_len4096",(32,4096):"e128k_len4096",
(8,8192):"e32k_len8192",(8,16360):"e32k_len16360",(8,32744):"e32k_len32744",
(16,16360):"e64k_len16360",(16,32744):"e64k_len32744",(16,65512):"e64k_len65512",
(32,32744):"e128k_len32744",(32,65512):"e128k_len65512"}
PRIOR={(4,8192):84.0,(4,16360):40.0} # measured on the prior box, same haystack
base={}
for (sc,pl),d in SRC.items():
v=read(f"{d}/summary.csv","base680")
if v is not None: base[(sc,pl)]=v
base.update(PRIOR)
fig,(axA,axB,axC)=plt.subplots(1,3,figsize=(19.5,5.6),facecolor=S["surface"])
PL=[1024,4096,8192,16360,32744,65512]
cols=plt.cm.viridis([i/(len(PL)-1) for i in range(len(PL))])
for pl,c in zip(PL,cols):
pts=sorted((sc,v) for (sc,p),v in base.items() if p==pl)
if len(pts)<2: continue
axA.plot([p[0] for p in pts],[p[1] for p in pts],"o-",color=c,lw=2,ms=6,
label=f"{pl:,} tok")
style(axA); axA.set_xscale("log",base=2); axA.set_xticks([4,8,16,32])
axA.set_xticklabels(["x4","x8","x16","x32"]); axA.set_ylim(0,100)
axA.set_xlabel("YARN_SCALE (ORIG_MAXPOS held at 4096)",color=S["ink2"],fontsize=10)
axA.set_ylabel("niah_single accuracy (%)",color=S["ink2"],fontsize=10)
axA.set_title("Same untrained weights. Only the RoPE scale changes.",
color=S["ink"],fontsize=12,loc="left",pad=10)
axA.legend(frameon=False,fontsize=8.5,labelcolor=S["ink2"],loc="lower right")
for pl,c in zip(PL,cols):
pts=sorted((pl/sc,v) for (sc,p),v in base.items() if p==pl)
if not pts: continue
axB.plot([p[0] for p in pts],[p[1] for p in pts],"o-",color=c,lw=2,ms=6,label=f"{pl:,} tok")
style(axB); axB.set_xscale("log",base=2); axB.set_ylim(0,100)
axB.set_xlabel("virtual position of last token (real / scale)",color=S["ink2"],fontsize=10)
axB.set_ylabel("niah_single accuracy (%)",color=S["ink2"],fontsize=10)
axB.axvspan(2000,4200,color="#f3eceb",zorder=0)
axB.text(2900,6,"pretrain p90\n~2,721 tok",ha="center",fontsize=7.6,color=S["ink2"],style="italic")
axB.set_title("Collapses onto virtual position; spread = distractor load",
color=S["ink"],fontsize=12,loc="left",pad=10)
axB.legend(frameon=False,fontsize=8.5,labelcolor=S["ink2"],loc="lower left")
fig.text(.005,-.02,"680M base680, ZERO long-context training. niah_single N=100, greedy, "
"identical haystack (199,754,694 B) across every point. Wilson 95% CI is +/-9-13pp at N=100.",
fontsize=7.8,color=S["ink2"])
# ---- panel C: the 1.06B control -- config helps, but training still wins at 64k
PLC=[16360,32744,65512]
own=[70,50,31]; corr=[78,73,65]; cliff_own=[54,31,18]; cliff_corr=[66,70,65]
m680=[92,79,46]
x=range(len(PLC))
axC.plot(x,own,"o--",color="#b0b0aa",lw=1.6,ms=6,label="1.06B, its own config")
axC.plot(x,corr,"o-",color="#b3243b",lw=2.4,ms=7,label="1.06B, corrected config")
axC.plot(x,cliff_corr,"s--",color="#eb6834",lw=1.6,ms=6,label='1.06B "collapsed", corrected')
axC.plot(x,m680,"^-",color="#12508f",lw=2.4,ms=8,label="680M base, no long-ctx training")
style(axC); axC.set_xticks(list(x)); axC.set_xticklabels(["16k","32k","64k"])
axC.set_ylim(0,100); axC.set_xlabel("prompt length",color=S["ink2"],fontsize=10)
axC.set_ylabel("niah_single accuracy (%)",color=S["ink2"],fontsize=10)
axC.set_title("Control: training wins where it trained (64k)",color=S["ink"],fontsize=12,loc="left",pad=10)
axC.annotate("1.06B trained to 64k\nbeats the 680M here\nby 19pp",xy=(2,65),xytext=(1.05,30),
fontsize=8.2,color=S["ink2"],arrowprops=dict(arrowstyle="->",color=S["ink2"],lw=1))
axC.legend(frameon=False,fontsize=8,labelcolor=S["ink2"],loc="upper right")
fig.tight_layout(); fig.savefig("plots_p1/scale_sweep.png",dpi=155,bbox_inches="tight")
print(" wrote plots_p1/scale_sweep.png")
print(f" points plotted: {len(base)}")
|