diff options
| -rw-r--r-- | docs/campaign/CASCADE_ABLATION_PLAN.md | 22 | ||||
| -rw-r--r-- | ep_run/casc_bp_train.py | 6 | ||||
| -rw-r--r-- | ep_run/casc_eq_train.py | 19 |
3 files changed, 41 insertions, 6 deletions
diff --git a/docs/campaign/CASCADE_ABLATION_PLAN.md b/docs/campaign/CASCADE_ABLATION_PLAN.md index 8e0ac43..19f97bb 100644 --- a/docs/campaign/CASCADE_ABLATION_PLAN.md +++ b/docs/campaign/CASCADE_ABLATION_PLAN.md @@ -283,3 +283,25 @@ fires at the early analog read (nf step 2500) or all-done. tok), qk_norm + beta_floor 3e-4 + cosine, warmup 500, 2.4 it/s solo -> ~6.8 h. The "neng kan" generation demo. Watcher fires at step 10000 (first generation-worthy ckpt) / done / death. Then Stage 2 (FineWeb-Edu) -> Stage 3 (OLMo2). + +### RESULT 5 (2026-07-10 07:3x): Stage-1 epoch BLEW UP @step 12100 — root-cause diagnosis (sig story REFUTED). +The qk_norm+beta_floor+cosine epoch was healthy to ~11400 (best val 1.6669) then blew up (val 1.67->7.6, +gn pre-clip 0.5->53) and oscillated in a degraded regime. **My first guess (sig_tok growth -> SNR +collapse -> add final_ln) was WRONG, refuted by its own telemetry:** +- sig rose only +8% (29.8@10000 -> 32.2@12000) then PLATEAUED; it was already ~30 at step 10000 when + everything was healthy. An 8% change cannot cause a catastrophic transition. +- cos was FINE (0.9935) until step 11900; the cos drop is a CONSEQUENCE of the blowup, not the cause. +- grad-clip is ALREADY present (clip 1.0); gn=53 is pre-clip telemetry. Not a magnitude-spike issue. +**LEADING INDICATOR = skips (drift-guard rejections = nudged fb relaxation drift>0.5 = CONVERGENCE +FAILURE).** skips accelerate from ~step 11000 (4->13 by 11400) BEFORE gn (11700), cos (12000), val +(12100). **Diagnosis: a CONTRACTIVITY BIFURCATION in the nudged fb relaxation** -- as training sharpens +the operator (block Jacobians grow), an increasing fraction of batches have a non-contractive nudged +iteration -> skipped -> gradient bias -> a marginally-converged batch emits a bad step -> over the edge. +**This is the cascade analog of the looped-EP Hopf wall** (non-conservative attention loses +contractivity as CE drops -- documented in ep-c512-residual-defense-fix). 4000-step runs never saw it +(operator not sharp enough yet; edge ~step 11400). Right fix = CONTRACTIVITY control (resreg/jacreg or +geta<1 damping), NOT final_ln. +**CONFIRMATORY A/B/C (resume from ckpt-10000, pre-bifurcation, beta floored 3e-4 via --sig0 1.6):** + A=control (K3,lr1e-3) -> should reproduce skip-climb+blowup; B=K8 (does more fb rounds hold skips? + = marginal-contractivity test); C=lr3e-4 (slower sharpening -> delayed edge? = driver test). + Code added: --resume, --sig0, --final_ln, --qk_norm(CausalSelfAttn/SDPA). Watcher diag_watch.sh armed. 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 |