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
202
203
204
205
206
207
208
209
210
211
212
213
|
"""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) # free-phase draw
qblocksN = copy.deepcopy(blocks) # independent nudged-phase draw (split mode)
_qsrc = [p.detach().clone() for p in blocks.parameters()]
def requantize(target):
"""stochastic rounding, fresh draw each call (the per-step operator-noise ensemble)"""
if a.qbits <= 0: return
with torch.no_grad():
for p, s in zip(target.parameters(), _qsrc):
if p.dim() == 2:
sc = s.abs().max() / (2 ** (a.qbits - 1) - 1)
x = s / sc
p.copy_((torch.floor(x + torch.rand_like(x))) * sc)
else:
p.copy_(s)
def ep_and_bp(x, y, beta, mode='frozen'):
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)]
requantize(qblocks)
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
QB = qblocks
if mode == 'split':
requantize(qblocksN); QB = qblocksN
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 = QB[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(QB.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 (names[i].split('.')[0], names[i].split('.')[1]) in
(('0','attn'), ('8','attn'), ('6','ff'), ('11','ff'))]
betas = [float(s) for s in a.betas.split(',')]
REP = 6
x, y = get_batch()
g_bp_ref = None
print(f'qbits={a.qbits} stochastic rounding | one batch, {REP} rounding draws per estimate', flush=True)
for mode in ('frozen', 'split'):
store = {i: {b: [] for b in betas} for i in SEL}
for b in betas:
for r in range(REP):
g_bp, g_ep = ep_and_bp(x, y, b, mode)
if g_bp_ref is None: g_bp_ref = [g.detach().cpu() for g in g_bp]
for i in SEL: store[i][b].append(g_ep[i].detach().cpu())
del g_bp, g_ep
torch.cuda.empty_cache()
print(f'--- mode={mode} (frozen=trainer-faithful one draw/estimate; split=decorrelated free vs nudged) ---', flush=True)
for i in SEL:
m, n = params[i].shape
gb = g_bp_ref[i]
s1 = float(torch.linalg.svdvals(gb)[0])
u1 = torch.linalg.svd(gb, full_matrices=False).U[:, 0]
cells, Rs = [], []
for b in betas:
Xi = torch.stack(store[i][b])
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[i][b]]
Rs.append(R); cells.append(f'{R:8.2f}({np.mean(ovs):.3f})')
lb = np.log(np.array(betas)); lR = np.log(np.maximum(Rs, 1e-12))
tag = 'no-crossing(R<1 all)' if max(Rs) < 1 else ('no-crossing(R>1 all)' if min(Rs) > 1 else '')
if not tag:
for j in range(len(betas) - 1):
if (lR[j]) * (lR[j + 1]) <= 0 and lR[j] != lR[j + 1]:
t = (0.0 - lR[j]) / (lR[j + 1] - lR[j]); tag = f'beta*={float(np.exp(lb[j] + t * (lb[j + 1] - lb[j]))):.2e}'; break
print(f'{names[i]:22s} {m:5d}x{n:<5d} s1 {s1:9.3g} ' + ' '.join(cells) + f' {tag}', flush=True)
print('BBP3_DONE', flush=True)
|
