diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_eq_train.py | 20 |
1 files changed, 16 insertions, 4 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index f3b8736..5a2677c 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -37,6 +37,7 @@ ap.add_argument('--sig0', type=float, default=-1.0) # override SIG0 (beta-s ap.add_argument('--olmo2', action='store_true') # OLMo2-standard block: norm-AFTER-sublayer RMSNorm, full-width QK-norm, RoPE(500k), SwiGLU, no-bias, untied head, final RMSNorm, 0.02 init ap.add_argument('--wd', type=float, default=-1.0) # >=0: grouped weight decay (linear weights+head decay; embeddings/norm-gains none). <0 = legacy uniform 1e-4 ap.add_argument('--zloss', type=float, default=0.0) # z-loss coefficient on train objective (OLMo2-style logit regularizer); 0 = off +ap.add_argument('--kretry', type=int, default=0) # >0: on drift-reject, RETRY the batch once with this many fb rounds (diag B: K8 converges the marginal batches) instead of dropping it 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; <=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) @@ -290,9 +291,20 @@ def ep_step(x, y): drift = sum(float((a - b).norm()) for a, b in zip(zp, zs_free)) / max( sum(float(b.norm()) for b in zs_free), 1e-9) if (not math.isfinite(drift)) or (drift > 0.5 and not args.noguard): - GOV['skd'] = GOV.get('skd', 0) + 1 # drift-guard reject (relaxation non-convergence) - for p in all_params: p.grad = None - return free_ce, beta_t, GOV['K'], False + ok_retry = False + if args.kretry > 0 and math.isfinite(drift) and not args.noguard: + GOV['skr'] = GOV.get('skr', 0) + 1 # marginal batch: retry once with deeper relaxation + z0, zs, ins, outs = free_states_graphed(x) + zs_free = [z.clone() for z in zs] + zp, last_outs = relax(z0, zs, ins, outs, y, +beta_t, args.kretry, x) + with torch.no_grad(): + drift = sum(float((a - b).norm()) for a, b in zip(zp, zs_free)) / max( + sum(float(b.norm()) for b in zs_free), 1e-9) + ok_retry = math.isfinite(drift) and drift <= 0.5 + if not ok_retry: + GOV['skd'] = GOV.get('skd', 0) + 1 # drift-guard reject (relaxation non-convergence) + for p in all_params: p.grad = None + return free_ce, beta_t, GOV['K'], False GOV['drift'] = drift E = 0.0 for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach() - o) ** 2).sum() @@ -368,7 +380,7 @@ for step in range(start_step, args.steps + 1): val = evaluate(); best = min(best, val) gtag = '' if math.isnan(gcos) else f' cos={gcos:.4f}' print(f'step {step:5d}/{args.steps} | train {ce:.4f} val {val:.4f} (best {best:.4f}) ' - f'| beta={beta_t:.2e} K={rounds} skips={skips}(d{GOV.get("skd",0)}/g{GOV.get("skg",0)}){gtag} ' + f'| beta={beta_t:.2e} K={rounds} skips={skips}(d{GOV.get("skd",0)}/g{GOV.get("skg",0)}/r{GOV.get("skr",0)}){gtag} ' f'drift={GOV["drift"]:.3f} gn={GOV["gn"]:.2e} sig={GOV["sig"]:.1f} | {step/max(time.time()-t0,1e-9):.3f} it/s', flush=True) if wb is not None: try: wb.log({'train_ce': ce, 'val_ce': val, 'best': best, 'beta_t': beta_t, |