diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 24 |
1 files changed, 16 insertions, 8 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 4eba267..f1f26a4 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -27,7 +27,8 @@ ap.add_argument('--tok_init', type=float, default=0.0) # >0: init tok/pos with ap.add_argument('--compile', action='store_true') # torch.compile each block (free speed where supported) ap.add_argument('--sig_every', type=int, default=25) # tok-sigma refresh interval (amortized) ap.add_argument('--dtop_every', type=int, default=1) # 1 = exact (DEFAULT, BP-parity); 2 = fast mode (~20% cheaper, ~4% CE tax at high lr) -ap.add_argument('--gate_every', type=int, default=200) # in-training cos(EP,BP) telemetry +ap.add_argument('--gate_every', type=int, default=200) # in-training cos(EP,BP) telemetry; <=0 = fully BP-free (no bp_gate at all) +ap.add_argument('--gate_govern', action='store_true') # let gate cos adjust K/bscale (default: observe-only => training control is BP-free) args = ap.parse_args() torch.manual_seed(args.seed) dev = 'cuda' if torch.cuda.is_available() else 'cpu' @@ -129,14 +130,20 @@ def relax(z0, zs, ins, outs, y, beta, K, x): return zs, outs def dFdtheta(zs, x, y, beta): - """dF/dtheta at fixed relaxed states (z0 rebuilt WITH graph so emb gets its E-path grad).""" + """theta-readout at FIXED states. Not used by the training loop (relax reuses its own + graphs); kept as the INVARIANT-TEST surface for test_bp_free.py. Self-sealing: inputs + are detached here so the local-graph property holds for any caller.""" + zs = [z.detach() for z in zs] prev = tok(x) + pos(torch.arange(args.T, device=dev))[None] E = 0.0 - for z, b in zip(zs, blocks): E = E + 0.5 * ((z - b(prev, mask)) ** 2).sum(); prev = z + for z, b in zip(zs, blocks): + E = E + 0.5 * ((z - b(prev, mask)) ** 2).sum() + prev = z # zs detached at entry => blocks l>0 get detached inputs; block 0 gets the graphed emb obj = E / NBT + beta * F.cross_entropy(readout(zs[-1]).reshape(-1, vocab), y.reshape(-1)) gs = torch.autograd.grad(obj, all_params, allow_unused=True) return [g if g is not None else None for g in gs] + SIG0 = None GOV = {'K': None, 'bscale': 1.0, 'gema': None, 'drift': 0.0, 'gn': 0.0, 'sig': 0.0} def ep_step(x, y): @@ -215,17 +222,18 @@ for step in range(args.steps + 1): ce, beta_t, rounds, ok = ep_step(x, y) if not ok: skips += 1 gcos = float('nan') - if step % args.gate_every == 0 and ok: + if args.gate_every > 0 and step % args.gate_every == 0 and ok: gbp = bp_gate(x, y) num = den1 = den2 = 0.0 for p, g in zip(all_params, gbp): if p.grad is None or g is None: continue num += float((p.grad * g).sum()); den1 += float((p.grad ** 2).sum()); den2 += float((g ** 2).sum()) gcos = num / max((den1 ** 0.5) * (den2 ** 0.5), 1e-12) - if gcos < 0.97: # estimator governor: spend more - GOV['K'] = min(GOV['K'] + 2, args.kmax); GOV['bscale'] = max(GOV['bscale'] * 0.7, 0.05) - elif gcos > 0.995 and GOV['K'] > args.K: # relax back when quality is abundant - GOV['K'] -= 1; GOV['bscale'] = min(GOV['bscale'] * 1.05, 1.0) + if args.gate_govern: # opt-in: BP-informed control flow + if gcos < 0.97: + GOV['K'] = min(GOV['K'] + 2, args.kmax); GOV['bscale'] = max(GOV['bscale'] * 0.7, 0.05) + elif gcos > 0.995 and GOV['K'] > args.K: + GOV['K'] -= 1; GOV['bscale'] = min(GOV['bscale'] * 1.05, 1.0) torch.nn.utils.clip_grad_norm_(all_params, 1.0) opt.step(); sched.step(); opt.zero_grad(set_to_none=True) if step % args.log == 0: |