summaryrefslogtreecommitdiff
path: root/ep_run
diff options
context:
space:
mode:
Diffstat (limited to 'ep_run')
-rw-r--r--ep_run/fieldviz.py281
1 files changed, 281 insertions, 0 deletions
diff --git a/ep_run/fieldviz.py b/ep_run/fieldviz.py
new file mode 100644
index 0000000..f5e5e6a
--- /dev/null
+++ b/ep_run/fieldviz.py
@@ -0,0 +1,281 @@
+"""2D update-vector-field analysis (professor's suggestion): PCA plane through the
+training trajectories; loss contours + BP/EP-small-beta/EP-big-beta quiver fields;
+curl (non-conservativity) map; trajectory overlays; 1D ride<->BP interpolation.
+All from existing checkpoints. Outputs to ../assets/figs/."""
+import argparse, math, pickle, re
+import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
+from pathlib import Path
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+
+ap = argparse.ArgumentParser()
+ap.add_argument('--grid', type=int, default=13)
+ap.add_argument('--nb', type=int, default=2) # batches averaged per field eval
+ap.add_argument('--K', type=int, default=3)
+a = ap.parse_args()
+dev = 'cuda'
+torch.manual_seed(11)
+R = Path('/home/yurenh2/ept/ep_run/runs')
+DD = Path('/home/yurenh2/ept/ep_run/data/tinystories_bpe')
+vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size']
+B, T = 8, 256
+
+# ---- model (matches probe_blockcos) ----
+class RMSNorm(nn.Module):
+ def __init__(self, C, eps=1e-6):
+ super().__init__(); self.g = nn.Parameter(torch.ones(C)); self.eps = eps
+ def forward(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.g
+
+class SwiGLU(nn.Module):
+ def __init__(self, C):
+ super().__init__()
+ h = ((8 * C // 3) + 63) // 64 * 64
+ self.w1 = nn.Linear(C, h, bias=False); self.w3 = nn.Linear(C, h, bias=False)
+ self.w2 = nn.Linear(h, C, bias=False)
+ def forward(self, x):
+ return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+class Attn(nn.Module):
+ def __init__(self, C, H, T):
+ super().__init__()
+ self.H, self.hd = H, C // H
+ self.qkv = nn.Linear(C, 3 * C, bias=False); self.proj = nn.Linear(C, C, bias=False)
+ self.qn, self.kn = RMSNorm(C), RMSNorm(C)
+ inv = 1.0 / (500000.0 ** (torch.arange(0, self.hd, 2).float() / self.hd))
+ fr = torch.outer(torch.arange(T).float(), inv)
+ self.register_buffer('rc', fr.cos(), persistent=False)
+ self.register_buffer('rs', fr.sin(), persistent=False)
+ def rope(self, x):
+ Tn = x.shape[2]
+ x1, x2 = x[..., ::2], x[..., 1::2]
+ c, s = self.rc[None, None, :Tn].to(x.dtype), self.rs[None, None, :Tn].to(x.dtype)
+ return torch.stack((x1 * c - x2 * s, x1 * s + x2 * c), dim=-1).flatten(-2)
+ def forward(self, x):
+ Bn, Tn, C = x.shape
+ q, k, v = self.qkv(x).split(C, dim=2)
+ q, k = self.qn(q), self.kn(k)
+ q = self.rope(q.view(Bn, Tn, self.H, self.hd).transpose(1, 2))
+ k = self.rope(k.view(Bn, Tn, self.H, self.hd).transpose(1, 2))
+ v = v.view(Bn, Tn, self.H, self.hd).transpose(1, 2)
+ y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+ return self.proj(y.transpose(1, 2).contiguous().view(Bn, Tn, C))
+
+class Block(nn.Module):
+ def __init__(self, C, H, T):
+ super().__init__()
+ self.attn = Attn(C, H, T); self.ff = SwiGLU(C)
+ self.na, self.nf = RMSNorm(C), RMSNorm(C)
+ def forward(self, z):
+ z = z + self.na(self.attn(z))
+ return z + self.nf(self.ff(z))
+
+C_, H_, L_ = 512, 8, 12
+tok = nn.Embedding(vocab, C_).to(dev)
+blocks = nn.ModuleList([Block(C_, H_, T) for _ in range(L_)]).to(dev)
+W_out = torch.zeros(vocab, C_, device=dev)
+ln_f = RMSNorm(C_).to(dev)
+params = [p for p in blocks.parameters()]
+NBT = B * T
+
+def flat_from_ckpt(path):
+ ck = torch.load(path, map_location='cpu', weights_only=False)
+ blocks.load_state_dict(ck['blocks'], strict=False)
+ tok.load_state_dict(ck['tok'])
+ global W_out
+ W_out = ck['wout'].to(dev)
+ ln_f.load_state_dict(ck['lnf'])
+ vecs = [p.detach().reshape(-1).cpu() for p in blocks.parameters()]
+ vecs += [tok.weight.detach().reshape(-1).cpu(), W_out.detach().reshape(-1).cpu(),
+ ln_f.g.detach().reshape(-1).cpu()]
+ return torch.cat(vecs).float(), ck
+
+def set_flat(v):
+ i = 0
+ with torch.no_grad():
+ for p in params:
+ n = p.numel()
+ p.copy_(v[i:i + n].view_as(p).to(dev)); i += n
+ for t_ in (tok.weight, W_out, ln_f.g):
+ n = t_.numel()
+ t_.copy_(v[i:i + n].view_as(t_).to(dev)); i += n
+
+def load_aux(ck):
+ tok.load_state_dict(ck['tok'])
+ global W_out
+ W_out = ck['wout'].to(dev)
+ ln_f.load_state_dict(ck['lnf'])
+
+def readout(z): return ln_f(z) @ W_out.t()
+
+data = np.memmap(DD / 'val.bin', dtype=np.uint16, mode='r')
+def get_batch(seed):
+ g = torch.Generator().manual_seed(seed)
+ ix = torch.randint(len(data) - T - 1, (B,), generator=g)
+ x = torch.stack([torch.from_numpy(data[i:i + T].astype(np.int64)) for i in ix])
+ y = torch.stack([torch.from_numpy(data[i + 1:i + 1 + T].astype(np.int64)) for i in ix])
+ return x.to(dev), y.to(dev)
+
+def val_loss(batches):
+ with torch.no_grad():
+ tot = 0.0
+ for x, y in batches:
+ z = tok(x)
+ for b in blocks: z = b(z)
+ tot += float(F.cross_entropy(readout(z).reshape(-1, vocab), y.reshape(-1)))
+ return tot / len(batches)
+
+def bp_grad(batches):
+ acc = None
+ for x, y in batches:
+ z = tok(x)
+ for b in blocks: z = b(z)
+ ce = F.cross_entropy(readout(z).reshape(-1, vocab), y.reshape(-1))
+ gs = torch.autograd.grad(ce, params)
+ v = torch.cat([g.reshape(-1) for g in gs])
+ acc = v if acc is None else acc + v
+ return (acc / len(batches)).detach()
+
+def ep_grad(batches, beta):
+ acc = None
+ for x, y in batches:
+ ins, outs, zs = [], [], []
+ prev = tok(x).detach()
+ for b in blocks:
+ i = prev.detach().requires_grad_(True)
+ o = b(i)
+ ins.append(i); outs.append(o); zs.append(o.detach().float())
+ prev = zs[-1]
+ d = [None] * L_
+ for k in range(a.K):
+ zc = zs[L_ - 1].detach().requires_grad_(True)
+ ce = F.cross_entropy(readout(zc).reshape(-1, vocab), y.reshape(-1))
+ d[L_ - 1] = (-beta * NBT * torch.autograd.grad(ce, zc)[0]).detach().float()
+ for l in range(L_ - 2, -1, -1):
+ d[l] = torch.autograd.grad(outs[l + 1], ins[l + 1], grad_outputs=d[l + 1],
+ retain_graph=True)[0].detach().float()
+ prev = tok(x).detach()
+ n_ins, n_outs = [], []
+ for l in range(L_):
+ i = prev.detach().requires_grad_(True)
+ o = blocks[l](i)
+ n_ins.append(i); n_outs.append(o)
+ zs[l] = o.detach().float() + d[l]
+ prev = zs[l]
+ ins, outs = n_ins, n_outs
+ E = 0.0
+ for z, o in zip(zs, outs): E = E + 0.5 * ((z.detach() - o.float()) ** 2).sum()
+ obj = E / (NBT * beta) + F.cross_entropy(readout(zs[-1].detach()).reshape(-1, vocab), y.reshape(-1))
+ gs = torch.autograd.grad(obj, params, allow_unused=True)
+ v = torch.cat([(g if g is not None else torch.zeros_like(p)).reshape(-1) for g, p in zip(gs, params)])
+ acc = v if acc is None else acc + v
+ return (acc / len(batches)).detach()
+
+# ---- collect snapshots ----
+runs = {
+ 'orig_smallbeta': [f'stage1b_ep_muon_s{s}.pt' for s in range(5000, 56000, 10000)],
+ 'plain_bigbeta': [f'stage1b_plain3e3_s{s}.pt' for s in range(10000, 59000, 10000)] ,
+ 'ride': [f'stage1b_ride_s{s}.pt' for s in range(10000, 59000, 10000)] ,
+ 'cent': [f'stage1b_cent_s{s}.pt' for s in range(10000, 51000, 10000)],
+ 'BP': [f'stage1b_bp_muon_s{s}.pt' for s in range(5000, 56000, 10000)],
+}
+snaps, labels = [], []
+for name, files in runs.items():
+ for f in files:
+ p = R / f
+ if p.exists():
+ v, ck = flat_from_ckpt(p)
+ snaps.append(v); labels.append((name, f))
+print(f'{len(snaps)} snapshots loaded')
+S = torch.stack(snaps) # N x P (cpu fp32)
+mu = S.mean(0)
+Sc = S - mu
+G = Sc @ Sc.t() # N x N Gram
+evals, evecs = torch.linalg.eigh(G)
+d1 = (Sc.t() @ evecs[:, -1]); d1 /= d1.norm()
+d2 = (Sc.t() @ evecs[:, -2]); d2 /= d2.norm()
+coords = torch.stack([Sc @ d1, Sc @ d2], 1) # N x 2
+print('plane variance share:', float((evals[-1] + evals[-2]) / evals.sum()))
+
+# anchor aux weights (tok/wout/lnf) from the ride final ckpt for all grid evals
+
+PB = sum(p.numel() for p in params) # block-subspace size; grads are computed on blocks
+d1b, d2b = d1[:PB].contiguous(), d2[:PB].contiguous()
+batches = [get_batch(1000 + i) for i in range(a.nb)]
+lo = coords.min(0).values; hi = coords.max(0).values
+pad = 0.18 * (hi - lo)
+A = torch.linspace(lo[0] - pad[0], hi[0] + pad[0], a.grid)
+Bx = torch.linspace(lo[1] - pad[1], hi[1] + pad[1], a.grid)
+LOSS = np.zeros((a.grid, a.grid))
+FB = np.zeros((a.grid, a.grid, 2)); FS = np.zeros_like(FB); FL = np.zeros_like(FB)
+for i, aa in enumerate(A):
+ for j, bb in enumerate(Bx):
+ set_flat(mu + aa * d1 + bb * d2)
+ LOSS[j, i] = val_loss(batches)
+ gb = bp_grad(batches)
+ gs_ = ep_grad(batches, 3e-4)
+ gl = ep_grad(batches, 1e-2)
+ for Fm, gv in ((FB, gb), (FS, gs_), (FL, gl)):
+ Fm[j, i, 0] = -float(gv.cpu() @ d1b); Fm[j, i, 1] = -float(gv.cpu() @ d2b)
+ print(f'row {i+1}/{a.grid} done', flush=True)
+
+def curl(Fm):
+ da = float(A[1] - A[0]); db = float(Bx[1] - Bx[0])
+ dFy_dx = np.gradient(Fm[:, :, 1], da, axis=1)
+ dFx_dy = np.gradient(Fm[:, :, 0], db, axis=0)
+ return dFy_dx - dFx_dy
+
+# ---- figure: contours + 3 quivers + curl maps ----
+AA, BB = np.meshgrid(A.numpy(), Bx.numpy())
+colors = {'orig_smallbeta': '#888888', 'plain_bigbeta': '#2c6fbb', 'ride': '#d95f02',
+ 'cent': '#7b5aa6', 'BP': '#2e7d32'}
+fig, axes = plt.subplots(2, 3, figsize=(16.5, 10))
+names = [('BP gradient field', FB), (r'EP field, $\beta=3\times10^{-4}$', FS), (r'EP field, $\beta=10^{-2}$', FL)]
+for ax, (ttl, Fm) in zip(axes[0], names):
+ cs = ax.contour(AA, BB, LOSS, levels=14, colors='#bbbbbb', linewidths=0.7)
+ ax.clabel(cs, fontsize=6, fmt='%.2f')
+ m = np.hypot(Fm[:, :, 0], Fm[:, :, 1])
+ ax.quiver(AA, BB, Fm[:, :, 0] / (m + 1e-12), Fm[:, :, 1] / (m + 1e-12), m,
+ cmap='viridis', scale=28, width=0.004)
+ for name in runs:
+ pts = np.array([coords[k].numpy() for k, (n, _) in enumerate(labels) if n == name])
+ if len(pts):
+ ax.plot(pts[:, 0], pts[:, 1], '-o', color=colors[name], ms=3, lw=1.2, label=name)
+ ax.set_title(ttl, fontsize=11)
+axes[0, 0].legend(fontsize=7, loc='upper left')
+cl_b, cl_s, cl_l = curl(FB), curl(FS), curl(FL)
+vmax = max(abs(cl_l).max(), abs(cl_b).max())
+for ax, (ttl, cl) in zip(axes[1], [('curl(BP) [conservative ref]', cl_b),
+ (r'curl(EP) @ $3\times10^{-4}$', cl_s),
+ (r'curl(EP) @ $10^{-2}$', cl_l)]):
+ im = ax.pcolormesh(AA, BB, cl, cmap='RdBu_r', vmin=-vmax, vmax=vmax)
+ ax.set_title(ttl, fontsize=11)
+ plt.colorbar(im, ax=ax, shrink=0.8)
+fig.suptitle('Update fields (block-subspace component) on the shared full-model PCA plane (42M, TinyStories)',
+ fontsize=13)
+fig.tight_layout()
+fig.savefig('/home/yurenh2/ept/assets/figs/fig_field2d.png', dpi=150)
+print('curl RMS: BP %.3g | EP small %.3g | EP big %.3g' %
+ (np.sqrt((cl_b**2).mean()), np.sqrt((cl_s**2).mean()), np.sqrt((cl_l**2).mean())))
+
+# ---- 1D interpolation ride-final <-> BP-final ----
+v_ride, ck_r = flat_from_ckpt(R / 'stage1b_ride_s50000.pt')
+v_bp, ck_b = flat_from_ckpt(R / 'stage1b_bp_muon_s55000.pt')
+ts = np.linspace(-0.15, 1.15, 27)
+ls = []
+big_batches = [get_batch(5000 + i) for i in range(6)]
+for t in ts:
+ set_flat((1 - t) * v_ride + t * v_bp)
+ ls.append(val_loss(big_batches))
+fig2, ax2 = plt.subplots(figsize=(7.2, 4.6))
+ax2.plot(ts, ls, '-o', color='#2c6fbb', ms=4)
+ax2.axvline(0, color='#d95f02', ls=':'); ax2.axvline(1, color='#2e7d32', ls=':')
+ax2.text(0, max(ls), ' ride final', color='#d95f02', fontsize=9, va='top')
+ax2.text(1, max(ls), ' BP final', color='#2e7d32', fontsize=9, va='top', ha='right')
+ax2.set_xlabel('interpolation t'); ax2.set_ylabel('val CE (6 batches)')
+ax2.set_title('Linear interpolation between EP(ride, s50000) and BP (s55000) weights — same basin?')
+fig2.tight_layout()
+fig2.savefig('/home/yurenh2/ept/assets/figs/fig_interp1d.png', dpi=150)
+print('interp endpoints: ride %.4f BP %.4f max-in-between %.4f' % (ls[3], ls[-4], max(ls[4:-4])))
+print('DONE_FIELDVIZ')
generated by cgit v1.2.3 (git 2.25.1) at 2026-08-18 13:31:07 +0000