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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
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}')
|
