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
|
"""Validation curves at 135M: EP with the readout fix, EP without it, and the backprop twins.
The point of the figure is that the first and third are the same curve, and the second stops
descending partway through. Run after fw135m_rlin completes; it reads the training logs directly.
"""
import re
import statistics as st
from pathlib import Path
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
RUNS = Path('/home/yurenh2/ept/ep_run/runs')
TOTAL = 440000
TOK_PER_STEP = 24 * 256 # effective batch x context
def curve(tag):
pts = []
for line in (RUNS / f'{tag}.log').read_bytes().decode(errors='replace').splitlines():
m = re.match(rf'step (\d+)/{TOTAL} \| train [\d.]+ val ([\d.]+)', line)
if m:
pts.append((int(m.group(1)), float(m.group(2))))
pts.sort()
return pts
def smooth(pts, win=25):
xs = [p[0] for p in pts]
ys = [p[1] for p in pts]
out = []
for i in range(len(ys)):
lo, hi = max(0, i - win // 2), min(len(ys), i + win // 2 + 1)
out.append(st.median(ys[lo:hi]))
return xs, out
def main():
# EP is drawn thick and underneath; the backprop seeds go on top as thin lines, so the reader
# sees them tracking rather than being hidden by whichever curve was plotted last.
arms = [
('fw135m_rlin', 'Equilibrium Propagation', '#1f6feb', 3.0, '-', 2),
('fw135m_bsign', 'Equilibrium Propagation,\nwithout the readout fix', '#d1442f', 2.6, '-', 2),
('fw135m_bp', 'Backprop, seed 1', '#2b3540', 1.0, '-', 4),
('fw135m_bp_s2', 'Backprop, seed 2', '#7d8794', 1.0, '-', 4),
]
fig, ax = plt.subplots(figsize=(7.6, 4.9))
finals = {}
for tag, label, color, lw, ls, zo in arms:
pts = curve(tag)
if not pts:
print(f'skip {tag}: no data')
continue
xs, ys = smooth(pts)
toks = [x * TOK_PER_STEP / 1e9 for x in xs]
ax.plot(toks, ys, color=color, lw=lw, ls=ls, label=label, zorder=zo)
tail = [v for s, v in pts if s >= 0.9 * TOTAL]
finals[tag] = st.mean(tail) if tail else float('nan')
print(f'{tag:14s} last step {pts[-1][0]:>7} tail mean {finals[tag]:.4f}')
ax.set_xlabel('training tokens (billions)')
ax.set_ylabel('validation cross-entropy')
ax.set_ylim(3.15, 4.35)
ax.set_xlim(0, TOTAL * TOK_PER_STEP / 1e9)
ax.grid(alpha=0.18, lw=0.7)
for side in ('top', 'right'):
ax.spines[side].set_visible(False)
h, l = ax.get_legend_handles_labels()
order = [2, 3, 0, 1] # backprop seeds first in the legend, matching how one reads the plot
leg = ax.legend([h[i] for i in order], [l[i] for i in order],
frameon=False, fontsize=9.5, loc='upper right', handlelength=1.6)
for t in leg.get_texts():
t.set_va('center')
if 'fw135m_bsign' in finals and 'fw135m_rlin' in finals:
ax.annotate('single-precision rounding in the contrast readout;\nthe run stops improving',
xy=(1.9, finals['fw135m_bsign'] + 0.005), xytext=(1.35, 3.86),
fontsize=9, color='#d1442f',
arrowprops=dict(arrowstyle='-', color='#d1442f', lw=0.9, alpha=0.8))
ax.annotate(f"EP {finals['fw135m_rlin']:.3f}\nBP {st.mean([finals['fw135m_bp'], finals['fw135m_bp_s2']]):.3f}",
xy=(2.62, 3.245), fontsize=9.5, color='#2b3540', ha='right')
ax.set_title('135M-parameter transformer language model, FineWeb-Edu',
fontsize=10.5, color='#333', pad=10, loc='left')
fig.tight_layout()
for ext in ('png', 'pdf'):
p = Path(f'/home/yurenh2/ept/assets/figs/fig_135m_curves.{ext}')
fig.savefig(p, dpi=300, bbox_inches='tight')
print('wrote', p)
if __name__ == '__main__':
main()
|
