lt_ep_train: optional W&B mirroring (--wandb/--wandb_run), failure-proof logging of val/ema/best/res/jr/lr/rho/gov state

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn

1 files changed, 35 insertions, 3 deletions

diff --git a/ep_run/lt_ep_train.py b/ep_run/lt_ep_train.py
index 2e555f9..1c52740 100644
--- a/ep_run/lt_ep_train.py
+++ b/ep_run/lt_ep_train.py
@@ -17,8 +17,11 @@ import argparse, math, pickle, time, json, os, numpy as np, torch, torch.nn.func
 from pathlib import Path
 
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
-DD = Path('/home/yurenh2/ept/ep_run/data/tinystories_bpe')
-vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size']
+DD = Path(os.environ.get('EPT_DATA', '/home/yurenh2/ept/ep_run/data/tinystories_bpe'))
+try:
+    vocab = pickle.load(open(DD / 'meta.pkl', 'rb'))['vocab_size']
+except FileNotFoundError:                                # non-NFS hosts (Delta): main() re-loads via --data
+    vocab = 4096
 
 
 def get_batch(split, B, T):
@@ -439,6 +442,8 @@ def main():
     ap.add_argument('--T1', type=int, default=80); ap.add_argument('--T2', type=int, default=15)
     ap.add_argument('--eps', type=float, default=0.1); ap.add_argument('--beta', type=float, default=0.02)
     ap.add_argument('--lr', type=float, default=1e-3); ap.add_argument('--log', type=int, default=100)
+    ap.add_argument('--wandb', type=str, default='')           # W&B project name ('' = off); run name = --wandb_run or ckpt stem
+    ap.add_argument('--wandb_run', type=str, default='')
     ap.add_argument('--warmup', type=int, default=0)           # linear lr warmup steps (big-model stability)
     ap.add_argument('--state', type=str, default='')           # periodic FULL-state path (weights+opt+sched+step)
     ap.add_argument('--resume', action='store_true')           # resume from --state if it exists (Colab timeouts)
@@ -496,6 +501,7 @@ def main():
     ap.add_argument('--gov_k', type=int, default=100)          # scout cadence (steps)
     ap.add_argument('--gov_scout', type=float, default=0.985)  # scout rho threshold to trigger ARPACK confirmation
     ap.add_argument('--gov_boost', type=float, default=3.0)    # resreg multiplier during rescue mode
+    ap.add_argument('--gov_keepjr', action='store_true')       # keep frozen jr floor DURING the free phase (redx-faithful; C1280 lesson: truly-free dies by step 800)
     ap.add_argument('--rt_final', type=float, default=0.0)     # anneal res_target to this (0=off), 25%-75% of run
     ap.add_argument('--nudge_brake', type=float, default=0.0)  # kappa: anchor spring during nudge (Tikhonov adjoint)
     ap.add_argument('--init_ckpt', type=str, default='')       # warm-start weights from a saved ckpt
@@ -639,6 +645,16 @@ def main():
         print(f"[fingerprint] ckpt={cfg.init_ckpt or 'scratch'} | res={fp['res']:.2e} cos(EP,BPTT)={fp['cos']:.4f} "
               f"rho={fp['rho']:.5f} Re_mu={fp['mu_re']:+.4f} val={fp['val']:.4f}", flush=True)
         return
+    wb = None
+    if cfg.wandb:                                        # optional W&B mirror of the log lines; never allowed to kill training
+        try:
+            import wandb as _wandb
+            wb = _wandb.init(project=cfg.wandb,
+                             name=cfg.wandb_run or (Path(cfg.ckpt).stem if cfg.ckpt else None),
+                             config=vars(cfg), resume='allow',
+                             id=(cfg.wandb_run or (Path(cfg.ckpt).stem if cfg.ckpt else None)))
+        except Exception as e:
+            print(f"[wandb] disabled ({e})", flush=True); wb = None
     gov = {'engaged': not cfg.governor, 'boost_until': 0, 'hi': 0, 'cache': {}}
     for step in range(start_step, cfg.steps + 1):
         if cfg.governor and step % cfg.gov_k == 0 and step > 0:
@@ -674,8 +690,9 @@ def main():
         if cfg.mode == 'ep':
             sw = hw_swap() if hw_on else None
             dly = (step < cfg.reg_delay) or (cfg.governor and not gov['engaged'])
+            _jr = jr if (cfg.gov_keepjr and not (step < cfg.reg_delay)) else (0.0 if dly else jr)
             _rr = cfg.resreg * (cfg.gov_boost if step < gov['boost_until'] else 1.0)
-            grads, res = ep_step(blk, idx, y, cfg.T1, cfg.T2, cfg.eps, cfg.beta, 0.0 if dly else jr,
+            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,
                                  0.0 if dly else _rr,
@@ -772,8 +789,23 @@ def main():
                             'step': step, 'best': best}, cfg.ckpt)
             ftag = f" rho={blk._floss_rho:.4f}" if cfg.floss > 0 and hasattr(blk, '_floss_rho') else ""
             print(f"step {step:4d}/{cfg.steps} | val CE {v:.4f}{etag} (best {best:.4f}) | jr={jr:.1f} res={res:.1e}{ftag} | {step/(time.time()-t0):.2f} it/s", flush=True)
+            if wb is not None:
+                try:
+                    rec = {'val_ce': v, 'best': best, 'res': res, 'jr': jr,
+                           'lr': opt.param_groups[0]['lr'], 'it_per_s': step / max(time.time() - t0, 1e-9)}
+                    if pema is not None: rec['ema_ce'] = ve
+                    if cfg.floss > 0 and hasattr(blk, '_floss_rho'): rec['rho'] = blk._floss_rho
+                    if cfg.governor: rec['gov_engaged'] = int(gov['engaged'])
+                    wb.log(rec, step=step)
+                except Exception:
+                    pass
             save_state(step)                       # full-state checkpoint each log interval (Colab resume)
     print(f"[{cfg.mode}] DONE best val CE {best:.4f}  (random baseline ln({vocab})={math.log(vocab):.3f})", flush=True)
+    if wb is not None:
+        try:
+            wb.summary['best_val_ce'] = best; wb.finish()
+        except Exception:
+            pass
     out_dir = Path('runs')
     out_dir.mkdir(exist_ok=True)
     json.dump({'mode': cfg.mode, 'best_val_ce': best}, open(out_dir / f'H2_{cfg.mode}.json', 'w'))