"""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')