diff options
Diffstat (limited to 'ep_run/lt_ep_train.py')
| -rw-r--r-- | ep_run/lt_ep_train.py | 18 |
1 files changed, 13 insertions, 5 deletions
diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py index 99a1811..d56af1d 100644 --- a/ep_run/lt_ep_train.py +++ b/ep_run/lt_ep_train.py @@ -407,6 +407,8 @@ def main(): ap.add_argument('--floss_q', type=int, default=10) # unroll horizon (steps past z_T1, with graph, sub-batch) ap.add_argument('--floss_rho', type=float, default=0.995) # per-step contraction target (one-sided; matches eig_margin) ap.add_argument('--floss_bsub', type=int, default=4) # sub-batch rows for the graphed unroll (memory) + ap.add_argument('--reg_delay', type=int, default=0) # reg-free first N steps (resreg/jacreg/floss/adaptc all off) — magic-s2000 test + ap.add_argument('--noadaptc', action='store_true') # kill the hidden jacreg==0 damping feedback (clean single-reg arms) ap.add_argument('--diag_cos', type=int, default=0) # #1: every N steps, log cos(EP grad, exact BPTT grad) + res ap.add_argument('--fingerprint', action='store_true') # load --init_ckpt, print (res,cos,abscissa,val) fingerprint, exit ap.add_argument('--opt', choices=['adamw', 'lion', 'lionlars', 'sgdm', 'sgdsai'], default='adamw') @@ -558,12 +560,18 @@ def main(): idx, y = get_batch('train', cfg.B, cfg.T) if cfg.mode == 'ep': sw = hw_swap() if hw_on else None - grads, res = ep_step(blk, idx, y, cfg.T1, cfg.T2, cfg.eps, cfg.beta, jr, cfg.holo, cfg.hr, - cfg.t1max, cfg.res_est, cfg.t2sel, cfg.corr_every, cfg.res_gate, cfg.resreg, - cfg.eigreg, cfg.eig_margin, cfg.floss, cfg.floss_q, cfg.floss_rho, cfg.floss_bsub) + dly = step < cfg.reg_delay # reg-free early phase (magic-ckpt hypothesis: reach the edge unleashed) + grads, res = ep_step(blk, idx, y, cfg.T1, cfg.T2, cfg.eps, cfg.beta, 0.0 if dly else jr, + cfg.holo, cfg.hr, + cfg.t1max, cfg.res_est, cfg.t2sel, cfg.corr_every, cfg.res_gate, + 0.0 if dly else cfg.resreg, + cfg.eigreg, cfg.eig_margin, 0.0 if dly else cfg.floss, + cfg.floss_q, cfg.floss_rho, cfg.floss_bsub) if sw is not None: hw_restore(sw) - if cfg.jacreg > 0: # continuous controller: drive residual -> res_target (smooth) + if dly: # delay phase: no controller, no damping feedback either + pass + elif cfg.jacreg > 0: # continuous controller: drive residual -> res_target (smooth) flo = cfg.jacreg if cfg.jr_floor is None else cfg.jr_floor if cfg.jr_lrcouple: flo *= sched.get_last_lr()[0] / cfg.lr @@ -573,7 +581,7 @@ def main(): rtgt = math.exp((1 - u) * math.log(cfg.res_target) + u * math.log(cfg.rt_final)) rs = res if rs is None else cfg.res_ema * rs + (1 - cfg.res_ema) * res jr = min(cfg.jr_max, max(flo, jr * math.exp(0.3 * math.log((rs + 1e-9) / rtgt)))) - else: # damping feedback (no jacreg) + elif not cfg.noadaptc: # damping feedback (no jacreg) — DISABLE for clean single-reg ablations if res > 1e-3: blk.c = min(cfg.ccap, blk.c * 1.3) elif res < 2e-4: |