summaryrefslogtreecommitdiff
path: root/ep_run
diff options
context:
space:
mode:
Diffstat (limited to 'ep_run')
-rw-r--r--ep_run/build_notebook.py141
-rw-r--r--ep_run/casc_eq_train.py43
-rw-r--r--ep_run/fig_curves.py94
3 files changed, 278 insertions, 0 deletions
diff --git a/ep_run/build_notebook.py b/ep_run/build_notebook.py
new file mode 100644
index 0000000..7ef23dd
--- /dev/null
+++ b/ep_run/build_notebook.py
@@ -0,0 +1,141 @@
+"""Build the demo notebook: samples from the EP-trained 135M model beside its backprop twin.
+
+The notebook ships with its outputs already filled in, so it can be read without a checkout or a
+GPU, and the code in it is real, so it also runs if you have the checkpoints.
+"""
+import json
+import re
+from pathlib import Path
+
+RUN = Path('/home/yurenh2/ept/ep_run')
+OUT = Path('/home/yurenh2/ept/assets/EPT_135M_samples.ipynb')
+PROMPTS = [
+ 'The main difference between a virus and a bacterium is',
+ 'To find the area of a circle, you',
+ 'In 1815, the eruption of Mount Tambora',
+ 'Photosynthesis is the process by which',
+]
+
+
+def parse(path):
+ txt = Path(path).read_text(encoding='utf-8', errors='replace')
+ parts = re.split(r'\n--- prompt (\d+), sample (\d+) ---\n', txt)
+ out = {}
+ for i in range(1, len(parts), 3):
+ out[(int(parts[i]), int(parts[i + 1]))] = parts[i + 2].strip()
+ return out
+
+
+def md(source):
+ return {'cell_type': 'markdown', 'metadata': {}, 'source': source.splitlines(keepends=True)}
+
+
+def code(source, stdout=None):
+ outputs = []
+ if stdout is not None:
+ outputs.append({'output_type': 'stream', 'name': 'stdout',
+ 'text': stdout.splitlines(keepends=True)})
+ return {'cell_type': 'code', 'execution_count': None, 'metadata': {},
+ 'source': source.splitlines(keepends=True), 'outputs': outputs}
+
+
+def main():
+ ep, bp = parse(RUN / 'runs/gen_rlin.txt'), parse(RUN / 'runs/gen_bp.txt')
+ cells = [md("""# A 135M language model trained without backpropagation
+
+Every parameter update in the model below comes from Equilibrium Propagation. There is no backward
+pass anywhere in its training, and no layer, block, or output head is trained with a backprop rule.
+The model is an ordinary 12 layer transformer, OLMo2 style blocks, 32k vocabulary, trained from
+scratch on 2.7B tokens of FineWeb-Edu.
+
+Beside it, for the same prompts and the same sampling seeds, is its backprop twin: identical
+architecture, tokenizer, data order, optimizer, step budget, and evaluation, differing only in the
+training rule. Final validation cross-entropy is 3.209 for EP and 3.207 for backprop, against a
+spread of 0.006 between backprop seeds.
+
+Sampling is an ordinary forward pass. Equilibrium Propagation appears only during training, so
+nothing unusual happens here at inference time.
+"""),
+ code("""from pathlib import Path
+import re, subprocess
+
+RUN = Path('ep_run') # your checkout
+CKPT = {'EP': 'runs/fw135m_rlin_s440000.pt', # trained with Equilibrium Propagation
+ 'BP': 'runs/fw135m_bp_s440000.pt'} # the matched backprop twin
+
+def samples(which, n=2, new=110, temp=0.8, topk=40):
+ \"\"\"Return {(prompt, sample): text}. Uses the cached generation if it is present.\"\"\"
+ cached = RUN / 'runs' / ('gen_rlin.txt' if which == 'EP' else 'gen_bp.txt')
+ if cached.exists():
+ txt = cached.read_text(encoding='utf-8', errors='replace')
+ else:
+ txt = subprocess.run([
+ 'python3', 'casc_eq_train.py', '--gen', str(n), '--gen_new', str(new),
+ '--gen_temp', str(temp), '--gen_topk', str(topk), '--tag', f'gen_{which}',
+ '--resume', CKPT[which], '--data', 'fineweb_edu', '--wandb', '', '--untie',
+ '--L', '12', '--C', '768', '--H', '12', '--T', '256', '--B', '4',
+ '--olmo2', '--steps', '440000'], cwd=RUN, capture_output=True, text=True).stdout
+ parts = re.split(r'\\n--- prompt (\\d+), sample (\\d+) ---\\n', txt)
+ return {(int(parts[i]), int(parts[i+1])): parts[i+2].strip() for i in range(1, len(parts), 3)}
+
+EP, BP = samples('EP'), samples('BP')
+print(f'{len(EP)} samples from the EP model, {len(BP)} from the backprop twin')""",
+ f'{len(ep)} samples from the EP model, {len(bp)} from the backprop twin\n'),
+ md("""## Side by side
+
+Same prompt, same sampling seed, same temperature and top-k. The left column was trained without
+backpropagation.
+"""),
+ ]
+
+ def render(pi):
+ lines = [f'PROMPT: {PROMPTS[pi-1]}', '']
+ for si in (1, 2):
+ lines += [f' [ trained with Equilibrium Propagation, sample {si} ]',
+ ' ' + ep.get((pi, si), '').replace('\n', '\n '), '',
+ f' [ trained with backpropagation, sample {si} ]',
+ ' ' + bp.get((pi, si), '').replace('\n', '\n '), '']
+ return '\n'.join(lines) + '\n'
+
+ src = """def show(pi):
+ print(f'PROMPT: {PROMPTS[pi-1]}\\n')
+ for si in (1, 2):
+ for name, d in (('Equilibrium Propagation', EP), ('backpropagation', BP)):
+ print(f' [ trained with {name}, sample {si} ]')
+ print(' ' + d[(pi, si)].replace('\\n', '\\n '), '\\n')
+
+PROMPTS = %r
+show(%d)"""
+ cells.append(code(src % (PROMPTS, 1), render(1)))
+ for pi in (2, 3, 4):
+ cells.append(code(f'show({pi})', render(pi)))
+
+ cells.append(md("""## What this does and does not show
+
+The two models are close in quality because Equilibrium Propagation is designed to compute the same
+gradient as backpropagation in the small nudge limit, and our measurements confirm that it does:
+cosine to the backprop gradient stays near 0.9997 through training. The interest is not that the
+gradient is different. It is that this update is local, so a physical substrate can perform it
+without a global backward pass, and that it survives at this scale under the constraints such a
+substrate imposes, including 8 bit weights, injected device noise, and the precision of the
+contrast readout.
+
+That last one turned out to matter more than we expected. An earlier version of this run plateaued
+around 3.58 because the contrast was recovered by subtracting two large states, which destroys the
+part of the nudge that falls below single precision resolution. Reading the contrast from the
+stored displacement instead removes the effect at no cost, and is what separates the two EP curves
+in the accompanying figure.
+"""))
+
+ import uuid
+ for c in cells:
+ c['id'] = uuid.uuid4().hex[:8] # nbformat 4.5 requires cell ids
+ nb = {'cells': cells, 'metadata': {'kernelspec': {'display_name': 'Python 3',
+ 'language': 'python', 'name': 'python3'}, 'language_info': {'name': 'python'}},
+ 'nbformat': 4, 'nbformat_minor': 5}
+ OUT.write_text(json.dumps(nb, indent=1, ensure_ascii=False))
+ print('wrote', OUT, f'({OUT.stat().st_size/1024:.1f} KB, {len(cells)} cells)')
+
+
+if __name__ == '__main__':
+ main()
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index ed164ce..05c1e4e 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -53,6 +53,11 @@ ap.add_argument('--dgain_geo', type=float, default=0.0) # >0: per-layer geomet
# gain = geo^l (layer 0 = x1), optionally capped
ap.add_argument('--dgain_geo_cap', type=float, default=0.0) # >0: cap for the geometric profile
ap.add_argument('--dgain_rand', type=float, default=0.0)
+ap.add_argument('--gen', type=int, default=0) # >0: sample text from the resumed checkpoint and exit
+ap.add_argument('--gen_prompts', default='') # '|'-separated prompts; empty = built-in set
+ap.add_argument('--gen_temp', type=float, default=0.8)
+ap.add_argument('--gen_topk', type=int, default=40)
+ap.add_argument('--gen_new', type=int, default=120)
ap.add_argument('--probe_dgspec', type=int, default=0) # >0: M1 spectroscopy, value = n batches; exits before training
ap.add_argument('--probe_gains', default='1,2,4,8,16,32,64,128,256')
ap.add_argument('--probe_f64', action='store_true') # fp64 states+model in the probe: the fp-floor decisive arm # >1: per-STEP log-uniform dgain_top in
@@ -859,6 +864,44 @@ if args.ddp_grad_test:
import sys
sys.exit(0)
+if args.gen > 0:
+ # Sampling from a resumed checkpoint. Inference here is an ordinary forward pass through the
+ # blocks, which is the point: EP appears only in training. The same code path loads an EP or a
+ # backprop checkpoint, since both scripts save the same keys, so the two can be sampled side by
+ # side under identical settings.
+ import sys
+ from tokenizers import Tokenizer as _Tok
+ _tk = _Tok.from_file(str(DD / 'tokenizer.json'))
+ prompts = [p for p in args.gen_prompts.split('|') if p] or [
+ 'The main difference between a virus and a bacterium is',
+ 'To find the area of a circle, you',
+ 'In 1815, the eruption of Mount Tambora',
+ 'Photosynthesis is the process by which',
+ ]
+ @torch.no_grad()
+ def _sample(prompt, n_new, temp, topk, seed):
+ torch.manual_seed(seed)
+ ids = _tk.encode(prompt).ids[:args.T - n_new - 1]
+ idx = torch.zeros(1, args.T, dtype=torch.long, device=dev)
+ L = len(ids)
+ idx[0, :L] = torch.tensor(ids, device=dev)
+ for _ in range(n_new):
+ if L >= args.T: break
+ z = emb(idx)
+ for b in blocks: z = b(z, mask)
+ lg = readout(z)[0, L - 1].float() / max(temp, 1e-6)
+ v, _i = torch.topk(lg, topk)
+ lg[lg < v[-1]] = -float('inf')
+ nt = torch.multinomial(F.softmax(lg, -1), 1).item()
+ idx[0, L] = nt; L += 1
+ return _tk.decode(idx[0, :L].tolist())
+ print(f'# samples from {args.resume} (temp {args.gen_temp}, top-k {args.gen_topk})', flush=True)
+ for pi, p in enumerate(prompts):
+ for s in range(args.gen):
+ print(f'\n--- prompt {pi + 1}, sample {s + 1} ---', flush=True)
+ print(_sample(p, args.gen_new, args.gen_temp, args.gen_topk, 1000 * pi + s), flush=True)
+ sys.exit(0)
+
if args.probe_dgspec > 0:
# M1 DGAIN SPECTROSCOPY: per-block leak vector vs uniform read-displacement gain.
# Paired design: L_l(g) = mean_b[gEP_l(g) - gBP_l] on the SAME batch — batch-sampling
diff --git a/ep_run/fig_curves.py b/ep_run/fig_curves.py
new file mode 100644
index 0000000..9185702
--- /dev/null
+++ b/ep_run/fig_curves.py
@@ -0,0 +1,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()
generated by cgit v1.2.3 (git 2.25.1) at 2026-08-18 13:27:18 +0000