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

1 files changed, 148 insertions, 0 deletions

diff --git a/ep_run/fig_primer_block.py b/ep_run/fig_primer_block.py
new file mode 100644
index 0000000..ba3b5d0
--- /dev/null
+++ b/ep_run/fig_primer_block.py
@@ -0,0 +1,148 @@
+"""Primer figure 2: one OLMo2 block at op granularity (companion to fig_primer_arch).
+Color semantics: weight-matmuls (crossbar-mappable) vs activation-activation matmuls
+vs norms vs elementwise. Vector PDF + PNG to ../assets/."""
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+plt.rcParams['svg.fonttype'] = 'none'   # keep text editable in PPT
+from matplotlib.patches import FancyBboxPatch, Circle
+
+INK   = '#3a3a3a'
+GRAY  = '#8a8a8a'
+BOXF  = '#f4f4f4'
+BOXE  = '#9a9a9a'
+BLUE  = '#2c6fbb'
+BLUEF = '#eaf1fa'
+PURP  = '#7b5aa6'
+PURPF = '#f1ebf8'
+CREAM = '#faf6ee'
+CREME = '#c9b895'
+RAIL  = '#555555'
+
+fig = plt.figure(figsize=(10.0, 14.2))
+ax = fig.add_axes([0, 0, 1, 1]); ax.set_xlim(0, 100); ax.set_ylim(0, 142); ax.axis('off')
+
+def box(xc, yc, w, s, fc=BOXF, ec=BOXE, size=8.6, lw=1.1, h=3.2, tc=INK):
+    b = FancyBboxPatch((xc - w / 2, yc - h / 2), w, h,
+                       boxstyle='round,pad=0.25,rounding_size=0.55', fc=fc, ec=ec, lw=lw)
+    ax.add_patch(b)
+    ax.text(xc, yc, s, fontsize=size, color=tc, ha='center', va='center')
+
+def txt(x, y, s, size=8.5, color=INK, ha='center', va='center', w='normal', style='normal', rot=0):
+    ax.text(x, y, s, fontsize=size, color=color, ha=ha, va=va,
+            fontweight=w, fontstyle=style, rotation=rot)
+
+def arrow(x0, y0, x1, y1, color=INK, lw=1.15, ms=10):
+    ax.annotate('', xy=(x1, y1), xytext=(x0, y0),
+                arrowprops=dict(arrowstyle='-|>', color=color, lw=lw,
+                                mutation_scale=ms, shrinkA=0.4, shrinkB=0.4))
+
+def rail(pts, color=RAIL, lw=1.4):
+    xs, ys = zip(*pts)
+    ax.plot(xs, ys, color=color, lw=lw, solid_capstyle='round', zorder=1)
+
+# ---------------- title ----------------
+txt(50, 138.8, 'One transformer block, op by op', 12.5, w='bold')
+txt(50, 135.9, 'OLMo2 ordering — RMSNorm after each sublayer, inside the residual.'
+               '     $C{=}512$   $H{=}8$   $h_d{=}64$   $T{=}256$   $h_{ff}{=}1408$', 8.8, GRAY)
+
+# ---------------- attention sublayer ----------------
+txt(3.2, 100, 'attention sublayer', 9.5, GRAY, style='italic', rot=90)
+
+txt(40, 131.6, r'input $z$   $(B,T,C)$', 9.5, w='bold')
+rail([(40, 130.4), (88, 130.4), (88, 71.8)]); arrow(88, 72.4, 42.0, 71.8, color=RAIL, lw=1.4)
+txt(90.2, 101, 'residual rail', 8, GRAY, rot=90)
+
+arrow(40, 130.2, 40, 128.6)
+box(40, 126.9, 30, r'matmul  $W_{qkv}: C \to 3C$', fc=BLUEF, ec=BLUE)
+txt(56.5, 126.9, r'$(B,T,3C)$ — split into $q,k,v$', 7.6, GRAY, ha='left')
+
+qx, kx, vx = 16, 40, 64
+for x, s in [(qx, '$q$'), (kx, '$k$'), (vx, '$v$')]:
+    txt(x, 123.3, s, 9.5, w='bold')
+arrow(37, 125.3, qx + 1, 122.5); arrow(40, 125.3, kx, 122.5); arrow(43, 125.3, vx - 1, 122.5)
+
+box(qx, 120.9, 17.5, r'RMSNorm$_q$', fc=CREAM, ec=CREME)
+box(kx, 120.9, 17.5, r'RMSNorm$_k$', fc=CREAM, ec=CREME)
+box(qx, 116.2, 17.5, 'split heads ($H{=}8$)')
+box(kx, 116.2, 17.5, 'split heads ($H{=}8$)')
+box(vx, 116.2, 17.5, 'split heads ($H{=}8$)')
+box(qx, 111.5, 17.5, 'RoPE')
+box(kx, 111.5, 17.5, 'RoPE')
+txt(50.2, 111.5, r'$\theta = 5{\cdot}10^5$', 7.6, GRAY, ha='left')
+for x in (qx, kx):
+    arrow(x, 119.3, x, 117.8); arrow(x, 114.6, x, 113.1)
+arrow(vx, 121.9, vx, 117.8)
+
+arrow(qx, 109.9, 30, 107.2); arrow(kx, 109.9, 40, 107.2); arrow(vx, 114.6, 50, 107.2)
+box(40, 105.5, 32, r'matmul  $S = q\,k^{\top} \cdot 1/\sqrt{h_d}$', fc=PURPF, ec=PURP)
+txt(57.5, 105.5, r'$(B,H,T,T)$', 7.6, GRAY, ha='left')
+arrow(40, 103.9, 40, 102.4)
+box(40, 100.8, 32, 'causal mask')
+arrow(40, 99.2, 40, 97.7)
+box(40, 96.1, 32, 'softmax over keys')
+arrow(40, 94.5, 40, 93.0)
+box(40, 91.4, 32, r'matmul  $y = A\,v$', fc=PURPF, ec=PURP)
+txt(57.5, 91.4, r'$(B,H,T,h_d)$', 7.6, GRAY, ha='left')
+arrow(40, 89.8, 40, 88.3)
+box(40, 86.7, 32, 'merge heads')
+txt(57.5, 86.7, r'$(B,T,C)$', 7.6, GRAY, ha='left')
+arrow(40, 85.1, 40, 83.6)
+box(40, 82.0, 32, r'matmul  $W_{proj}: C \to C$', fc=BLUEF, ec=BLUE)
+arrow(40, 80.4, 40, 78.9)
+box(40, 77.3, 32, r'RMSNorm$_{attn}$', fc=CREAM, ec=CREME)
+txt(57.5, 77.3, '◄  OLMo2: norm on the sublayer OUTPUT,\n     not on its input', 7.6, INK, ha='left')
+arrow(40, 75.7, 40, 73.8)
+c1 = Circle((40, 71.8), 1.9, fc='white', ec=INK, lw=1.3, zorder=3); ax.add_patch(c1)
+txt(40, 71.8, '+', 11, w='bold')
+
+arrow(40, 69.9, 40, 67.3)
+txt(42.2, 68.7, r"$z'$", 10, w='bold', ha='left')
+
+# ---------------- SwiGLU sublayer ----------------
+txt(3.2, 50, 'SwiGLU sublayer', 9.5, GRAY, style='italic', rot=90)
+
+rail([(40, 67.0), (88, 67.0), (88, 37.2)]); arrow(88, 37.8, 42.0, 37.2, color=RAIL, lw=1.4)
+
+arrow(38, 66.4, 27, 64.0); arrow(42, 66.4, 53, 64.0)
+box(26, 62.4, 21, r'matmul  $W_1: C \to h_{ff}$', fc=BLUEF, ec=BLUE)
+box(54, 62.4, 21, r'matmul  $W_3: C \to h_{ff}$', fc=BLUEF, ec=BLUE)
+txt(65.7, 62.4, r'$h_{ff} = 1408$', 7.6, GRAY, ha='left')
+arrow(26, 60.8, 26, 59.3)
+box(26, 57.7, 12, 'SiLU')
+arrow(26, 56.1, 33.5, 54.1); arrow(54, 60.8, 46.5, 54.1)
+box(40, 52.4, 23, '⊙  elementwise gate')
+arrow(40, 50.8, 40, 49.3)
+box(40, 47.1, 25, r'matmul  $W_2: h_{ff} \to C$', fc=BLUEF, ec=BLUE)
+arrow(40, 45.5, 40, 44.0)
+box(40, 42.4, 25, r'RMSNorm$_{ff}$', fc=CREAM, ec=CREME)
+txt(54, 42.4, '◄  norm after the sublayer, again', 7.6, INK, ha='left')
+arrow(40, 40.8, 40, 39.2)
+c2 = Circle((40, 37.2), 1.9, fc='white', ec=INK, lw=1.3, zorder=3); ax.add_patch(c2)
+txt(40, 37.2, '+', 11, w='bold')
+arrow(40, 35.3, 40, 32.8)
+txt(40, 31.3, r'output $z^{\prime\prime}$  —  one full $f_l$ done', 9.5, w='bold')
+
+# ---------------- legend + notes ----------------
+def chip(x, y, fc, ec):
+    b = FancyBboxPatch((x, y - 0.9), 3.4, 1.8, boxstyle='round,pad=0.15,rounding_size=0.35',
+                       fc=fc, ec=ec, lw=1.1)
+    ax.add_patch(b)
+
+chip(5, 25.5, BLUEF, BLUE)
+txt(9.5, 25.5, 'weight × activation matmul — 7 logical ($W_{qkv}$ fuses $W_q,W_k,W_v$); crossbar-mappable on analog hardware', 8.2, ha='left')
+chip(5, 22.1, PURPF, PURP)
+txt(9.5, 22.1, r'activation × activation matmul — 2 ($q\,k^{\top}$ and $A\,v$); computed on the fly: the non-crossbar part of attention', 8.2, ha='left')
+chip(5, 18.7, CREAM, CREME)
+txt(9.5, 18.7, 'RMSNorm (learned gain)', 8.2, ha='left')
+chip(38, 18.7, BOXF, BOXE)
+txt(42.5, 18.7, 'elementwise / reshape', 8.2, ha='left')
+
+txt(5, 14.6, 'No biases anywhere in the network.  The $1/\\sqrt{h_d}$ scale is folded into $S$; mask and softmax run along the key axis.', 8.2, ha='left')
+txt(5, 11.9, r'$h_{ff} = 1408 \approx 8C/3$, rounded up to a multiple of 64.  QK-norm acts on the full width $C$, before the head split.', 8.2, ha='left')
+txt(5, 9.2, r'The boxed pipeline, input $z$ to output $z^{\prime\prime}$, is exactly $f_l$ on the companion page; the state $z_l$ lives on the residual rail.', 8.2, ha='left')
+
+fig.savefig('/home/yurenh2/ept/assets/fig_primer_block.pdf')
+fig.savefig('/home/yurenh2/ept/assets/fig_primer_block.png', dpi=185)
+fig.savefig('/home/yurenh2/ept/assets/fig_primer_block.svg')
+print('saved fig_primer_block.{pdf,png,svg}')