Stage-1 epoch blowup @12100: sig story REFUTED (only +8%, cos fine till after); leading indicator=drift-guard skips -> contractivity bifurcation in nudged relaxation (cascade Hopf wall); add resume/sig0/final_ln; A/B/C diagnostic launched

2 files changed, 19 insertions, 6 deletions

diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py
index c056fed..7647fae 100644
--- a/ep_run/casc_bp_train.py
+++ b/ep_run/casc_bp_train.py
@@ -20,6 +20,7 @@ ap.add_argument('--tok_init', type=float, default=0.0)  # >0: init tok/pos std (
 ap.add_argument('--cosine', action='store_true')         # warmup then cosine decay to lr_min_ratio*lr over --steps (long runs)
 ap.add_argument('--lr_min_ratio', type=float, default=0.1)
 ap.add_argument('--qk_norm', action='store_true')        # RMS-norm q,k per head before scores (OLMo2-style; bounds logits, analog-friendly)
+ap.add_argument('--final_ln', action='store_true')       # final LayerNorm before readout (standard GPT; bounds sig_tok growth -> keeps beta/estimator healthy on long runs)
 args = ap.parse_args()
 torch.manual_seed(args.seed)
 dev = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -73,7 +74,8 @@ if args.tok_init > 0:
         tok.weight.normal_(0, args.tok_init); pos.weight.normal_(0, args.tok_init)
 blocks = nn.ModuleList([Block(args.C, args.H, args.qk_norm) for _ in range(args.L)]).to(dev)
 mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1)
-params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters())
+ln_f = nn.LayerNorm(args.C).to(dev) if args.final_ln else nn.Identity()
+params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) + list(ln_f.parameters())
 if args.opt == 'muon':
     from muon import build_hybrid
     opt, sched = build_hybrid(blocks, params, args.lr, args.muon_lr, args.warmup)
@@ -91,7 +93,7 @@ else:
 def fwd(x):
     z = tok(x) + pos(torch.arange(args.T, device=dev))[None]
     for b in blocks: z = b(z, mask)
-    return z @ tok.weight.t()
+    return ln_f(z) @ tok.weight.t()
 
 @torch.no_grad()
 def evaluate(nb=6):
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index f6ae71b..74d6a3b 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -31,6 +31,9 @@ ap.add_argument('--beta_fixed', action='store_true')     # disable sig^2 schedul
 ap.add_argument('--cosine', action='store_true')         # warmup then cosine decay to lr_min_ratio*lr over --steps (long runs)
 ap.add_argument('--lr_min_ratio', type=float, default=0.1)
 ap.add_argument('--qk_norm', action='store_true')        # RMS-norm q,k per head before scores (OLMo2-style; bounds logits, analog-friendly)
+ap.add_argument('--final_ln', action='store_true')       # final LayerNorm before readout (standard GPT; bounds sig_tok growth -> keeps beta/estimator healthy on long runs)
+ap.add_argument('--resume', default='')                  # path to a ckpt (tok/pos/blocks) to continue from; step taken from ckpt
+ap.add_argument('--sig0', type=float, default=-1.0)      # override SIG0 (beta-schedule ref); needed on resume to restore original beta regime
 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)
@@ -94,8 +97,15 @@ if args.compile:
         print(f'[compile] disabled ({e})', flush=True)
 mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1)
 W_out = nn.Parameter(torch.randn(vocab, args.C, device=dev) * 0.02) if args.untie else None
-readout = (lambda z: z @ W_out.t()) if args.untie else (lambda z: z @ tok.weight.t())
-all_params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) + ([W_out] if args.untie else [])
+ln_f = nn.LayerNorm(args.C).to(dev) if args.final_ln else nn.Identity()
+readout = (lambda z: ln_f(z) @ W_out.t()) if args.untie else (lambda z: ln_f(z) @ tok.weight.t())
+all_params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) + list(ln_f.parameters()) + ([W_out] if args.untie else [])
+start_step = 0
+if args.resume:
+    _ck = torch.load(args.resume, map_location=dev, weights_only=False)
+    tok.load_state_dict(_ck['tok']); pos.load_state_dict(_ck['pos']); blocks.load_state_dict(_ck['blocks'])
+    start_step = int(_ck.get('step', 0))
+    print(f'[resume] loaded {args.resume} at step {start_step}', flush=True)
 if args.opt == 'muon':
     from muon import build_hybrid
     opt, sched = build_hybrid(blocks, all_params, args.lr, args.muon_lr, args.warmup)
@@ -189,7 +199,7 @@ def ep_step(x, y):
         GOV['sig'] = tok_sigma()
     GOV['step'] = GOV.get('step', 0) + 1
     sig = GOV['sig']
-    if SIG0 is None: SIG0 = sig
+    if SIG0 is None: SIG0 = args.sig0 if args.sig0 > 0 else sig
     beta_t = args.beta * GOV['bscale'] * (SIG0 * SIG0) / max(sig * sig, 1e-9)
     if args.beta_fixed: beta_t = args.beta * GOV['bscale']
     if args.beta_floor > 0.0: beta_t = max(beta_t, args.beta_floor)
@@ -253,7 +263,8 @@ print(f'[{args.tag}] cascade-EP(EQUILIBRIUM/fb) L{args.L} C{args.C} T{args.T} be
       f'K={args.K} geta={args.geta} | {n/1e6:.2f}M | {dev}', flush=True)
 best, t0 = 1e9, time.time()
 skips = 0
-for step in range(args.steps + 1):
+for _ in range(start_step): sched.step()   # advance LR schedule to the resumed step
+for step in range(start_step, args.steps + 1):
     x, y = get_batch('train')
     ce, beta_t, rounds, ok = ep_step(x, y)
     if not ok: skips += 1