"""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}')