From d5804f62d488be9dbe7222c8b67bb09fa578c36d Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Wed, 15 Jul 2026 21:39:08 -0500 Subject: 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 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/casc_eq_train.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) (limited to 'ep_run/casc_eq_train.py') 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) -- cgit v1.2.3