diff options
Diffstat (limited to 'src/zbp_scaling/zbp/hardware.py')
| -rw-r--r-- | src/zbp_scaling/zbp/hardware.py | 176 |
1 files changed, 176 insertions, 0 deletions
diff --git a/src/zbp_scaling/zbp/hardware.py b/src/zbp_scaling/zbp/hardware.py new file mode 100644 index 0000000..32ae482 --- /dev/null +++ b/src/zbp_scaling/zbp/hardware.py @@ -0,0 +1,176 @@ +"""Analog-hardware simulation layer for hardware-in-the-loop ZBP (paper part 3). + +Models the fully-analog machine sketched in the plan: crossbar MVMs with programming quantization and +write noise, per-device saturating transfer functions, input DACs, a Walsh-dither + lock-in measurement +channel (gain error / DC offset / readout noise that trades against integration time), sample-and-hold +error transport with droop, and pulse-quantized weight updates. Everything is exposed to training only +through forward evaluations -- the simulation's autograd is used exclusively by the validation harness. +""" +import math +from dataclasses import dataclass, replace + +import torch +import torch.nn as nn + + +@dataclass +class HWConfig: + w_bits: int = 7 # crossbar programming resolution (0 = ideal) + w_write_noise: float = 0.01 # relative programming noise per write + w_range: float = 2.0 # programmable weight range [-w_range, w_range] + dev_gain_std: float = 0.10 # per-unit transfer-function variation phi_i(z) = g tanh(a z + c) + d + dev_a_std: float = 0.10 + dev_c_std: float = 0.05 + dev_d_std: float = 0.02 + dac_bits: int = 8 # input DAC (0 = ideal) + dac_range: float = 4.0 + meas_gain_std: float = 0.02 # lock-in channel: per-unit gain error + meas_offset: float = 0.01 # per-unit DC offset (cancels in the +/- difference) + meas_sigma0: float = 0.01 # readout noise std at integration time 1 + t_int: float = 1.0 # integration time (noise scales as sigma0 / sqrt(t_int)) + sh_droop: float = 0.0 # fraction of the held error lost during a block's probe phase + sh_noise: float = 0.0 # additive noise on the held error + update_lsb: float = 0.0 # weight-update quantum (0 = continuous) + + def replace(self, **kw): + return replace(self, **kw) + + +def _quantize(x, bits, rng): + if bits <= 0: + return x.clamp(-rng, rng) + step = 2 * rng / (2 ** bits - 1) + return (x.clamp(-rng, rng) / step).round() * step + + +class AnalogLinear(nn.Module): + """Crossbar MVM: the ideal parameter W is 'programmed' into W_eff = quantize(W) + write noise. + program() is called after every optimizer step (one write per update).""" + + def __init__(self, d_in, d_out, hw: HWConfig, gen=None): + super().__init__() + self.hw = hw + self.weight = nn.Parameter(torch.empty(d_out, d_in)) + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + self.register_buffer("w_eff", torch.zeros_like(self.weight)) + self.gen = gen + self.program() + + @torch.no_grad() + def program(self): + w = _quantize(self.weight, self.hw.w_bits, self.hw.w_range) + if self.hw.w_write_noise > 0: + scale = self.weight.abs().mean().clamp_min(1e-12) + noise = torch.randn(w.shape, generator=self.gen) # CPU generator; move to the weight's device + w = w + self.hw.w_write_noise * scale * noise.to(w.device) + self.w_eff.copy_(w) + + def forward(self, x): + # physics uses the programmed weights; the ideal parameter only exists inside the digital optimizer. + # (straight-through: gradients w.r.t. `weight` flow as if W_eff = weight, which is exactly the local + # in-block rule a physical outer-product update implements.) + return x @ (self.w_eff + (self.weight - self.weight.detach())).t() + + +class DeviceNonlinearity(nn.Module): + """phi_i(z) = g_i tanh(a_i z + c_i) + d_i with fixed per-device variation (no analytic form assumed + by training: only this forward is ever called).""" + + def __init__(self, d, hw: HWConfig, gen=None): + super().__init__() + g = 1 + hw.dev_gain_std * torch.randn(d, generator=gen) + a = 1 + hw.dev_a_std * torch.randn(d, generator=gen) + c = hw.dev_c_std * torch.randn(d, generator=gen) + dd = hw.dev_d_std * torch.randn(d, generator=gen) + for n, t in [("g", g), ("a", a), ("c", c), ("d", dd)]: + self.register_buffer(n, t) + + def forward(self, z): + return self.g * torch.tanh(self.a * z + self.c) + self.d + + +class DAC(nn.Module): + def __init__(self, hw: HWConfig): + super().__init__() + self.hw = hw + + def forward(self, x): + q = _quantize(x, self.hw.dac_bits, self.hw.dac_range) + return x + (q - x).detach() # straight-through for the validation autograd only + + +class AnalogBranch(nn.Module): + """DAC -> crossbar -> device nonlinearity -> crossbar (one residual branch of the analog machine).""" + + def __init__(self, d, hidden, hw: HWConfig, gen=None, out_scale=1.0): + super().__init__() + self.dac = DAC(hw) + self.a1 = AnalogLinear(d, hidden, hw, gen) + self.phi = DeviceNonlinearity(hidden, hw, gen) + self.a2 = AnalogLinear(hidden, d, hw, gen) + with torch.no_grad(): + self.a2.weight.mul_(out_scale) + self.a2.program() + + def forward(self, x): + return self.a2(self.phi(self.a1(self.dac(x)))) + + +class Measure: + """Lock-in measurement channel applied to the block output before the digital dot product with v: + y_meas = (1 + gamma) * y + offset + sigma0/sqrt(t_int) * xi (fresh xi per query).""" + + def __init__(self, d, hw: HWConfig, gen=None, device="cpu"): + self.gamma = (hw.meas_gain_std * torch.randn(d, generator=gen)).to(device) + self.offset = (hw.meas_offset * torch.randn(d, generator=gen)).to(device) + self.sigma = hw.meas_sigma0 / math.sqrt(max(hw.t_int, 1e-12)) + self.device = device + + def __call__(self, y): + out = (1 + self.gamma) * y + self.offset + if self.sigma > 0: + out = out + self.sigma * torch.randn_like(y) + return out + + +class SampleHold: + """Error-transport channel: the held error droops and picks up noise while the block is probed.""" + + def __init__(self, hw: HWConfig): + self.droop, self.noise = hw.sh_droop, hw.sh_noise + + def __call__(self, v): + out = (1 - self.droop) * v + if self.noise > 0: + out = out + self.noise * v.std() * torch.randn_like(v) + return out + + +class PulseQuantizedSGD: + """Wraps an optimizer: applied weight changes are rounded to multiples of update_lsb, and every + AnalogLinear is re-programmed (quantize + write noise) after the update.""" + + def __init__(self, opt, model, hw: HWConfig): + self.opt, self.model, self.hw = opt, model, hw + self._prev = None + + def zero_grad(self, set_to_none=True): + self.opt.zero_grad(set_to_none=set_to_none) + + @torch.no_grad() + def _snapshot(self): + return [p.detach().clone() for p in self.model.parameters()] + + def step(self): + if self.hw.update_lsb > 0: + prev = self._snapshot() + self.opt.step() + with torch.no_grad(): + for p, q in zip(self.model.parameters(), prev): + delta = p.detach() - q + p.copy_(q + (delta / self.hw.update_lsb).round() * self.hw.update_lsb) + else: + self.opt.step() + for m in self.model.modules(): + if isinstance(m, AnalogLinear): + m.program() |