RESULT 11: stage1b sealed (gap 0.050); beta ceiling unfound (3e-3 -> cos 0.9974); K exonerated on OLMo2; bf16 naive-cast dead at any beta (TF32 stays cost baseline); f3e3 CE-endpoint cont launched

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

1 files changed, 10 insertions, 2 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index 5a3e640..ad82b5e 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -40,6 +40,7 @@ ap.add_argument('--zloss', type=float, default=0.0)      # z-loss coefficient on
 ap.add_argument('--kretry', type=int, default=0)         # >0: on drift-reject, RETRY the batch once with this many fb rounds (diag B: K8 converges the marginal batches) instead of dropping it
 ap.add_argument('--bf_late', type=float, default=0.0)    # >0: raise beta_floor to this value from step --bf_late_at (late-training SNR fix; dose-response 2026-07-10)
 ap.add_argument('--bf_late_at', type=int, default=25000)
+ap.add_argument('--bf16', action='store_true')           # cast model to bf16 (E-accumulation + tok_sigma stay fp32) — the x0.5 cost lever, GATE before production
 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)
@@ -173,6 +174,13 @@ if args.resume:
     if _ck.get('lnf') is not None and not isinstance(ln_f, nn.Identity): ln_f.load_state_dict(_ck['lnf'])
     start_step = int(_ck.get('step', 0))
     print(f'[resume] loaded {args.resume} at step {start_step}', flush=True)
+if args.bf16:
+    for _m in (tok, pos, blocks):
+        _m.to(torch.bfloat16)
+    if not isinstance(ln_f, nn.Identity): ln_f.to(torch.bfloat16)
+    if args.untie:
+        with torch.no_grad(): W_out.data = W_out.data.to(torch.bfloat16)
+    print('[bf16] model cast to bfloat16 (E-accum + sigma stay fp32)', 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,
@@ -220,7 +228,7 @@ def free_states_graphed(x):
 @torch.no_grad()
 def tok_sigma(iters=8):
     """top singular value of tok.weight (power iteration on the raw matrix)."""
-    W = W_out if args.untie else tok.weight
+    W = (W_out if args.untie else tok.weight).float()
     v = torch.randn(W.shape[1], device=dev); v /= v.norm()
     sig = 1.0
     for _ in range(iters):
@@ -312,7 +320,7 @@ def ep_step(x, y):
             return free_ce, beta_t, GOV['K'], False
     GOV['drift'] = drift
     E = 0.0
-    for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach() - o) ** 2).sum()
+    for z, o in zip(zp, last_outs): E = E + 0.5 * ((z.detach().float() - o.float()) ** 2).sum()   # fp32 accumulation (bf16-safe; no-op in fp32)
     obj = E / (NBT * beta_t) + obj_loss(readout(zp[-1].detach()).reshape(-1, vocab), y.reshape(-1))
     gs = torch.autograd.grad(obj, all_params, allow_unused=True)
     gn = 0.0