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
|
"""Replicate the governor's rho meter offline: trainer-faithful NUDGED relax (derive-d via
top CE grad + down-chain vjp, rebuild via up-chain feedforward + d) at fixed ckpts, K sweeps,
fixed val batch. Reports per-sweep residual ratio rho (the exact quantity beta_cap gates on)
+ the per-block residual profile of the dominant mode (localization), per lineage x step.
This is the operator whose contraction collapse killed crown-3; free-state norm audits
(specaudit, statej) could not see it — the growth may live in curvature/alignment."""
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('--ckpts', default=('plain:150000,plain:185000,plain:195000,plain:200000,'
'plain:210000,plain:230000,cent:150000,cent:185000,'
'cent:195000,cent:200000,cent:210000,cent:230000'))
ap.add_argument('--K', type=int, default=30)
ap.add_argument('--beta', type=float, default=3e-3)
a = ap.parse_args()
dev = 'cuda'
torch.manual_seed(7)
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))
first = True
for spec in a.ckpts.split(','):
lineage, step = spec.split(':'); step = int(step)
p = f'runs/fw72m_{lineage}_s{step}.pt'
try:
ck = torch.load(p, map_location='cpu', weights_only=False)
except FileNotFoundError:
print(f'{lineage} s{step}: MISSING', flush=True); continue
cfg = ck['config']; C, H, L = cfg['C'], cfg['H'], cfg['L']
if first:
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')
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]).to(dev)
y = torch.stack([torch.from_numpy(data[i + 1:i + 1 + T].astype(np.int64)) for i in ix]).to(dev)
first = False
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
def readout(z): return ln_f(z) @ W_out.t()
# free feedforward pass (block-wise graphs, EP-style detach between blocks)
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
rhos, res_list = [], []
prof = None
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] = (-a.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 = [], []
rnum, rden, pblk = 0.0, 0.0, []
for l in range(L):
i = prev.detach().requires_grad_(True)
o = blocks[l](i)
n_ins.append(i); n_outs.append(o)
znew = o.detach().float() + d[l]
dn = float((znew - zs[l]).norm())
rnum += dn; rden += float(zs[l].norm()); pblk.append(dn)
zs[l] = znew
prev = zs[l]
ins, outs = n_ins, n_outs
res = rnum / max(rden, 1e-9)
if res_list: rhos.append(res / max(res_list[-1], 1e-12))
res_list.append(res)
prof = pblk
if not np.isfinite(res) or res > 1e3:
print(f'{lineage} s{step//1000}k: DIVERGED at sweep {k} (res {res:.2e})', flush=True)
break
tail = rhos[-5:] if len(rhos) >= 5 else rhos
pn = np.array(prof) / (np.sum(prof) + 1e-30)
print(f'{lineage} s{step//1000}k | res0 {res_list[0]:.4f} resK {res_list[-1]:.2e} | '
f'rho tail-med {np.median(tail):.4f} max {max(rhos):.4f} | mode blk-profile ' +
' '.join(f'{v:.2f}' for v in pn), flush=True)
del tok, blocks, W_out, ln_f, ins, outs, f_ins, f_outs, zs, d
torch.cuda.empty_cache()
print('RHORELAX_DONE', flush=True)
|
