diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_eq_train.py | 16 |
1 files changed, 11 insertions, 5 deletions
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): |