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
|
"""bf16 mixed-precision GATE for cascade-EP (the /2 cost lever; COST_MODEL.md consequence #3).
Distinct from the DEAD naive-cast (--bf16, RESULT 11): here params/states/displacements/E-accum
stay fp32; ONLY matmul-heavy forwards run under torch.autocast(bf16). Modes:
fp32 — reference EP (expect cos(EP,BP_fp32) ~0.97, the E-tier clean band)
amp_all — autocast on every forward incl. the last rebuild (E/theta-read graph is bf16)
amp_last — autocast on free pass + intermediate rebuilds + force chains; LAST rebuild fp32
(the E-readout subtraction (z-o) is the wall-1-sensitive additive term -> exact)
bp_amp — BP under autocast, vs BP fp32 (the fair 'amp hurts everyone' baseline)
Gate: cos(EP_amp, BP_fp32) within the fp32-EP band (>=0.96) at beta_floor 1e-3 => PASS.
Runs on native-bf16 hardware only (A6000/sm86+); Pascal emulation is not representative.
Usage: amp_gate.py [--ckpt ...] [--K 3] [--nb 4]
"""
import argparse, math, 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/stage1b_ep_muon_s55000.pt')
ap.add_argument('--K', type=int, default=3)
ap.add_argument('--nb', type=int, default=4)
a = ap.parse_args()
dev = 'cuda'
torch.manual_seed(7)
assert torch.cuda.is_bf16_supported(), 'gate requires native bf16 hardware'
DD = Path('/home/yurenh2/ept/ep_run/data/tinystories_bpe')
vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size']
B, T = 8, 256
def get_batch():
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])
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)
# ---- model (OLMo2 cascade, matches trainer/etier_probe) ----
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=dev, weights_only=False)
cfg = ck['config']; C, H, L = cfg['C'], cfg['H'], cfg['L']
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()
def AC(on): return torch.autocast('cuda', dtype=torch.bfloat16, enabled=on)
def ep_grad(x, y, beta, amp='off'):
"""single-sided fb EP grad; amp='off'|'all'|'last'. States/d/E stay fp32 in ALL modes;
autocast wraps block forwards only. amp='last' -> final rebuild (E/theta graph) fp32."""
relax_amp = amp in ('all', 'last')
# graphed free pass (cascade free equilibrium = one forward sweep)
ins, outs, zs = [], [], []
prev = tok(x).detach()
with AC(relax_amp):
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)) # top force fp32 always
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].to(outs[l + 1].dtype))[0].detach().float()
last = (k + 1 == a.K)
round_amp = (amp == 'all') or (relax_amp and not last)
prev = tok(x).detach()
n_ins, n_outs = [], []
with AC(round_amp):
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().float() - o.float()) ** 2).sum()
obj = E / (NBT * beta)
params = [p for p in blocks.parameters()]
gs = torch.autograd.grad(obj, params, allow_unused=True)
return [g.float() if g is not None else torch.zeros(1, device=dev) for g in gs]
def bp_grad(x, y, amp=False):
with AC(amp):
z = tok(x)
for b in blocks: z = b(z)
ce = F.cross_entropy(readout(z.float()).reshape(-1, vocab), y.reshape(-1))
params = [p for p in blocks.parameters()]
gs = torch.autograd.grad(ce, params, allow_unused=True)
return [g.float() if g is not None else torch.zeros(1, device=dev) for g in gs]
def val_ce(x, y, amp=False):
with torch.no_grad(), AC(amp):
z = tok(x)
for b in blocks: z = b(z)
return float(F.cross_entropy(readout(z.float()).reshape(-1, vocab), y.reshape(-1)))
def cos(ga, gb):
va = torch.cat([g.reshape(-1).double() for g in ga]); vb = torch.cat([g.reshape(-1).double() for g in gb])
return float((va @ vb) / (va.norm() * vb.norm() + 1e-30))
batches = [get_batch() for _ in range(a.nb)]
ce32 = sum(val_ce(x, y) for x, y in batches) / a.nb
ce16 = sum(val_ce(x, y, amp=True) for x, y in batches) / a.nb
print(f"[fwd] valCE fp32 {ce32:.4f} | bf16-amp {ce16:.4f} (Δ{ce16-ce32:+.4f})", flush=True)
gref = [[g.cpu() for g in bp_grad(x, y)] for x, y in batches] # BP fp32 = truth
torch.cuda.empty_cache()
cb = [cos([g.cpu() for g in bp_grad(x, y, amp=True)], gref[i]) for i, (x, y) in enumerate(batches)]
print(f"[bp_amp] cos(BP_amp, BP_fp32) {sum(cb)/a.nb:.4f} <- amp-hurts-everyone baseline", flush=True)
torch.cuda.empty_cache()
for mode, beta in [('off', 1e-3), ('all', 1e-3), ('all', 3e-3), ('all', 1e-2), ('last', 1e-3), ('last', 3e-3)]:
cs = []
for i, (x, y) in enumerate(batches):
ge = ep_grad(x, y, beta, amp=mode)
cs.append(cos([g.cpu() for g in ge], gref[i]))
del ge; torch.cuda.empty_cache()
tag = {'off': 'fp32', 'all': 'amp_all', 'last': 'amp_last'}[mode]
print(f"[{tag} beta={beta:g}] cos(EP, BP_fp32) {sum(cs)/a.nb:.4f}", flush=True)
print("DONE_AMPGATE", flush=True)
|
