summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-07-18 23:04:14 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-07-18 23:04:14 -0500
commit71096d092978f7e2b7466b4a98dbc06ef67aabbd (patch)
treee9b9c914e5dc3d4a3fdb84fff502a54345ec7bb1
parentc97c991a4c338026cc77f91d1298e7d6713ca453 (diff)
RESULT 49: 回路解剖=23段乘积近临界(下行×72/上行×13k/顶读1e-6),plain只热b8一段(h6 logit 94 vs cent 75)边际越线,灾变=硬底响应打伤b1(入口vjp 1.3→14.9); batch噪声对照0.39/0.45=Q1封口; --logit_knee 80实装+knee臂排c190后
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
-rw-r--r--docs/campaign/CASCADE_ABLATION_PLAN.md22
-rw-r--r--ep_run/casc_eq_train.py16
-rw-r--r--ep_run/probe_cycle.py198
3 files changed, 235 insertions, 1 deletions
diff --git a/docs/campaign/CASCADE_ABLATION_PLAN.md b/docs/campaign/CASCADE_ABLATION_PLAN.md
index 268ab91..f4471d6 100644
--- a/docs/campaign/CASCADE_ABLATION_PLAN.md
+++ b/docs/campaign/CASCADE_ABLATION_PLAN.md
@@ -568,6 +568,28 @@ direction), not training-under-fault; wave-2 = co-training with faults injected
+### RESULT 49 (2026-07-18): CYCLE-GAIN DECOMPOSITION — the explosion anatomized (user
+directive). The loop is a ~23-stage product sitting near-critical by construction; plain heats
+ONE stage (b8's sharp head h6, logit 94 vs cent 75) until a MARGINAL crossing; the catastrophe
+was the governor's hard-bottom RESPONSE, which damaged b1 (vjp-entry gain 1.3 -> 14.9).
+probe_cycle (mode-aligned per-stage gains, sweeps 15/16 diffs, beta=3e-3):
+ structure (BOTH lineages): down-chain vjp product ~50-72x, up-chain rebuild ~13-15k x,
+ top read gH ~1e-6 (beta*NBT*H_top) -> round-trip O(1). Biggest single amplifier = b8:
+ 2.9 (vjp) x 3.8 (fwd) ~ 11x round-trip; b7 next. Same map as R44/R47.
+ plain vs cent @195k: same structure ~1.2-1.4x hotter, concentrated at the b8 stage
+ (2.92 vs 2.50) and its carrier head b8-h6 (share 0.21, max|logit| 94 vs cent 75, ent 0.4);
+ healthy runs CARRY 50-57-logit heads fine (cent b10-h7=57, b11-h0=51) -> threshold ~80+.
+ plain @210k (sick): gV entry-stage 0<-1 EXPLODED 1.31 -> 14.91 (b1 sigma(J) 16->34->44 in
+ statej) = the noise-walk DAMAGE, not the original cause. Pre-collapse growth was modest
+ (res0 +55% over 40k steps); G's 300x conflates marginal crossing with post-crossing damage.
+BP-batch-noise control: per-batch BP u1 vs mean-BP u1 overlap 0.388/0.446 == EP arms' 0.37-0.41
+-> the "shared noise" is minibatch sampling, identical for BP; Q1 sealed.
+FIX ARM ARMED: --logit_knee 80 (piecewise: identity below 80, slope 0.2 above; healthy heads
+untouched incl cent's own carriers; cuts the runaway head's switching Jacobian 5x).
+fw72m_plain_knee = exact plain_nf replay + knee, queued behind c190 (smoke then DDP launch).
+Also on the menu (user decision): est_auto = centered-on-demand only when rho_ema near cap
+(pays 1.39x on ~10-20% of steps, avg ~1.05x) — mechanism-blind but c190-proven lever.
+
### RESULT 48c (2026-07-18): WHY the digital floor can't be BBP — the THETA-READ HAS NO 1/beta
CHANNEL, structurally. The BBP transition lives exactly where readout is DIFFERENTIAL (analog).
probe_bbp3 (6-bit STOCHASTIC rounding, fresh draw per estimate; frozen = one draw for free+
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index cca37b5..c1593f3 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -68,6 +68,11 @@ ap.add_argument('--drift_adapt', type=float, default=0.0) # >0: adaptive drift
ap.add_argument('--beta_cap_rho', type=float, default=0.0) # >0: LOOP-GAIN CAP on beta — if per-sweep residual
# ratio rho^ exceeds this, bscale *= 0.8 (beta backs off
# under the wall-2 ceiling); recovers x1.02 when rho^ low
+ap.add_argument('--logit_knee', type=float, default=0.0) # >0: piecewise clamp of attn logits above the
+ # knee (slope knee_slope) — cuts the switching-regime
+ # Jacobian of runaway sharp heads (b8-class carriers)
+ # without touching healthy logits below the knee
+ap.add_argument('--knee_slope', type=float, default=0.2)
ap.add_argument('--cap_floor', type=float, default=0.05) # hard bottom of the rho-cap; 0 = pure ceiling-tracking
# (cap follows the measured ceiling all the way down; a
# pinned bottom above the true ceiling = disguised wall-2)
@@ -211,6 +216,8 @@ class Olmo2Attn(nn.Module):
fr = torch.outer(torch.arange(T).float(), inv)
self.register_buffer('rc', fr.cos(), persistent=False)
self.register_buffer('rs', fr.sin(), persistent=False)
+ if args.logit_knee > 0:
+ self.register_buffer('cmask', torch.ones(T, T, dtype=torch.bool).tril(), persistent=False)
def rope(self, x):
x1, x2 = x[..., ::2], x[..., 1::2]
c, s = self.rc[None, None], self.rs[None, None]
@@ -222,7 +229,14 @@ class Olmo2Attn(nn.Module):
q = self.rope(q.view(B, T, self.H, self.hd).transpose(1, 2))
k = self.rope(k.view(B, T, self.H, self.hd).transpose(1, 2))
v = v.view(B, T, self.H, self.hd).transpose(1, 2)
- y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+ if args.logit_knee > 0:
+ kn = args.logit_knee
+ lg = (q @ k.transpose(-2, -1)) * (self.hd ** -0.5)
+ lg = torch.where(lg > kn, kn + args.knee_slope * (lg - kn), lg)
+ lg = lg.masked_fill(~self.cmask[:T, :T], float('-inf'))
+ y = lg.softmax(-1) @ v
+ else:
+ y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
return self.proj(y.transpose(1, 2).contiguous().view(B, T, C))
class Olmo2Block(nn.Module):
diff --git a/ep_run/probe_cycle.py b/ep_run/probe_cycle.py
new file mode 100644
index 0000000..fb06069
--- /dev/null
+++ b/ep_run/probe_cycle.py
@@ -0,0 +1,198 @@
+"""Cycle-gain decomposition of the nudge loop at the DOMINANT MODE (user directive: measure
+the explosion's cause). Run the trainer-faithful nudged relax, let the residual converge to the
+dominant mode, then difference consecutive sweeps to get per-stage amplification factors ALONG
+THE ACTUAL MODE (norms audits failed; alignment is everything):
+ gH = |dd_11| / |dz_11| top stage: beta*NBT*CE-curvature read
+ gV(l) = |dd_l| / |dd_{l+1}| down-chain vjp through block l+1
+ gF(l) = |do_l| / |dz_{l-1}| up-chain rebuild through block l
+ closure: gH * prod(gV) ~ dd-chain, rebuild adds dd_l -> next dz; rho_meas from residuals.
+Plus per-head carrier analysis for blocks 8-11 (share of the mode through each head, its
+max|logit| and attention entropy) -> names the physical carrier (switching-regime heads?).
+Control column for the batch-noise question: per-batch BP u1 vs 8-batch-mean BP u1 overlap."""
+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:195000,cent:195000,plain:210000')
+ap.add_argument('--K', type=int, default=16)
+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 heads(self, x):
+ """per-head attention outputs (pre-proj) + logit stats + entropy"""
+ 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)
+ lg = (q @ k.transpose(-2, -1)) / (self.hd ** 0.5)
+ mask = torch.ones(Tn, Tn, dtype=torch.bool, device=x.device).tril()
+ lgm = lg.masked_fill(~mask, float('-inf'))
+ p = lgm.softmax(-1)
+ y = p @ v # (B,H,T,hd)
+ ent = -(p.clamp_min(1e-12).log() * p).sum(-1).mean(dim=(0, 2)) # (H,)
+ mx = lg.masked_fill(~mask, 0).abs().amax(dim=(0, 2, 3)) # (H,)
+ return y, mx, ent
+ 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)
+ ck = torch.load(f'runs/fw72m_{lineage}_s{step}.pt', map_location='cpu', weights_only=False)
+ 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()
+
+ 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
+ hist = [] # (zs_copy, d_copy, os_copy) per sweep
+ res_seq = []
+ 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, os_now, rnum, rden = [], [], [], 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]
+ os_now.append(o.detach().float())
+ rnum += float((znew - zs[l]).norm()); rden += float(zs[l].norm())
+ zs[l] = znew
+ prev = zs[l]
+ ins, outs = n_ins, n_outs
+ res_seq.append(rnum / max(rden, 1e-9))
+ hist.append(([z.clone() for z in zs], [di.clone() for di in d], os_now))
+ rho = res_seq[-1] / max(res_seq[-2], 1e-12)
+ zA, dA, oA = hist[-2]; zB, dB, oB = hist[-1]
+ dz = [zB[l] - zA[l] for l in range(L)]
+ dd = [dB[l] - dA[l] for l in range(L)]
+ do = [oB[l] - oA[l] for l in range(L)]
+ gH = float(dd[L - 1].norm() / max(dz[L - 1].norm(), 1e-12))
+ gV = [float(dd[l].norm() / max(dd[l + 1].norm(), 1e-12)) for l in range(L - 1)]
+ gF = [float(do[l].norm() / max(dz[l - 1].norm(), 1e-12)) for l in range(1, L)]
+ print(f'{lineage} s{step//1000}k | rho {rho:.4f} | gH(top read) {gH:.4f}', flush=True)
+ print(' gV vjp l<-l+1 (11<-top .. 0<-1): ' + ' '.join(f'{v:.2f}' for v in reversed(gV)), flush=True)
+ print(' gF fwd l-1->l (1 .. 11): ' + ' '.join(f'{v:.2f}' for v in gF), flush=True)
+ # per-head carriers, blocks 8-11: mode share through each head + logit/entropy state
+ for l in range(8, 12):
+ at = blocks[l].attn
+ xin = zA[l - 1]
+ with torch.no_grad():
+ y0, mx, ent = at.heads(xin)
+ y1, _, _ = at.heads(xin + dz[l - 1])
+ share = torch.linalg.vector_norm(y1 - y0, dim=(0, 2, 3))
+ share = (share / share.sum()).cpu().numpy()
+ order = np.argsort(-share)[:3]
+ cells = ' '.join(f'h{h}:share {share[h]:.2f} maxlg {float(mx[h]):.0f} ent {float(ent[h]):.2f}'
+ for h in order)
+ print(f' blk{l} carriers: {cells}', flush=True)
+ del tok, blocks, W_out, ln_f, ins, outs, f_ins, f_outs, zs, d, hist
+ torch.cuda.empty_cache()
+
+# batch-noise control (the "shared with BP" question): per-batch BP u1 vs mean-BP u1
+ck = torch.load('runs/fw72m_plain_s195000.pt', map_location='cpu', 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'])
+params = list(blocks.parameters())
+names = [n for n, _ in blocks.named_parameters()]
+SEL = [i for i, n in enumerate(names) if n in ('0.attn.qkv.weight', '8.attn.qkv.weight')]
+gs = {i: [] for i in SEL}
+for _ in range(8):
+ ixb = torch.randint(len(data) - T - 1, (B,))
+ xb = torch.stack([torch.from_numpy(data[i:i + T].astype(np.int64)) for i in ixb]).to(dev)
+ yb = torch.stack([torch.from_numpy(data[i + 1:i + 1 + T].astype(np.int64)) for i in ixb]).to(dev)
+ z = tok(xb)
+ for b in blocks: z = b(z)
+ ce = F.cross_entropy((ln_f(z) @ W_out.t()).reshape(-1, vocab), yb.reshape(-1))
+ g = torch.autograd.grad(ce, params, allow_unused=True)
+ for i in SEL: gs[i].append(g[i].detach().cpu())
+for i in SEL:
+ gbar = torch.stack(gs[i]).mean(0)
+ u = torch.linalg.svd(gbar, full_matrices=False).U[:, 0]
+ ovs = [abs(float(u @ torch.linalg.svd(g, full_matrices=False).U[:, 0])) for g in gs[i]]
+ print(f'BP-batch-noise control {names[i]}: per-batch u1 vs mean u1 overlap = {np.mean(ovs):.3f}', flush=True)
+print('CYCLE_DONE', flush=True)