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
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
|
"""BBP floor v2 (generalized, measured-spectrum): production AMP semantics (bf16 block
forwards, trainer-faithful) + NO structural noise assumption. Per layer & beta:
signal spike = sigma1( mean_batches g_BP_fp32 )
noise edge = mean_batches sigma1( Xi_fluct ), Xi_fluct = (g_EP_amp - g_BP_fp32) - mean_batches(...)
(systematic truncation bias = the batch-constant mean component -> distortion, removed;
the fluctuation spectrum IS the detection noise, whatever its structure - generalized BBP/BGN)
R(beta) = spike/edge; beta* = crossing of R=1 (log-interp); plus empirical u1-overlap vs beta.
v1 (fp32, iid-additive fit) measured the wrong ensemble: fp32 rounding 2^-23 -> a=0 artifact.
Production floor lives in bf16 (2^-8): RESULT 11 naive-cast death + amp-gate rising cos-vs-beta
+ quant beta-buyback are all BBP signatures. Old docstring below.
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='1e-4,3e-4,1e-3,3e-3,1e-2')
ap.add_argument('--nb', type=int, default=8)
ap.add_argument('--qbits', 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())
qparams = None # set after qblocks exists
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, plist=None):
plist = plist if plist is not None else params
obj = sum((o * c.detach()).sum() for o, c in zip(outs_list, cots))
gs = torch.autograd.grad(obj, plist, 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, plist)]
AMP = torch.autocast('cuda', dtype=torch.bfloat16)
import copy
qblocks = copy.deepcopy(blocks)
if a.qbits > 0:
with torch.no_grad():
for p in qblocks.parameters():
if p.dim() == 2:
sc = p.abs().max() / (2 ** (a.qbits - 1) - 1)
p.copy_(torch.round(p / sc) * sc)
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 qblocks:
i = prev.detach().requires_grad_(True)
with AMP:
o = b(i)
o = o.float()
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)
with AMP:
o = qblocks[l](i)
o = o.float()
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, list(qblocks.parameters()))
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')]
SEL = [i for i in SEL if int(names[i].split('.')[0]) in (0, 6, 8, 11)]
betas = [float(s) for s in a.betas.split(',')]
batches = [get_batch() for _ in range(a.nb)]
# pass 1: per batch/beta store g_ep; per batch store g_bp (fp32 reference)
store_ep = {i: {b: [] for b in betas} for i in SEL}
store_bp = {i: [] for i in SEL}
for (x, y) in batches:
for bi, b in enumerate(betas):
g_bp, g_ep = ep_and_bp(x, y, b)
for i in SEL:
store_ep[i][b].append(g_ep[i].detach().cpu())
if bi == 0: store_bp[i].append(g_bp[i].detach().cpu())
del g_bp, g_ep
torch.cuda.empty_cache()
print('layer m x n spike=s1(gbar) ' +
' '.join(f'R@{b:g}(ov)' for b in betas) + ' beta*(R=1)', flush=True)
for i in SEL:
m, n = params[i].shape
gbar = torch.stack(store_bp[i]).mean(0)
s1 = float(torch.linalg.svdvals(gbar)[0])
u1 = torch.linalg.svd(gbar, full_matrices=False).U[:, 0]
Rs, cells = [], []
for b in betas:
Xi = torch.stack([ge - gb for ge, gb in zip(store_ep[i][b], store_bp[i])])
Xif = Xi - Xi.mean(0, keepdim=True)
edge = float(np.mean([torch.linalg.svdvals(Xif[j])[0] for j in range(Xif.shape[0])]))
R = s1 / max(edge, 1e-30)
ovs = [abs(float(u1 @ torch.linalg.svd(ge, full_matrices=False).U[:, 0])) for ge in store_ep[i][b]]
Rs.append(R); cells.append(f'{R:8.2f}({np.mean(ovs):.3f})')
bstar = float('nan')
lb = np.log(np.array(betas)); lR = np.log(np.maximum(Rs, 1e-12))
for j in range(len(betas) - 1):
if (lR[j] - 0.0) * (lR[j + 1] - 0.0) <= 0 and lR[j] != lR[j + 1]:
t = (0.0 - lR[j]) / (lR[j + 1] - lR[j]); bstar = float(np.exp(lb[j] + t * (lb[j + 1] - lb[j]))); break
print(f'{names[i]:22s} {m:5d}x{n:<5d} {s1:12.4g} ' + ' '.join(cells) +
f' {bstar:.2e}' if bstar == bstar else f'{names[i]:22s} {m:5d}x{n:<5d} {s1:12.4g} ' + ' '.join(cells) + ' R>1 everywhere', flush=True)
print('BBP2Q_DONE', flush=True)
|
