diff options
| -rw-r--r-- | ep_run/casc_eq_train.py | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 38b687e..10206f6 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -24,6 +24,7 @@ ap.add_argument('--untie', action='store_true') # separate readout matri ap.add_argument('--tok_init', type=float, default=0.0) # >0: init tok/pos with this std (GPT-standard 0.02) 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 args = ap.parse_args() torch.manual_seed(args.seed) @@ -100,7 +101,7 @@ def relax(z0, zs, ins, outs, y, beta, K, x): reuses them instead of re-running a full graphed chain.""" d = [None] * args.L for k in range(K): - if k % 2 == 0 or d[args.L - 1] is None: + if k % args.dtop_every == 0 or d[args.L - 1] is None: zc = zs[args.L - 1].detach().requires_grad_(True) ce = F.cross_entropy(readout(zc).reshape(-1, vocab), y.reshape(-1)) d[args.L - 1] = (-beta * NBT * torch.autograd.grad(ce, zc)[0]).detach() |