From 6db80927045e6c8ca91df9dd505f07463cc3d73d Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Mon, 20 Jul 2026 06:48:29 -0500 Subject: =?UTF-8?q?RESULT=2054:=20=E4=B8=8A=E7=95=8C=E2=88=9D1/=CF=83?= =?UTF-8?q?=C2=B2=E4=B8=8B=E6=B2=89(=CF=83242=E2=86=92473=E7=BF=BB?= =?UTF-8?q?=E5=80=8D,=CE=B2*=E4=BB=8E0.12=E8=B7=8C=E5=88=B0<0.03),?= =?UTF-8?q?=E5=9B=BA=E5=AE=9A=CE=B2=3D=E5=81=9C=E6=BB=9E=E9=9D=9E=E5=AE=89?= =?UTF-8?q?=E5=85=A8(plain2=E6=94=B6=E5=AE=9889%skip,best=E5=85=B6?= =?UTF-8?q?=E5=AE=9E=E6=98=AF200k=E6=88=90=E7=BB=A9);=20K=3D30=E5=AF=86?= =?UTF-8?q?=E6=8E=A2=E9=92=88=E6=8E=A8=E7=BF=BBCodex=E7=9A=84=CE=B2*=3D0.7?= =?UTF-8?q?(K=3D8=E8=AF=AF=E5=88=A4=E6=85=A2=E5=8F=91=E6=95=A3=E4=B8=BA?= =?UTF-8?q?=E6=94=B6=E6=95=9B);=20=CF=83-scaling=E5=8E=9F=E8=AE=BE?= =?UTF-8?q?=E8=AE=A1=E5=AF=B9=E4=BD=86=E8=A2=ABfloor=E9=92=89=E6=AD=BB+?= =?UTF-8?q?=E6=97=A9=E6=9C=9F=CF=83=E6=9A=B4=E6=B6=A8=E5=BC=A0=E5=8A=9B;?= =?UTF-8?q?=20--beta=5Fcos=5Fmin=E5=AE=9E=E8=A3=85(=E6=8C=89=E8=BF=9B?= =?UTF-8?q?=E7=A8=8Bcosine=E9=99=8D=CE=B2=E8=B7=9F=E8=B8=AA=E4=B8=8B?= =?UTF-8?q?=E6=B2=89=E4=B8=8A=E7=95=8C);=20betacos=E7=8E=8B=E5=86=A02e-2?= =?UTF-8?q?=E2=86=928e-4=E5=9C=A8=E9=A3=9E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/casc_eq_train.py | 19 ++++++- ep_run/probe_cx3_rhogrid.py | 133 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 149 insertions(+), 3 deletions(-) create mode 100644 ep_run/probe_cx3_rhogrid.py (limited to 'ep_run') 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) -- cgit v1.2.3