diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_eq_train.py | 17 |
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: |