summaryrefslogtreecommitdiff
path: root/ep_run/fig_primer_pack2.py
blob: 74850c8f46d8e41a16864f4e5e5e089a55aa9aae (plain)
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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
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/figs/'

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([104.8], [3.7117], marker='*', ms=14, color=GREEN, zorder=5)
axv.text(102, 3.52, 'best 3.7117 @105k — never improved again:\n90k healthy steps of plateau before the blow', fontsize=8.5, color=GREEN, ha='left')
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 (SEALED): rode 3e-3 for all 234k steps — zero guard events, final gap 0.043', fontsize=8.6, color=ORAN, ha='left')
ax2.text(186.5, 1.62e-3, '✓', fontsize=12, 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')