Stage-0 BP mirror arms: qup/qcomp flags in casc_bp_train; bp_qctl + 6 quant arms chained after qcomp (decision metric = Delta_EP vs Delta_BP)

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

1 files changed, 23 insertions, 0 deletions

diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py
index bed8d1c..7bafc16 100644
--- a/ep_run/casc_bp_train.py
+++ b/ep_run/casc_bp_train.py
@@ -27,6 +27,8 @@ ap.add_argument('--resume', default='')                  # path to a ckpt (tok/p
 ap.add_argument('--olmo2', action='store_true')          # OLMo2-standard block (see casc_eq_train.py)
 ap.add_argument('--wd', type=float, default=-1.0)        # >=0: grouped weight decay; <0 = legacy uniform 1e-4
 ap.add_argument('--zloss', type=float, default=0.0)      # z-loss coefficient; 0 = off
+ap.add_argument('--qup_bits', type=int, default=0)        # STAGE-0 mirror: naked resident-cell writes
+ap.add_argument('--qcomp_bits', type=int, default=0)      # STAGE-0 mirror: compute on DAC grid, fp32 master
 ap.add_argument('--data', default='tinystories_bpe')     # dataset dir under ep_run/data
 args = ap.parse_args()
 if args.olmo2 and args.tok_init <= 0: args.tok_init = 0.02
@@ -200,14 +202,35 @@ outdir = Path('runs'); outdir.mkdir(exist_ok=True)
 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')
+    if args.qcomp_bits > 0:
+        with torch.no_grad():
+            QSAVE = [p.detach().clone() for p in params]
+            for p in params:
+                rng = float(p.abs().max())
+                if rng <= 0: continue
+                g_ = rng / (2 ** (args.qcomp_bits - 1))
+                p.copy_((p / g_).round() * g_)
     with torch.autocast('cuda', dtype=torch.bfloat16, enabled=args.amp):
         logits = fwd(x).reshape(-1, vocab)
         loss = F.cross_entropy(logits, y.reshape(-1))
         if args.zloss > 0:
             loss = loss + args.zloss * (torch.logsumexp(logits.float(), -1) ** 2).mean()
     opt.zero_grad(set_to_none=True); loss.backward()
+    if args.qcomp_bits > 0:
+        with torch.no_grad():
+            for p, q in zip(params, QSAVE): p.copy_(q)
     torch.nn.utils.clip_grad_norm_(params, 1.0)
     opt.step(); sched.step()
+    if args.qup_bits > 0:
+        with torch.no_grad():
+            for p in params:
+                if p.ndim < 1: continue
+                rng = float(p.abs().max())
+                if rng <= 0: continue
+                g_ = rng / (2 ** (args.qup_bits - 1))
+                q = p / g_
+                fl = q.floor()
+                p.copy_((fl + (torch.rand_like(p) < (q - fl)).float()) * g_)
     if step % args.log == 0:
         val = evaluate(); best = min(best, val)
         print(f'step {step:5d}/{args.steps} | train {loss.item():.4f} val {val:.4f} (best {best:.4f}) '