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
|
"""BBP floor: is wall-1 a spiked-matrix detectability transition? Per-layer model
g_hat(beta) = g_true + Xi/beta, Xi = EP-specific error (additive component).
Measure across batches x betas: (i) entry-std s(beta) of (g_EP - g_BP), fit s = a/beta (+) b
to split additive a (BBP-active) from multiplicative b (co-scaling, exempt per r-sweep);
(ii) sigma1(g_BP) per layer; -> BBP/BGN threshold beta*_l = a_l*(sqrt(m)+sqrt(n))/2 / sigma1_l
(iid-noise convention: bulk edge of Xi/beta at std nu=a/beta per entry is nu*(sqrt(m)+sqrt(n))/
sqrt(mn)*sqrt(mn)= a/beta*(sqrt m + sqrt n); spike detaches iff sigma1 > that /2..1 band —
report both edge conventions); (iii) EMPIRICAL overlap cos(u1(g_hat), u1(g_BP)) vs beta —
the BBP order parameter, compare its rise against beta*.
Layers: per-block attn.qkv + ff.w2 (the two families), blocks 0..11."""
import argparse, pickle
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
from pathlib import Path
ap = argparse.ArgumentParser()
ap.add_argument('--ckpt', default='runs/fw72m_plain_s150000.pt')
ap.add_argument('--K', type=int, default=3)
ap.add_argument('--betas', default='3e-4,1e-3,3e-3,1e-2')
ap.add_argument('--nb', type=int, default=6)
a = ap.parse_args()
dev = 'cuda'
torch.manual_seed(11)
B, T = 8, 256
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 Olmo2Attn(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 Olmo2Block(nn.Module):
def __init__(self, C, H, T):
super().__init__()
self.attn = Olmo2Attn(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))
ck = torch.load(a.ckpt, map_location='cpu', weights_only=False)
cfg = ck['config']; C, H, L = cfg['C'], cfg['H'], cfg['L']
DD = Path('/home/yurenh2/ept/ep_run/data') / cfg.get('data', 'fineweb_edu')
vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size']
data = np.memmap(DD / 'val.bin', dtype=np.uint16, mode='r')
tok = nn.Embedding(vocab, C).to(dev); tok.load_state_dict(ck['tok'])
blocks = nn.ModuleList([Olmo2Block(C, H, T) for _ in range(L)]).to(dev)
blocks.load_state_dict(ck['blocks'], strict=False)
W_out = ck['wout'].to(dev)
ln_f = RMSNorm(C).to(dev); ln_f.load_state_dict(ck['lnf'])
NBT = B * T
params = list(blocks.parameters())
names = [n for n, _ in blocks.named_parameters()]
def readout(z): return ln_f(z) @ W_out.t()
def get_batch():
ix = torch.randint(len(data) - T - 1, (B,))
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 grads_from(outs_list, cots):
obj = sum((o * c.detach()).sum() for o, c in zip(outs_list, cots))
gs = torch.autograd.grad(obj, params, allow_unused=True, retain_graph=True)
return [g.float() if g is not None else torch.zeros_like(p) for g, p in zip(gs, params)]
def ep_and_bp(x, y, beta):
z = tok(x)
zs_bp = []
for b in blocks:
z = b(z); zs_bp.append(z)
ce = F.cross_entropy(readout(zs_bp[-1]).reshape(-1, vocab), y.reshape(-1))
g_bp = [g.float() for g in torch.autograd.grad(ce, params, retain_graph=True, allow_unused=False)]
f_ins, f_outs = [], []
prev = tok(x).detach()
for b in blocks:
i = prev.detach().requires_grad_(True)
o = b(i)
f_ins.append(i); f_outs.append(o)
prev = o.detach()
ins, outs = f_ins, f_outs
zs = [o.detach().float() for o in f_outs]
d = [None] * L
for k in range(a.K):
zc = zs[L - 1].detach().requires_grad_(True)
ce_k = F.cross_entropy(readout(zc).reshape(-1, vocab), y.reshape(-1))
d[L - 1] = (-beta * NBT * torch.autograd.grad(ce_k, 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
md = [-di / (beta * NBT) for di in d] # normalize so EP grad is on BP scale
g_ep = grads_from(outs, md)
return g_bp, g_ep
SEL = [i for i, n in enumerate(names) if n.endswith('attn.qkv.weight') or n.endswith('ff.w2.weight')]
betas = [float(s) for s in a.betas.split(',')]
batches = [get_batch() for _ in range(a.nb)]
# collect per layer: err entry-std per beta (across batches), sigma1(gbp), top-vec overlap per beta
acc = {i: {b: {'s': [], 'ov': []} for b in betas} for i in SEL}
sig1 = {i: [] for i in SEL}
for (x, y) in batches:
ref = None
for b in betas:
g_bp, g_ep = ep_and_bp(x, y, b)
for i in SEL:
gb, ge = g_bp[i], g_ep[i]
if ref is None: pass
err = ge - gb
acc[i][b]['s'].append(float(err.std()))
try:
ub = torch.linalg.svd(gb, full_matrices=False).U[:, 0]
ue = torch.linalg.svd(ge, full_matrices=False).U[:, 0]
acc[i][b]['ov'].append(abs(float(ub @ ue)))
except Exception:
acc[i][b]['ov'].append(float('nan'))
for i in SEL:
if b == betas[0]: sig1[i].append(float(torch.linalg.svdvals(g_bp[i])[0]))
del g_bp, g_ep
torch.cuda.empty_cache()
print('layer m x n sigma1(gBP) a(add) b(mult) beta*_edge ov@' +
' ov@'.join(f'{b:g}' for b in betas), flush=True)
for i in SEL:
m, n = params[i].shape
s1 = float(np.mean(sig1[i]))
ss = np.array([np.mean(acc[i][b]['s']) for b in betas])
X = np.vstack([1.0 / np.array(betas), np.ones(len(betas))]).T
coef, *_ = np.linalg.lstsq(X, ss, rcond=None)
aa, bb = max(coef[0], 0.0), max(coef[1], 0.0)
edge = aa * (np.sqrt(m) + np.sqrt(n))
bstar = edge / max(s1, 1e-30)
ovs = ' '.join(f'{np.nanmean(acc[i][b]["ov"]):.3f}' for b in betas)
print(f'{names[i]:22s} {m:5d}x{n:<5d} {s1:11.4g} {aa:10.3g} {bb:10.3g} {bstar:11.3g} {ovs}', flush=True)
print('BBP_DONE', flush=True)
|
