diff options
| -rw-r--r-- | docs/campaign/CASCADE_ABLATION_PLAN.md | 24 | ||||
| -rw-r--r-- | ep_run/casc_eq_train.py | 19 | ||||
| -rw-r--r-- | ep_run/probe_cx3_rhogrid.py | 133 |
3 files changed, 173 insertions, 3 deletions
diff --git a/docs/campaign/CASCADE_ABLATION_PLAN.md b/docs/campaign/CASCADE_ABLATION_PLAN.md index f0a82be..6dd93e5 100644 --- a/docs/campaign/CASCADE_ABLATION_PLAN.md +++ b/docs/campaign/CASCADE_ABLATION_PLAN.md @@ -568,6 +568,30 @@ direction), not training-under-fault; wave-2 = co-training with faults injected +### RESULT 54 (2026-07-20): THE CEILING SINKS ∝1/σ² AS TRAINING SHARPENS — no fixed β survives; +"fixed β never explodes" was FALSE (it STALLS: endgame skip-storm = frozen training). β must DESCEND. +User caught it: "固定β永不炸个屁！你一直skip还训不训了？" — dead right. +EVIDENCE: plain2 fixed-3e-3 full run — best 3.3437 hit at step 199800, then skip rate SURGES: +84%/step @224k, 89%/step @231k. Endgame (200-234k) = 34k steps of ~zero effective updates; +the "234k crown 3.3437" is really a 200k result. skip ≠ safe, skip = STALL. +DENSE ρ-PROBE (probe_cx3_rhogrid, K=30 asymptotic — corrects Codex's K=8 β*=0.7 which mistook +slow divergence for convergence): true ceiling β* (max β whose resK reaches the fp floor) SINKS: + s35000 β*≈0.12 (σ242) | s95000 β*≈0.12 | s150000 β*≈0.08 (σ443) | s230000 β*<0.03 (σ473). +MECHANISM: training sharpens the model (σ 242→473, ~2×); a sharper block Jacobian makes the nudge +relaxation harder to converge → ceiling drops. It's ∝1/σ^~2. The sunk endgame ceiling (<0.03, +approaching the 3e-3 training β) is WHY fixed 3e-3 skip-storms late. +CODE: σ-scaling EXISTS (line 473: β=β0·SIG0²/σ²) — the ORIGINAL wall-1 design already tracks σ — +but line 477 floor=3e-3 PINS it (plain2 wanted β↓1.94e-3 @σ473, floor forced 3e-3). My beta_simple +"full ownership" that DELETED σ-scaling was exactly backwards. But naive un-flooring fails too: +early σ 3.9→242 (62×, network growing structure NOT sharpening) would crush β∝1/σ² to 8e-7 and +starve training — which is WHY the floor existed. The real tension: early σ-growth (structure, +don't cut β) vs late σ-growth (sharpening, DO cut β); σ magnitude can't separate them. +FIX (least-assumption): SCHEDULE β descent by progress like LR — --beta_cos_min added (cosine +β from --beta to min, bypasses σ-scaling+floor). Calibrated to the measured ceiling: early 2e-2 +(<0.12), endgame 8e-4 (<sunk ceiling). Conservative=safe (CE flat in-corridor → undershoot free, +overshoot skips). fw72m_betacos (2e-2→8e-4, 234k, guards silent) IN FLIGHT vs plain2 3.3437 — +decisive test = does endgame KEEP LEARNING instead of skip-stalling. + ### RESULT 53 (2026-07-20): THE β DOGMA REBUILT — ceiling is REAL (β*≈0.7 @s35000) but the v1 collapse did NOT cross it; fixed β never explodes (only the controller does); policy = FIXED SMALL β. Codex ρ-probe (probe_cx2_rho, s35000, one relax, no training): ρ = res_k/res_{k-1} vs β: diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index c77b92c..86a8e06 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -29,6 +29,10 @@ ap.add_argument('--compile', action='store_true') # torch.compile each blo ap.add_argument('--sig_every', type=int, default=25) # tok-sigma refresh interval (amortized) ap.add_argument('--beta_floor', type=float, default=0.0) # >0: floor beta_t (anti finite-beta SNR collapse at depth) ap.add_argument('--beta_fixed', action='store_true') # disable sig^2 schedule, hold beta_t = args.beta constant +ap.add_argument('--beta_cos_min', type=float, default=0.0) # >0: cosine-descend beta from --beta to this over + # --steps (tracks the sinking ceiling by progress; + # bypasses sigma-scaling AND floor). Set --beta to the + # early value (below early ceiling ~0.12). ap.add_argument('--cosine', action='store_true') # warmup then cosine decay to lr_min_ratio*lr over --steps (long runs) ap.add_argument('--lr_min_ratio', type=float, default=0.1) ap.add_argument('--qk_norm', action='store_true') # RMS-norm q,k per head before scores (OLMo2-style; bounds logits, analog-friendly) @@ -472,9 +476,18 @@ def ep_step(x, y): if SIG0 is None: SIG0 = args.sig0 if args.sig0 > 0 else sig beta_t = args.beta * GOV['bscale'] * (SIG0 * SIG0) / max(sig * sig, 1e-9) if args.beta_fixed: beta_t = args.beta * GOV['bscale'] - fl = args.beta_floor - if args.bf_late > 0.0 and GOV.get('step', 0) >= args.bf_late_at: fl = args.bf_late - if fl > 0.0: beta_t = max(beta_t, fl) + if args.beta_cos_min > 0: + # SCHEDULED beta descent (07-20): the ceiling sinks ~1/sigma^2 as training sharpens the + # model (sigma 242->473); no fixed beta stays under it (endgame skip-stall). Descend beta + # by progress like LR — early large (below the high early ceiling), late small (below the + # sunk endgame ceiling). Skips sigma-scaling AND floor entirely. Conservative = safe: CE + # is flat across the in-corridor band, so undershoot costs nothing, overshoot skips. + prog = min(GOV.get('step', 0) / max(args.steps, 1), 1.0) + beta_t = args.beta_cos_min + 0.5 * (args.beta - args.beta_cos_min) * (1 + math.cos(math.pi * prog)) + else: + fl = args.beta_floor + if args.bf_late > 0.0 and GOV.get('step', 0) >= args.bf_late_at: fl = args.bf_late + if fl > 0.0: beta_t = max(beta_t, fl) if args.beta_ride > 1.0: # ride-v2(a): floor jumps must not compose with a pre-charged cap — rescale cap so the # EFFECTIVE beta is continuous across any floor change (the 0.09-at-20k bug, RESULT 37) diff --git a/ep_run/probe_cx3_rhogrid.py b/ep_run/probe_cx3_rhogrid.py new file mode 100644 index 0000000..7c27a2d --- /dev/null +++ b/ep_run/probe_cx3_rhogrid.py @@ -0,0 +1,133 @@ +"""Dense rho(beta) trend probe (user: measure more divergence points, see the shape). +Trainer-faithful nudged relax (matches probe_rhorelax.py, validated to reproduce GOV meter), +K=30 sweeps to read the ASYMPTOTIC rho (not the 8-sweep transient), dense beta grid through and +past the ceiling, on MULTIPLE ckpts to see the ceiling sink with training. Reports per (ckpt,beta): +res0 (drive, should be proportional to beta if the loop is linear), asymptotic rho (tail-median of +res ratios), and the divergence verdict. GPU, read-only, no training.""" +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='fw72m_plain2:35000,fw72m_plain2:95000,fw72m_plain2:150000,fw72m_plain2:230000') +ap.add_argument('--betas', default='0.03,0.06,0.1,0.15,0.2,0.3,0.4,0.5,0.6,0.7,0.8,1.0,1.3,1.7,2.5') +ap.add_argument('--K', type=int, default=30) +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)) + +betas = [float(s) for s in a.betas.split(',')] +first = True +for spec in a.ckpts.split(','): + tag, step = spec.split(':'); step = int(step) + p = f'runs/{tag}_s{step}.pt' + try: + ck = torch.load(p, map_location=dev, weights_only=False) + except FileNotFoundError: + print(f'{tag} 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() + print(f'\n=== {tag} s{step} (K={a.K} sweeps) ===', flush=True) + print(f'{"beta":>7} {"res0":>10} {"res1":>10} {"resK":>11} {"rho_tail":>9} {"verdict":>10}', flush=True) + bstar = None + for beta in betas: + 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 + res_list = [] + 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)) + 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], + retain_graph=True)[0].detach().float() + prev = tok(x).detach(); n_ins, n_outs = [], []; rnum = rden = 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] + rnum += float((znew - zs[l]).norm()); rden += float(zs[l].norm()) + zs[l] = znew; prev = zs[l] + ins, outs = n_ins, n_outs + res_list.append(rnum / max(rden, 1e-9)) + if not np.isfinite(res_list[-1]) or res_list[-1] > 1e4: break + ratios = [res_list[i] / res_list[i - 1] for i in range(1, len(res_list)) + if res_list[i - 1] > 1e-7] + tail = ratios[-6:] if len(ratios) >= 6 else ratios + rho_tail = float(np.median(tail)) if tail else float('nan') + diverged = (not np.isfinite(res_list[-1])) or res_list[-1] > 1e-2 or rho_tail > 1.0 + verdict = 'DIVERGE' if diverged else 'converge' + if diverged and bstar is None: bstar = beta + print(f'{beta:>7.3f} {res_list[0]:>10.2e} {(res_list[1] if len(res_list)>1 else float("nan")):>10.2e} ' + f'{res_list[-1]:>11.2e} {rho_tail:>9.4f} {verdict:>10}', flush=True) + print(f' -> ceiling beta* (first DIVERGE) = {bstar}', flush=True) + del tok, blocks, W_out, ln_f; torch.cuda.empty_cache() +print('\nRHOGRID_DONE', flush=True) |