Stage-0 gate #1b: --qcomp_bits (compute on DAC-grid weights, fp32 master = T64 word-streaming / shadow accumulation); qcomp8/6/4 launched

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

1 files changed, 17 insertions, 0 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index cbf94b9..17a0c70 100644
--- a/ep_run/casc_eq_train.py
+++ b/ep_run/casc_eq_train.py
@@ -65,6 +65,9 @@ ap.add_argument('--est', choices=['single', 'centered', 'richardson'], default='
 ap.add_argument('--est_late', choices=['', 'centered'], default='')
 ap.add_argument('--est_late_at', type=int, default=0)     # switch --est -> --est_late at this step (process-local,
                                                           # bf_late_at semantics); centered is TAIL medicine
+ap.add_argument('--qcomp_bits', type=int, default=0)      # STAGE-0 T64 scenario: forward/transpose COMPUTE
+                                                          # on grid-snapped weights, fp32 master gets updates
+                                                          # (= word-streaming / shadow accumulation)
 ap.add_argument('--qup_bits', type=int, default=0)        # STAGE-0: quantize weights to an absolute
                                                           # per-tensor grid after each update (stochastic
                                                           # rounding); emulates finite analog cell levels
@@ -613,7 +616,21 @@ skips = 0
 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:
+        # STAGE-0 HW GATE #1b (T64 scenario): COMPUTE runs on weights snapped to the DAC
+        # grid (deterministic round-to-nearest); the fp32 master (DDR / shadow accumulator)
+        # receives the update. Equivalent to word-streaming and to resident-cell + shadow.
+        with torch.no_grad():
+            QSAVE = [p.detach().clone() for p in all_params]
+            for p in all_params:
+                rng = float(p.abs().max())
+                if rng <= 0: continue
+                g_ = rng / (2 ** (args.qcomp_bits - 1))
+                p.copy_((p / g_).round() * g_)
     ce, beta_t, rounds, ok = ep_step(x, y)
+    if args.qcomp_bits > 0:
+        with torch.no_grad():
+            for p, q in zip(all_params, QSAVE): p.copy_(q)
     if not ok: skips += 1
     gcos = float('nan')
     if args.gate_every > 0 and step % args.gate_every == 0 and ok: