From c81d6a5b80122dd7f53dc2ff876b296742bcc25b Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Fri, 10 Jul 2026 08:09:14 -0500 Subject: AUDIT: retract confounded Muon verdict; tone down parity claims (n=3, best-of-noisy-val); scope depth-tax claim to 4k horizon; guard-split skip telemetry (skd/skg); add BP control arm diag_D_bp (+--resume in BP trainer); record resume confounds Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/casc_bp_train.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) (limited to 'ep_run/casc_bp_train.py') diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py index 7647fae..069d858 100644 --- a/ep_run/casc_bp_train.py +++ b/ep_run/casc_bp_train.py @@ -21,6 +21,7 @@ ap.add_argument('--cosine', action='store_true') # warmup then cosine de 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 args = ap.parse_args() torch.manual_seed(args.seed) dev = 'cuda' if torch.cuda.is_available() else 'cpu' @@ -76,6 +77,12 @@ blocks = nn.ModuleList([Block(args.C, args.H, args.qk_norm) for _ in range(args. mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1) 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()) +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, params, args.lr, args.muon_lr, args.warmup) @@ -116,7 +123,8 @@ n = sum(p.numel() for p in params) print(f'[{args.tag}] cascade-BP L{args.L} C{args.C} H{args.H} T{args.T} | {n/1e6:.2f}M params | {dev}', flush=True) best, t0 = 1e9, time.time() outdir = Path('runs'); outdir.mkdir(exist_ok=True) -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') loss = F.cross_entropy(fwd(x).reshape(-1, vocab), y.reshape(-1)) opt.zero_grad(set_to_none=True); loss.backward() -- cgit v1.2.3