Stage-0 gate #1: --qup_bits (absolute-grid stochastic-rounding weight quantization = finite analog cell levels); qup8/qup6/qup10 arms launched on GPU0

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

1 files changed, 16 insertions, 0 deletions

diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py
index dfb7098..cbf94b9 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('--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
 ap.add_argument('--centfast', action='store_true')        # centered via ONE doubled batch [x;x], +beta/-beta halves
                                                           # (shared kernels; math identical to sequential centered)
 args = ap.parse_args()
@@ -627,6 +630,19 @@ for step in range(start_step, args.steps + 1):
                 GOV['K'] -= 1; GOV['bscale'] = min(GOV['bscale'] * 1.05, 1.0)
     torch.nn.utils.clip_grad_norm_(all_params, 1.0)
     opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
+    if args.qup_bits > 0:
+        # STAGE-0 HW GATE: finite conductance levels. Snap every weight to an ABSOLUTE
+        # per-tensor grid (range/2^bits) with stochastic rounding (unbiased) — emulates
+        # analog cell writes; per-step deltas below one level survive only in expectation.
+        with torch.no_grad():
+            for p in all_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 DDP and args.sync_check > 0 and step % args.sync_check == 0 and step > 0:
         with torch.no_grad():
             h = torch.stack([torch.stack((p.double().sum(), (p.double() ** 2).sum())) for p in all_params]).sum(0)