From 7ef12f8006ea1b19657fab6aa189785929b1f4a4 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Thu, 30 Jul 2026 15:50:25 -0500 Subject: =?UTF-8?q?RESULT=2076:=20=E8=B0=B1=E4=BB=AA=E8=AF=BB=E5=87=BA=20?= =?UTF-8?q?=E2=80=94=20=E9=80=90=E5=9D=97=E6=B3=84=E6=BC=8F=E5=89=96?= =?UTF-8?q?=E9=9D=A2(C768=E4=B8=8A6=E5=9D=97-0.70..-0.23,C512=E5=9D=978-11?= =?UTF-8?q?=E4=B9=9F=E7=97=85)+=E5=81=B6=E6=80=A7=E5=AE=9E=E9=94=A4(-?= =?UTF-8?q?=CE=B2=E5=90=8C=E5=8F=B7)+disp=5Fratio=E9=98=88=E5=80=BC?= =?UTF-8?q?=E8=B7=A8=E5=AE=BD=E5=BA=A6=E5=A1=8C=E7=BC=A9(~2e-7=E2=89=882?= =?UTF-8?q?=C3=97fp32=20eps)+fp32=E5=9C=B0=E6=9D=BF=E5=81=87=E8=AF=B4(?= =?UTF-8?q?=E4=B8=80=E4=B8=BE=E8=A7=A3=E9=87=8A=E5=85=A8=E9=83=A8=E7=BA=A6?= =?UTF-8?q?=E6=9D=9F);=20=E9=A2=84=E6=B3=A8=E5=86=8C=E5=8F=8C=E8=AF=81?= =?UTF-8?q?=E4=BC=AA=E5=B7=B2=E5=8F=91=E5=B0=84(=CE=B2-scan=E5=B9=B3?= =?UTF-8?q?=E7=A7=BB+fp64=E5=86=B3=E5=AE=9A=E8=87=82);=20--probe=5Ff64?= =?UTF-8?q?=E5=AE=9E=E8=A3=85?= 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 | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) (limited to 'ep_run/casc_eq_train.py') diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index a9218a0..76b96fd 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -54,7 +54,8 @@ ap.add_argument('--dgain_geo', type=float, default=0.0) # >0: per-layer geomet ap.add_argument('--dgain_geo_cap', type=float, default=0.0) # >0: cap for the geometric profile ap.add_argument('--dgain_rand', type=float, default=0.0) ap.add_argument('--probe_dgspec', type=int, default=0) # >0: M1 spectroscopy, value = n batches; exits before training -ap.add_argument('--probe_gains', default='1,2,4,8,16,32,64,128,256') # >1: per-STEP log-uniform dgain_top in +ap.add_argument('--probe_gains', default='1,2,4,8,16,32,64,128,256') +ap.add_argument('--probe_f64', action='store_true') # fp64 states+model in the probe: the fp-floor decisive arm # >1: per-STEP log-uniform dgain_top in # [1, this] (spread-spectrum probing of the # decade-spread threshold distribution) ap.add_argument('--dgain_top', type=float, default=1.0) # amplify d in STATE FORMATION for blocks @@ -291,6 +292,7 @@ class Olmo2Block(nn.Module): z = z + self.na(self.attn(z)) return z + self.nf(self.ff(z)) +SDT = torch.float64 if args.probe_f64 else torch.float32 # state dtype; fp64 only in probe mode tok = nn.Embedding(vocab, args.C).to(dev) pos = nn.Embedding(args.T, args.C).to(dev) if args.tok_init > 0: @@ -375,7 +377,7 @@ def free_states_graphed(x): for b in blocks: i = prev.detach().requires_grad_(True) o = b(i, mask) - ins.append(i); outs.append(o); zs.append(o.detach().float()) + ins.append(i); outs.append(o); zs.append(o.detach().to(SDT)) prev = zs[-1] return z0, zs, ins, outs @@ -410,7 +412,7 @@ def relax(z0, zs, ins, outs, y, beta, K, x, bmask=None): if bmask is not None: g = g * bmask d[args.L - 1] = (-beta * nbt_loc * g).detach() for l in range(args.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() + d[l] = torch.autograd.grad(outs[l + 1], ins[l + 1], grad_outputs=d[l + 1].to(outs[l + 1].dtype))[0].detach().to(SDT) def rebuild(last): nonlocal ins, outs @@ -435,7 +437,7 @@ def relax(z0, zs, ins, outs, y, beta, K, x, bmask=None): if args.dgain_geo_cap > 0: _dg = min(_dg, args.dgain_geo_cap) else: _dg = args.dgain_all * ((GOV.get('dgcur') or args.dgain_top) if l >= args.L // 2 else 1.0) - znew = o.detach().float() + (_dg * d[l] if _dg != 1.0 else d[l]) + znew = o.detach().to(SDT) + (_dg * d[l] if _dg != 1.0 else d[l]) # damped (under-relaxed) mixing: geta<1 restores contraction on stiff operators # (wall-2 toolkit); fixed point unchanged (z = z + geta*(o+d-z) <=> z = o+d) mixed = znew if g_eff >= 1.0 else (zs[l] + g_eff * (znew - zs[l])) @@ -855,6 +857,10 @@ if args.probe_dgspec > 0: # splits odd-in-beta FD bias from even-in-|displacement| leak. Run WITHOUT --amp: the # bf16 floor is width-blind (b15 lesson); fp32 is the instrument-grade path. assert WORLD == 1, 'probe_dgspec is single-GPU' + if args.probe_f64: + tok.double(); blocks.double(); ln_f.double() + if W_out is not None: W_out.data = W_out.data.double() + print('[dgspec] FP64 states+model active', flush=True) import json, sys gains = [float(t) for t in args.probe_gains.split(',')] cfgs = [(f'g{g:g}', g, +1.0) for g in gains] + [('g1neg', 1.0, -1.0)] @@ -886,7 +892,7 @@ if args.probe_dgspec > 0: for l in range(args.L): disp[l] += float(GOV['_last_d'][l].norm() / max(float(zp[l].norm()), 1e-12)) / NB E = 0.0 - for z, o in zip(zp, lo): E = E + 0.5 * ((z.detach().float() - o.float()) ** 2).sum() + for z, o in zip(zp, lo): E = E + 0.5 * ((z.detach().to(SDT) - o.to(SDT)) ** 2).sum() obj = E / (NBT * bt) + obj_loss(readout(zp[-1].detach()).reshape(-1, vocab), y.reshape(-1)) gs = torch.autograd.grad(obj, all_params, allow_unused=True) for j, ix in enumerate(bix): -- cgit v1.2.3