summaryrefslogtreecommitdiff
path: root/ep_run/fig_primer_pack2.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-07-15 10:29:06 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-07-15 10:29:06 -0500
commit82da86dbd49a093a24e0827331667e8e9e6217df (patch)
tree96c40877ab6888d66c6e7992a8d6bfebb0453ba2 /ep_run/fig_primer_pack2.py
parentbf83f39828fb79428fe65fe0fef49768351afc05 (diff)
Alexi primer deck: 10 slides (arch op-by-op, EP phases, schedule/parallelism, 72M blowup from real logs, loop-gain + beta-window hypotheses, diagnostics table, in-flight outcomes, $50k ladder) + figure scripts
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run/fig_primer_pack2.py')
-rw-r--r--ep_run/fig_primer_pack2.py201
1 files changed, 201 insertions, 0 deletions
diff --git a/ep_run/fig_primer_pack2.py b/ep_run/fig_primer_pack2.py
new file mode 100644
index 0000000..ff28476
--- /dev/null
+++ b/ep_run/fig_primer_pack2.py
@@ -0,0 +1,201 @@
+"""Primer deck figures: schedule, 72M blowup (real logs), loop-gain, beta window.
+PNGs to ../assets/ for the pptx build."""
+import re
+import matplotlib
+matplotlib.use('Agg')
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.patches import FancyBboxPatch
+
+INK, GRAY = '#3a3a3a', '#8a8a8a'
+BOXF, BOXE = '#f4f4f4', '#9a9a9a'
+BLUE, BLUEF = '#2c6fbb', '#eaf1fa'
+ORAN, ORANF = '#d95f02', '#fdeee2'
+RED = '#b03a2e'
+CREAM, CREME = '#faf6ee', '#c9b895'
+GREEN = '#2e7d32'
+A = '/home/yurenh2/ept/assets/'
+
+def newfig(w, h, xmax=100, ymax=62):
+ fig = plt.figure(figsize=(w, h))
+ ax = fig.add_axes([0, 0, 1, 1]); ax.set_xlim(0, xmax); ax.set_ylim(0, ymax); ax.axis('off')
+ return fig, ax
+
+def box(ax, xc, yc, w, s, fc=BOXF, ec=BOXE, size=10, lw=1.2, h=4.6, tc=INK, wt='normal'):
+ ax.add_patch(FancyBboxPatch((xc - w / 2, yc - h / 2), w, h,
+ boxstyle='round,pad=0.3,rounding_size=0.7', fc=fc, ec=ec, lw=lw))
+ ax.text(xc, yc, s, fontsize=size, color=tc, ha='center', va='center', fontweight=wt)
+
+def txt(ax, x, y, s, size=9.5, color=INK, ha='center', va='center', w='normal', style='normal'):
+ ax.text(x, y, s, fontsize=size, color=color, ha=ha, va=va, fontweight=w, fontstyle=style)
+
+def arrow(ax, x0, y0, x1, y1, color=INK, lw=1.4, ms=12, cs=None):
+ ax.annotate('', xy=(x1, y1), xytext=(x0, y0),
+ arrowprops=dict(arrowstyle='-|>', color=color, lw=lw, mutation_scale=ms,
+ shrinkA=0.5, shrinkB=0.5, connectionstyle=cs))
+
+# ================= fig S3: the update schedule =================
+fig, ax = newfig(11.8, 6.0)
+dx, zx = 15, 41
+txt(ax, dx, 58.5, 'mismatches derive TOP → BOTTOM', 10, ORAN, w='bold')
+txt(ax, zx, 58.5, 'states rebuild BOTTOM → TOP', 10, BLUE, w='bold')
+dl = [r'$d_3 = -\beta\,\nabla\mathrm{CE}$', r'$d_2 = J_3^{\top} d_3$', r'$d_1 = J_2^{\top} d_2$']
+dn = ['needs: loss only', 'needs: block 3 only', 'needs: block 2 only']
+for i, (s, n) in enumerate(zip(dl, dn)):
+ y = 51 - i * 9
+ box(ax, dx, y, 20, s, fc=ORANF, ec=ORAN)
+ txt(ax, dx, y - 3.4, n, 7.6, GRAY)
+ if i < 2: arrow(ax, dx, y - 2.4 - 2.2, dx, y - 9 + 2.5, color=ORAN)
+zl = [r'$z_1 = f_1(\mathrm{emb}) + d_1$', r'$z_2 = f_2(z_1) + d_2$', r'$z_3 = f_3(z_2) + d_3$']
+for i, s in enumerate(zl):
+ y = 33 + i * 9
+ box(ax, zx, y, 20, s, fc=BLUEF, ec=BLUE)
+ if i < 2: arrow(ax, zx, y + 2.4, zx, y + 9 - 2.5, color=BLUE)
+txt(ax, 28, 22.5, 'block $l$ touches NEIGHBORS only — $d$ from the block after it,\n'
+ '$z$ from the block before it. $\\Delta\\theta_l$ fires as soon as $d_l$ arrives.\n'
+ 'Repeat $K$ times ($K{=}3$ measured sufficient; $K{=}8$ identical).', 9)
+
+sx = 76
+txt(ax, sx, 58.5, 'the same settle, three schedules', 10.5, INK, w='bold')
+rows = [
+ ('sequential (our trainer today)', r'$O(K\!\cdot\!d)$', '3 × 12 = 36 block-ops', BOXF, BOXE),
+ ('pipelined wavefront (systolic; multi-GPU lever)', r'$O(K\!+\!d\!-\!1)$', '3 + 12 − 1 = 14 slots', BLUEF, BLUE),
+ ('continuous-time analog: all blocks relax AT ONCE', r'$\sim O(K)$', 'depth-serial factor gone;\n'
+ r'$K$ dissolves into settling time $\tau$', ORANF, ORAN),
+]
+for i, (name, comp, det, fc, ec) in enumerate(rows):
+ y = 50 - i * 12.5
+ box(ax, sx, y, 44, '', fc=fc, ec=ec, h=9.5)
+ txt(ax, sx - 20.5, y + 2.2, name, 9.3, ha='left', w='bold')
+ txt(ax, sx - 20.5, y - 1.8, det, 8.6, GRAY, ha='left')
+ txt(ax, sx + 15.5, y + 2.2, comp, 12)
+txt(ax, sx, 9.5, 'open question we track: does $\\tau$ itself grow with depth?\n'
+ '(one transit down the chain is unavoidable)', 8.4, GRAY)
+fig.savefig(A + 'fig_p3_schedule.png', dpi=170); plt.close(fig)
+
+# ================= fig S4: what happened at 72M (real logs) =================
+STEP_RE = re.compile(r'^step\s+(\d+)/\d+ \| train [\d.]+ val ([\d.]+) \(best ([\d.]+)\).*?beta=([\d.e+-]+)')
+def parse(fn):
+ out = {}
+ for line in open('/home/yurenh2/ept/ep_run/runs/' + fn, errors='replace'):
+ m = STEP_RE.match(line)
+ if m: out[int(m.group(1))] = (float(m.group(2)), float(m.group(4)))
+ return out
+
+base = parse('fw72m.log')
+resc = parse('fw72m_c.log'); resc.update(parse('fw72m_c2.log'))
+bs = np.array(sorted(base)); bval = np.array([base[s][0] for s in bs])
+rs = np.array(sorted(k for k in resc if k >= 195000)); rval = np.array([resc[s][0] for s in rs])
+print('base max val post-195k:', bval[bs > 195000].max(), '| resc pts:', len(rs))
+beta_steps = np.array(sorted([k for k in base if k < 195000] + list(rs)))
+beta_v = np.array([base[k][1] if k < 195000 else resc[k][1] for k in beta_steps])
+
+fig = plt.figure(figsize=(9.2, 6.0))
+axv = fig.add_axes([0.085, 0.12, 0.815, 0.78])
+axv.plot(bs / 1000, np.clip(bval, 0, 6.42), color=BLUE, lw=1.0, zorder=3)
+axv.plot(rs / 1000, np.clip(rval, 0, 6.42), color='#999999', lw=0.9, zorder=2)
+axv.text(233, 4.32, 'rescue resumes\n(new recipe from 195k / 215k)', fontsize=8, color='#777777', ha='right')
+axv.axhline(3.2884, color=GRAY, lw=1.0, ls='--')
+axv.text(2, 3.19, 'BP twin final 3.2884', fontsize=8.5, color=GRAY, va='top')
+axv.plot([185], [3.7117], marker='*', ms=14, color=GREEN, zorder=5)
+axv.text(180, 3.50, 'best 3.7117 @185k (never beaten again)', fontsize=8.5, color=GREEN, ha='right')
+axv.annotate('195k: relaxation diverges — val 3.9 → 6.0,\nguard storm (15k step-skips); the same\n$\\beta$ was safe for the previous 135k steps',
+ xy=(204, 5.95), xytext=(112, 5.62), fontsize=9.5, color=RED,
+ arrowprops=dict(arrowstyle='-|>', color=RED))
+axv.axvline(60, color=GRAY, lw=0.8, ls=':')
+axv.axvline(215, color=GRAY, lw=0.8, ls=':')
+axv.text(61, 6.30, 'floor raised\n3e-4 → 1e-3: fine', fontsize=8, color=GRAY, va='top')
+axv.text(146, 4.78, 'governor caps $\\beta$: survives,\nbut starved to 2e-5 — best frozen', fontsize=8.5, color=INK)
+axv.set_xlabel('step (×1000)', fontsize=10); axv.set_ylabel('val CE', fontsize=10, color=BLUE)
+axv.set_ylim(3.0, 6.5); axv.set_xlim(0, 238)
+axv.tick_params(labelsize=9)
+axb = axv.twinx()
+axb.plot(beta_steps / 1000, beta_v, color=ORAN, lw=1.3, alpha=0.9, drawstyle='steps-post')
+axb.set_yscale('log'); axb.set_ylim(8e-6, 6e-3)
+axb.set_ylabel(r'$\beta$ (log)', fontsize=10, color=ORAN)
+axb.tick_params(labelsize=8, colors=ORAN)
+for sp in ['top']: axv.spines[sp].set_visible(False); axb.spines[sp].set_visible(False)
+axv.set_title('fw72m — 72.11M params × 1.44B FineWeb tokens, fully BP-free (what actually happened)',
+ fontsize=11, color=INK, pad=10)
+fig.savefig(A + 'fig_p4_blowup.png', dpi=170); plt.close(fig)
+
+# ================= fig S5: the loop =================
+fig, ax = newfig(8.8, 6.2, xmax=100, ymax=70)
+bx = 26
+for i, name in enumerate(['block 1', 'block 2', 'block 3']):
+ box(ax, bx, 16 + i * 15, 22, name, fc=BOXF, ec=BOXE, h=8, size=11)
+arrow(ax, bx + 12, 20, bx + 12, 27, color=BLUE, lw=2.0); arrow(ax, bx + 12, 35, bx + 12, 42, color=BLUE, lw=2.0)
+arrow(ax, bx - 12, 42, bx - 12, 35, color=ORAN, lw=2.0); arrow(ax, bx - 12, 27, bx - 12, 20, color=ORAN, lw=2.0)
+ax.text(bx - 17.5, 31, 'errors flow DOWN $d_l = J_{l+1}^{\\top} d_{l+1}$', fontsize=9, color=ORAN,
+ ha='center', va='center', rotation=90)
+ax.text(bx + 17.5, 31, 'states rebuild UP $z_l = f_l(z_{l-1}) + d_l$', fontsize=9, color=BLUE,
+ ha='center', va='center', rotation=270)
+arrow(ax, bx + 9, 51.5, bx - 9, 51.5, color=RED, lw=2.0, cs='arc3,rad=0.45')
+txt(ax, bx, 60, 'top link: $-\\beta\\,\\nabla\\mathrm{CE}$, scale $\\sigma(W_{out})^2$', 9.5, RED)
+txt(ax, bx, 6, 'one sweep = down + up:\na closed loop through the stack', 9.5, GRAY)
+
+L = 50
+box(ax, 72.5, 56, 45, r'per-sweep gain $\rho \;\approx\; \beta\cdot\sigma(W_{out})^2\cdot\|J\mathrm{-chains}\|^2$',
+ fc=CREAM, ec=CREME, h=8, size=10.5)
+txt(ax, L, 46, r'$\rho < 1$: mismatch decays every sweep → settles', 10, GREEN, ha='left')
+txt(ax, L, 41, r'$\rho > 1$: every sweep AMPLIFIES the residual → diverges', 10, RED, ha='left')
+txt(ax, L, 34.8, 'and $\\rho$ RISES during training:', 10, INK, ha='left')
+txt(ax, L, 31.3, r'$\sigma(W_{out})$ and $\|J\|$ grow as the fit deepens', 10, INK, ha='left')
+txt(ax, L, 24, 'measured:', 9.5, INK, ha='left', w='bold')
+txt(ax, L, 20.3, '· damping refuted in 3 doses (spectrum monotone-positive)', 9, INK, ha='left')
+txt(ax, L, 16.8, '· lowering $\\beta$ restores convergence instantly', 9, INK, ha='left')
+txt(ax, L, 13.3, '· continuous-time flow is unconditionally stable →', 9, INK, ha='left')
+txt(ax, L + 1.5, 9.6, 'a DISCRETE-SOLVER disease; analog hardware\nhas no such ceiling', 9, ORAN, ha='left')
+fig.savefig(A + 'fig_p5_loop.png', dpi=170); plt.close(fig)
+
+# ================= fig S6: the beta window =================
+fig = plt.figure(figsize=(9.2, 6.0))
+ax2 = fig.add_axes([0.10, 0.13, 0.86, 0.76])
+x = np.linspace(0, 234, 500)
+cx = [0, 60, 120, 160, 195, 215, 222, 230, 234]
+cy = [1.2e-2, 1.0e-2, 6e-3, 2.5e-3, 9.5e-4, 4e-4, 2.6e-4, 4e-5, 2.2e-5]
+ceil = np.exp(np.interp(x, cx, np.log(cy)))
+fx_ = [0, 60, 120, 180, 234]; fy = [8e-5, 1.1e-4, 1.6e-4, 2.6e-4, 4.5e-4]
+floor = np.exp(np.interp(x, fx_, np.log(fy)))
+ax2.fill_between(x, ceil, 2e-2, color=RED, alpha=0.10)
+ax2.fill_between(x, 8e-6, floor, color='#666666', alpha=0.12)
+ax2.fill_between(x, floor, ceil, color=GREEN, alpha=0.07)
+ax2.plot(x, ceil, color=RED, lw=1.6)
+ax2.plot(x, floor, color='#555555', lw=1.6)
+ax2.text(6, 6.5e-3, r'wall-2: $\rho>1$ — relaxation diverges', fontsize=10, color=RED)
+ax2.text(6, 2.6e-5, 'wall-1: SNR < 1 — updates dissolve into noise', fontsize=10, color='#555555')
+ax2.text(96, 1.35e-3, r'$\beta_{max}(t) \propto \mathrm{margin}/(\sigma^2\|J\|^2)$', fontsize=9.5, color=RED, rotation=-14)
+ax2.text(120, 1.05e-4, r'$\beta_{min}(t) \propto \mathrm{noise}/|g|$', fontsize=9.5, color='#555555', rotation=6)
+ax2.text(150, 5.5e-4, 'operating window\n(narrows as fit deepens)', fontsize=9.5, color=GREEN, ha='center')
+# original schedule
+ax2.plot([0, 60], [3e-4, 3e-4], color=INK, lw=2.2)
+ax2.plot([60, 60], [3e-4, 1e-3], color=INK, lw=1.0, ls=':')
+ax2.plot([60, 195], [1e-3, 1e-3], color=INK, lw=2.2)
+ax2.plot([195], [1e-3], marker='x', ms=13, mew=3, color=RED, zorder6b=5) if False else ax2.plot([195], [1e-3], marker='x', ms=13, mew=3, color=RED, zorder=5)
+ax2.text(196.5, 1.25e-3, 'fixed floor crosses the\nfalling ceiling → blow', fontsize=8.5, color=RED)
+ax2.text(90, 7.6e-4, 'fw72m schedule (fixed floors)', fontsize=8.5, color=INK)
+# bcap rescue
+bs = x[(x >= 215) & (x <= 234)]
+ax2.plot(bs, np.exp(np.interp(bs, [215, 222, 230, 234], np.log([3e-4, 2.2e-4, 3.5e-5, 2e-5]))),
+ color=GRAY, lw=1.8, ls='-.')
+ax2.text(213, 4.8e-5, 'bcap rescue: survives, starved —\nthe window had CLOSED for the\nold recipe\'s margin', fontsize=8, color=GRAY, ha='right')
+# cent path
+xm = x[x <= 185]
+cent = np.where(xm < 20, 3e-4, 3e-3)
+cent = np.minimum(cent, np.exp(np.interp(xm, cx, np.log(cy))) * 0.85)
+ax2.plot(xm, cent, color=ORAN, lw=2.2, ls='--')
+ax2.text(24, 4.25e-3, 'fw72m_cent (in flight): centered + $\\beta$ rides ITS OWN ceiling via bcap 0.9 — where that ceiling sits is being measured now (48k: ahead)', fontsize=8.6, color=ORAN, ha='left')
+ax2.text(186.5, 1.62e-3, '→ ?', fontsize=10, color=ORAN, fontweight='bold')
+# measured ceiling points
+for (mx, my, lab) in [(195, 1e-3, '≤1e-3'), (222, 3e-4, '≤3e-4'), (228, 2.5e-5, '≈2e-5')]:
+ ax2.plot([mx], [my], marker='v', ms=7, color=RED, zorder=5)
+ax2.text(231, 2.3e-3, 'measured ceiling\ncrossings ▾', fontsize=8, color=RED, ha='right')
+ax2.set_yscale('log'); ax2.set_ylim(8e-6, 2e-2); ax2.set_xlim(0, 236)
+ax2.set_xlabel('training progress (×1000 steps)', fontsize=10)
+ax2.set_ylabel(r'$\beta$ (log scale)', fontsize=10)
+ax2.tick_params(labelsize=9)
+for sp in ['top', 'right']: ax2.spines[sp].set_visible(False)
+ax2.set_title(r'$\beta$ is squeezed from both sides — and the window narrows as the model fits deeper',
+ fontsize=11.5, color=INK, pad=10)
+fig.savefig(A + 'fig_p6_window.png', dpi=170); plt.close(fig)
+print('figs saved: p3_schedule, p4_blowup, p5_loop, p6_window')