A0.4 precision gate (TF32 harmless cos 0.9946==fp32; pure-bf16 cos 0.9427) + Muon hybrid optimizer (--opt muon) wired into both trainers; D1a flagship matrix launched (L12xC512, 3xBP + 3xEP + Muon arms)

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

1 files changed, 8 insertions, 2 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 10206f6..4eba267 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -21,6 +21,8 @@ ap.add_argument('--wandb', default=''); ap.add_argument('--wandb_run', default='
 ap.add_argument('--kmax', type=int, default=8)          # adaptive fb rounds cap
 ap.add_argument('--noguard', action='store_true')       # diagnosis: skip only non-finite grads
 ap.add_argument('--untie', action='store_true')         # separate readout matrix (untied from tok)
+ap.add_argument('--opt', choices=['adamw', 'muon'], default='adamw')
+ap.add_argument('--muon_lr', type=float, default=0.02)
 ap.add_argument('--tok_init', type=float, default=0.0)  # >0: init tok/pos with this std (GPT-standard 0.02)
 ap.add_argument('--compile', action='store_true')       # torch.compile each block (free speed where supported)
 ap.add_argument('--sig_every', type=int, default=25)    # tok-sigma refresh interval (amortized)
@@ -66,8 +68,12 @@ 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 [])
-opt = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=1e-4)
-sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(args.warmup, 1)))
+if args.opt == 'muon':
+    from muon import build_hybrid
+    opt, sched = build_hybrid(blocks, all_params, args.lr, args.muon_lr, args.warmup)
+else:
+    opt = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=1e-4)
+    sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(args.warmup, 1)))
 NBT = args.B * args.T
 
 def free_states_graphed(x):