Spaces:
Paused
Paused
| """Finite-state counter synapse — the memory lever for parameters+optimizer+gradients. | |
| A weight is a per-synapse finite-state automaton: a ternary visible weight t in {-1,0,+1} | |
| plus a residual counter c, packed into one small state. The optimizer state lives *inside* | |
| that state (plus a per-row scale and a per-row RMS second moment), so there is no FP master | |
| weight and no per-weight Adam moment. The update is fused into the layer's backward, so no | |
| full-MODEL gradient buffer is ever retained (only one layer's grad_w is live at a time). | |
| NOTE: by default (tile_rows=0 -> out_features) that one layer's grad_w is still a full | |
| [out,in] tile, materialized transiently and freed after the update; tile_rows>0 bounds it to | |
| [R,in], and the Triton update kernel forms it in registers (no dense grad_w at all). The | |
| FP-master/Adam-free property holds in every case; the "no dense grad_w" property is the | |
| kernel/tiled path, not the dense-PyTorch default. | |
| This is pure PyTorch: it runs on CPU and CUDA with stock torch, no custom engine. It is a | |
| *correctness / dynamics* implementation — it decodes states to ordinary tensors around the | |
| GEMM, so it validates learning and persistent-state accounting but does not yet realize the | |
| sub-byte bandwidth win (that needs a Triton/CUDA packed kernel; see README "Roadmap"). | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import torch | |
| import torch.distributed as dist | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| __all__ = [ | |
| "C_DEFAULT", | |
| "encode_state", | |
| "decode_state", | |
| "stochastic_round", | |
| "ternary_gradient_unbiased", | |
| "CompactCounterLinear", | |
| "RMSCounterLinear", | |
| ] | |
| def _allreduce_grad_w_(grad_w: torch.Tensor) -> None: | |
| """Average a counter weight-gradient across data-parallel ranks, in place. No-op unless a | |
| process group is initialized. Called inside backward so all replicas apply an identical | |
| counter update and their packed states stay synchronized (the counter has no Parameter | |
| gradient for DDP to handle -- the optimizer is the in-place state update itself).""" | |
| if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: | |
| dist.all_reduce(grad_w, op=dist.ReduceOp.AVG) | |
| # C=8 -> counter c in {-7..+7} (15 levels), 3*15 = 45 reachable states (fits 6 bits / uint8). | |
| # Larger C is allowed while 3*(2C-1) <= 256 (uint8); C=11 gives 63 states (best per ablation). | |
| C_DEFAULT = 8 | |
| def stochastic_round(x: torch.Tensor) -> torch.Tensor: | |
| """Unbiased stochastic rounding to integers, valid for positive and negative x.""" | |
| floor = torch.floor(x) | |
| return floor + (torch.rand_like(x) < (x - floor)).to(x.dtype) | |
| def encode_state(t: torch.Tensor, c: torch.Tensor, C: int = C_DEFAULT) -> torch.Tensor: | |
| """Encode t in {-1,0,1}, c in {-(C-1),...,C-1} into one uint8 state.""" | |
| levels = 2 * C - 1 | |
| code = (t.to(torch.int16) + 1) * levels + (c.to(torch.int16) + (C - 1)) | |
| return code.to(torch.uint8) | |
| def decode_state(state: torch.Tensor, C: int = C_DEFAULT) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Decode a uint8 state into int16 ternary weight t and residual counter c.""" | |
| levels = 2 * C - 1 | |
| z = state.to(torch.int16) | |
| t = torch.div(z, levels, rounding_mode="floor") - 1 | |
| c = torch.remainder(z, levels) - (C - 1) | |
| return t, c | |
| def ternary_gradient_unbiased(g: torch.Tensor) -> torch.Tensor: | |
| """Row-wise unbiased ternary estimator: output in {-a,0,+a} with E[Q(g)|g]=g, a=max|g_j|.""" | |
| amplitude = g.abs().amax(dim=1, keepdim=True) | |
| safe = amplitude.clamp_min(1e-30) | |
| p = (g.abs() / safe).clamp_(0.0, 1.0) | |
| event = (torch.rand_like(p) < p).to(g.dtype) | |
| return torch.sign(g) * event * amplitude | |
| class _FusedCounterLinearFn(torch.autograd.Function): | |
| """Linear whose weight update is fused into backward at row-tile granularity. | |
| A scalar `tap` (requires_grad=True) is threaded through purely so the output requires | |
| grad and backward runs even when neither the input nor any Parameter needs a gradient -- | |
| a counter layer fed raw input (e.g. the first layer) must still self-update. This mirrors | |
| the engine's "attach the update node regardless of x.requires_grad". | |
| """ | |
| def forward(ctx, x: torch.Tensor, module: "CompactCounterLinear", tap: torch.Tensor) -> torch.Tensor: | |
| if module._outstanding_forward: | |
| raise RuntimeError( | |
| "CompactCounterLinear was reused before its previous backward. " | |
| "Weight sharing and gradient accumulation need an explicit scheduler." | |
| ) | |
| module._outstanding_forward = True | |
| ctx.module = module | |
| y = module._forward_matmul(x) | |
| # Activation-memory lever: if act_save_bits is set, store an UNBIASED low-bit | |
| # quantization of x (codes + per-row scale) instead of fp x. The update only needs | |
| # E[Q(x)|x] = x, so it stays unbiased; the saved activation shrinks from fp to b bits. | |
| if module.act_save_bits: | |
| from .actquant import pack_int4, quantize_codes | |
| x2 = x.reshape(-1, x.shape[-1]) | |
| codes, scale = quantize_codes(x2, module.act_save_bits, dim=-1) | |
| if module.act_save_bits == 4: | |
| # true 4-bit packing: 2 codes per byte -> 0.5 byte/elem (codes in [-7,7]). | |
| store = pack_int4(codes) | |
| ctx.packed4 = True | |
| ctx.n_codes = x2.numel() | |
| else: | |
| # int8 (1 byte) for bits<=8, int16 for 9..15. | |
| store = codes.to(torch.int8) if module.act_save_bits <= 8 else codes.to(torch.int16) | |
| ctx.packed4 = False | |
| ctx.save_for_backward(store, scale) | |
| ctx.x_shape = x.shape | |
| ctx.quantized = True | |
| else: | |
| ctx.save_for_backward(x) | |
| ctx.quantized = False | |
| return y | |
| def backward(ctx, grad_out: torch.Tensor): | |
| module: CompactCounterLinear = ctx.module | |
| saved_i8 = None # (codes int8 [M,in], row scale [M,1]) | |
| if ctx.quantized: | |
| store, scale = ctx.saved_tensors | |
| if ctx.packed4: | |
| from .actquant import unpack_int4 | |
| codes = unpack_int4(store, ctx.n_codes).reshape(-1, module.in_features) | |
| else: | |
| codes = store | |
| if module.act_save_bits == 8: # int8 codes -> reusable by the int8 update | |
| saved_i8 = (store, scale) | |
| x2 = (codes.to(scale.dtype) * scale) # dequant Q(x): [-1, in] | |
| x = x2.reshape(ctx.x_shape) | |
| else: | |
| (x,) = ctx.saved_tensors | |
| x2 = x.reshape(-1, x.shape[-1]) | |
| go2 = grad_out.reshape(-1, grad_out.shape[-1]) | |
| # grad_x: a kernel subclass (Triton) computes it straight from packed state with no | |
| # dense weight at all; the base path accumulates it tile-by-tile in the loop below. | |
| use_kernel = module._has_fast_grad_x() | |
| if use_kernel: | |
| grad_x2 = module._backward_grad_x(go2) | |
| else: | |
| grad_x2 = torch.zeros((x2.shape[0], module.in_features), device=x.device, dtype=x.dtype) | |
| # The full weight-gradient never exists. Only [tile_rows, in_features] is live. | |
| # Adaptive decimation (memo M8): once a layer's flip-rate is tiny it is near-stable, so | |
| # apply the update only every _dec_period steps with lr scaled to compensate. Decided once | |
| # per backward; grad_x is always computed, only the update is skipped. | |
| do_update = module.training and module.update_enabled | |
| fire = module._decimation_apply() if do_update else False | |
| if fire: | |
| module._dec_flips_step = 0 # reset per-step flip accumulators (summed across tiles) | |
| module._dec_elems_step = 0 | |
| if do_update or not use_kernel: | |
| for lo in range(0, module.out_features, module.tile_rows): | |
| hi = min(lo + module.tile_rows, module.out_features) | |
| t_i, c_i = module._decode_rows(lo, hi) | |
| s_i = module.scale[lo:hi] | |
| if not use_kernel: | |
| w_i = s_i.to(x.dtype) * t_i.to(x.dtype) | |
| grad_x2.add_(go2[:, lo:hi] @ w_i) | |
| if fire: | |
| if module.update_compute == "int8" and saved_i8 is not None: | |
| # reuse the int8 activation saved in forward; quantize only grad_out | |
| from .int8_compute import int8_correlation_presaved | |
| grad_w_i = int8_correlation_presaved(go2[:, lo:hi], saved_i8[0], saved_i8[1]) | |
| elif module.update_compute == "int8": | |
| from .int8_compute import int8_correlation | |
| grad_w_i = int8_correlation(go2[:, lo:hi], x2) | |
| elif module.update_compute == "int4": | |
| from .int8_compute import int4_correlation | |
| grad_w_i = int4_correlation(go2[:, lo:hi], x2) | |
| elif module.update_compute == "fp8": | |
| from .int8_compute import fp8_correlation | |
| grad_w_i = fp8_correlation(go2[:, lo:hi], x2) | |
| else: | |
| grad_w_i = (go2[:, lo:hi].transpose(0, 1) @ x2).float() | |
| # proxy RMS: take the row second-moment from grad_out norms (* activation | |
| # energy) instead of the full ||G_o||^2 reduction (post-LN E[r_o^2]~||D_o||^2). | |
| proxy_gsq = None | |
| if getattr(module, "rms_mode", "exact") == "proxy": | |
| # E[r_o^2] ~ (sum_m D_mo^2) * E[X^2]: SUM over the M tokens, not mean -- | |
| # mean is off by 1/M, so the denominator scales ~1/sqrt(M) (batch/seq). | |
| proxy_gsq = go2[:, lo:hi].pow(2).sum(dim=0, keepdim=True).t() * x2.pow(2).mean() | |
| # Data-parallel: the counter optimizer lives in the state and is applied | |
| # in-place here, so there is no Parameter .grad for DDP to all-reduce. Sync | |
| # the counter gradient itself across ranks so every replica applies the same | |
| # update. Packed states then stay bit-identical -- VERIFIED empirically with | |
| # different data per rank and no per-step reseed (test_ddp_decimation; even | |
| # act_save_bits>0 gives 0 differing bytes). This holds because (a) grad_w is | |
| # averaged here, and (b) the torch SR (stochastic_round -> global torch.rand) | |
| # draws a SHAPE-determined, rank-symmetric number of samples each step, so the | |
| # global RNG stays phase-aligned. It is therefore CONDITIONAL: a data-dependent | |
| # random op (e.g. dropout with per-rank masks) would desync the torch.rand | |
| # stream and break bit-identity silently. The Triton kernel's hash-SR path | |
| # (per-element deterministic, no global RNG) is unconditionally DDP-safe. | |
| _allreduce_grad_w_(grad_w_i) | |
| # proxy RMS is built from this rank's raw grad_out/x, not the averaged grad_w, | |
| # so it must be reduced too -- otherwise the denominator differs per rank and | |
| # the packed states diverge across ranks. (exact/lagged read g_sq from the | |
| # already-averaged grad_w, so they need no extra reduce.) | |
| if proxy_gsq is not None: | |
| _allreduce_grad_w_(proxy_gsq) | |
| if not module._fused_update(lo, hi, grad_w_i): | |
| module._update_tile(lo, hi, grad_w_i, t_i, c_i, s_i, proxy_gsq=proxy_gsq) | |
| if fire: | |
| # per-layer flip-rate = flips summed over ALL tiles / elements over all tiles | |
| # (a single tile's rate would mis-drive decimation when tile_rows < out_features). | |
| module._dec_flip_rate = module._dec_flips_step / max(module._dec_elems_step, 1) | |
| module._decimation_observe() | |
| module._outstanding_forward = False | |
| # grads for (x, module, tap); tap's grad is unused. | |
| return grad_x2.reshape_as(x), None, None | |
| class CompactCounterLinear(nn.Module): | |
| """Ternary linear layer trained as a finite-state synaptic automaton. | |
| Persistent per-weight state in this implementation is 1 byte (uint8); the logical | |
| packed state is ceil(log2(states)) = 6 bits. Row scales stay FP32 (their O(out) | |
| cost is negligible). No FP master weight, no per-weight Adam moment. | |
| """ | |
| def __init__( | |
| self, | |
| in_features: int, | |
| out_features: int, | |
| *, | |
| C: int = C_DEFAULT, | |
| lr: float = 0.04, | |
| lr_scale: float = 2e-4, | |
| init_gain: float = 1.0, | |
| tile_rows: int = 0, | |
| local_grad_clip: float = 0.0, | |
| pulse_mode: str = "direct", | |
| act_save_bits: int = 0, | |
| decimate_updates: bool = False, | |
| cache_mode: str = "none", | |
| cache_patch: str = "full", | |
| update_compute: str = "fp", | |
| forward_compute: str = "fp", | |
| ) -> None: | |
| super().__init__() | |
| if 3 * (2 * C - 1) > 256: | |
| raise ValueError("C is too large for uint8 state encoding (need 3*(2C-1) <= 256)") | |
| self.in_features = int(in_features) | |
| self.out_features = int(out_features) | |
| self.C = int(C) | |
| self.lr = float(lr) | |
| self.lr_scale = float(lr_scale) | |
| # tile_rows=R: the update runs over R-row tiles, materializing only an [R,in] grad_w tile. | |
| # Default 0 = untiled (one [out,in] grad tile, one fused launch per layer). The T4 update | |
| # frontier shows R~256 BEATS untiled *for a large isolated layer* (2048x2048), but a full | |
| # s512 training step is FASTER untiled (17.1k vs 14.2k tok/s) because small layers turn | |
| # tiling into many tiny fused launches whose overhead dominates. So untiled stays the | |
| # default; tile_rows>0 is a memory knob (bounds the transient grad to R*in*4 bytes) worth | |
| # it only for very large layers or a tight VRAM budget. The result is identical either way. | |
| self.tile_rows = int(tile_rows) if int(tile_rows) > 0 else int(out_features) | |
| self.local_grad_clip = float(local_grad_clip) | |
| # 0 = store fp activation; >0 = store unbiased act_save_bits-bit Q(x) for the update. | |
| self.act_save_bits = int(act_save_bits) | |
| if pulse_mode not in {"direct", "ternary"}: | |
| raise ValueError("pulse_mode must be 'direct' or 'ternary'") | |
| self.pulse_mode = pulse_mode | |
| self.update_enabled = True | |
| self._outstanding_forward = False | |
| # Adaptive update decimation (memo M8): when on, a near-stable layer (tiny flip-rate) | |
| # updates only every _dec_period steps with lr scaled by the period to compensate. | |
| self.decimate_updates = bool(decimate_updates) | |
| self._dec_period = 1 | |
| self._dec_since = 0 | |
| self._dec_flip_rate = 1.0 | |
| self._dec_flips_step = 0 # flips accumulated across tiles THIS step (per-layer rate) | |
| self._dec_elems_step = 0 | |
| self._lr_mult = 1.0 | |
| t0 = torch.randint(-1, 2, (out_features, in_features), dtype=torch.int16) | |
| c0 = torch.zeros_like(t0) | |
| self.register_buffer("state", encode_state(t0, c0, C)) | |
| # Var(t)=2/3 for uniform {-1,0,1}; choose Var(w)=gain^2/fan_in. | |
| s0 = init_gain * math.sqrt(3.0 / (2.0 * in_features)) | |
| self.register_buffer("scale", torch.full((out_features, 1), s0, dtype=torch.float32)) | |
| # Diagnostics only: O(1) scalars, not per-weight optimizer state. | |
| self.register_buffer("update_events", torch.zeros((), dtype=torch.int64), persistent=False) | |
| self.register_buffer("weight_flips", torch.zeros((), dtype=torch.int64), persistent=False) | |
| # Derived visible-weight compute cache (memo M5): the visible ternary T is the only thing | |
| # forward/grad_x need; the hidden counter c is only the update's. cache_mode keeps T in a | |
| # GEMM-friendly dtype so the matmul never pays the 6-bit decode tax. It is a *derived view* | |
| # (not truth, not optimizer state): rebuilt from the state, refreshed on visible flips. | |
| if cache_mode not in {"none", "fp16", "int8"}: | |
| raise ValueError("cache_mode must be 'none', 'fp16' or 'int8'") | |
| # cache_patch (fusion-plan lever #4): how the persistent derived T-cache is refreshed after | |
| # an update. "full" rewrites the whole updated tile every step; "flip" writes ONLY the | |
| # elements whose visible ternary actually changed (the rest are already correct in the | |
| # persistent buffer). Bit-identical result; "flip" models the kernel's amortized-decode~0 | |
| # cost (with decimation almost no weight flips, so almost nothing is rewritten). The cache | |
| # is already persistent across steps (a persistent=False buffer, never rebuilt unless absent | |
| # or reloaded) -- "flip" is the cross-step flip-patching half of the lever. | |
| if cache_patch not in {"full", "flip"}: | |
| raise ValueError("cache_patch must be 'full' or 'flip'") | |
| self.cache_patch = cache_patch | |
| # diagnostics only: how many cache elements were patched (O(1) scalar, not optimizer state). | |
| self.register_buffer("cache_patches", torch.zeros((), dtype=torch.int64), persistent=False) | |
| if update_compute not in {"fp", "int8", "int4", "fp8"}: | |
| raise ValueError("update_compute must be 'fp', 'int8', 'int4' or 'fp8'") | |
| if forward_compute not in {"fp", "int8"}: | |
| raise ValueError("forward_compute must be 'fp' or 'int8'") | |
| # "int8"/"int4" form grad_w with an UNBIASED stochastic low-bit GEMM estimator (Tensor Cores | |
| # on CUDA) instead of the fp32 correlation -> training-neutral in expectation. "fp8" is a | |
| # deterministic fp8 (e4m3) estimator: slightly biased (round-to-nearest), parity-gated, the | |
| # error-feedback absorbs it (fusion-plan lever #5). All opt-in; fp is the default. | |
| self.update_compute = update_compute | |
| # forward_compute="int8" runs the forward Y=XT^T on the integer Tensor Cores (deterministic | |
| # round-to-nearest activation quant, so reversible/eager stay valid); pairs with the int8 | |
| # visible cache. grad_x and the update stay fp -- straight-through past the forward quant. | |
| self.forward_compute = forward_compute | |
| if forward_compute == "int8" and cache_mode == "none": | |
| cache_mode = "int8" # the int8 forward needs the int8 T view | |
| self.cache_mode = cache_mode | |
| if cache_mode != "none": | |
| # Built lazily on first use -- subclasses (e.g. the packed layout) finish constructing | |
| # their state after this base __init__. Rebuilt from the truth state after any load | |
| # (the cache is persistent=False, so it is never the source of truth). | |
| self.register_load_state_dict_post_hook( | |
| lambda m, _ic: m._build_t_cache() if hasattr(m, "_t_cache") else None) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| if torch.is_grad_enabled(): | |
| # tap forces the output to require grad so backward (and the self-update) runs | |
| # even when x and all Parameters need no gradient. | |
| tap = torch.zeros((), device=x.device, dtype=x.dtype, requires_grad=True) | |
| return _FusedCounterLinearFn.apply(x, self, tap) | |
| # inference: pure forward, no update, no graph. | |
| return self._forward_matmul(x) | |
| # --- storage abstraction (overridden by the packed-6bit subclass) ------------- | |
| # Base storage is one uint8 code per weight in `state` [out, in]. A subclass may store | |
| # `state` packed (4 codes / 3 bytes) and override these three to pack/unpack at the | |
| # boundary; the autograd Function and _update_tile go through them, so the update math | |
| # is shared and the persistent footprint is whatever the storage chooses. | |
| def _forward_matmul(self, x: torch.Tensor) -> torch.Tensor: | |
| # y = x @ W^T. Base path materializes the dense weight; a kernel subclass (Triton) | |
| # may override this to decode the packed state inside the GEMM with no dense weight. | |
| if self.forward_compute == "int8": | |
| from .int8_compute import int8_forward_ternary | |
| flat = x.reshape(-1, x.shape[-1]) | |
| # int8 GEMM on the Tensor Cores: (a_x q_x) T^T, then the per-output row scale s_o. | |
| # Deterministic quant (round-to-nearest) keeps F/G reversible-safe. | |
| y = int8_forward_ternary(flat, self._visible_t(torch.int8), stochastic=False) | |
| y = y * self.scale.reshape(1, self.out_features) | |
| return y.reshape(*x.shape[:-1], self.out_features).to(x.dtype) | |
| return F.linear(x, self._dense_weight(x.dtype)) | |
| def _has_fast_grad_x(self) -> bool: | |
| # Base path computes grad_x tile-by-tile in the backward loop. A kernel subclass | |
| # returns True and provides _backward_grad_x to form grad_x straight from state. | |
| return False | |
| def _backward_grad_x(self, grad_out2d: torch.Tensor) -> torch.Tensor: # pragma: no cover | |
| raise NotImplementedError | |
| def _fused_update(self, lo: int, hi: int, grad_w: torch.Tensor) -> bool: | |
| # A subclass with a one-launch fused update (Triton) returns True after applying it; | |
| # the base path always returns False so the caller runs the torch tile update. | |
| return False | |
| def _eff_lr(self) -> float: | |
| return self.lr * self._lr_mult | |
| def _decimation_apply(self) -> bool: | |
| """Advance the per-layer decimation clock; return whether to apply the update this step. | |
| Off (default) -> always fire at lr*1. On -> fire every _dec_period steps with lr scaled by | |
| the period (sum of r similar grads ~ r*grad). | |
| BIAS NOTE: dropping r-1 gradients and applying r*lr on the firing step is unbiased ONLY | |
| under stationarity (E[G] constant over the window); under non-stationary grads it is | |
| biased. Also the RMS-EMA (v) sees one update per period instead of per step, so its | |
| effective time-constant silently lengthens with _dec_period. Tested only on a stationary | |
| teacher, where both effects are invisible. Off by default; opt-in for near-stable layers.""" | |
| if not self.decimate_updates: | |
| self._lr_mult = 1.0 | |
| return True | |
| self._dec_since += 1 | |
| if self._dec_since >= self._dec_period: | |
| self._lr_mult = float(self._dec_period) | |
| self._dec_since = 0 | |
| return True | |
| return False | |
| def _decimation_observe(self) -> None: | |
| """Set the next update period from the flip-rate just observed (smaller flips -> rarer).""" | |
| if not self.decimate_updates: | |
| return | |
| r = self._dec_flip_rate | |
| # Data-parallel lockstep guard. Flips are counted on the ALREADY-averaged grad_w (and, in | |
| # proxy mode, an averaged proxy_gsq), so the per-rank flip-rate is normally identical and | |
| # the period stays synced on its own -> the `if fire` collective is symmetric. This MAX | |
| # all-reduce makes that invariant explicit (follow the most-active rank) so the period can | |
| # never desync even if a future change reintroduces a per-rank term into the flip count. | |
| # Cheap: one scalar, only on a `fire` step (itself symmetric) under world_size>1. | |
| if dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1: | |
| rt = torch.tensor(float(r), device=self.state.device, dtype=torch.float32) | |
| dist.all_reduce(rt, op=dist.ReduceOp.MAX) | |
| r = float(rt.item()) | |
| self._dec_period = 1 if r > 1e-3 else 2 if r > 1e-4 else 4 if r > 1e-5 else 8 | |
| def _build_t_cache(self) -> None: | |
| dt = torch.float16 if self.cache_mode == "fp16" else torch.int8 | |
| t, _ = self._decode_rows(0, self.out_features) | |
| cache = t.to(dt) | |
| if hasattr(self, "_t_cache"): | |
| self._t_cache.copy_(cache) | |
| else: | |
| self.register_buffer("_t_cache", cache, persistent=False) | |
| def _visible_t(self, dtype: torch.dtype) -> torch.Tensor: | |
| # the visible ternary weight T, from the derived cache when enabled (no decode), else decoded. | |
| if self.cache_mode != "none": | |
| if not hasattr(self, "_t_cache"): | |
| self._build_t_cache() | |
| return self._t_cache.to(dtype) | |
| t, _ = self._decode_rows(0, self.out_features) | |
| return t.to(dtype) | |
| def _dense_weight(self, dtype: torch.dtype) -> torch.Tensor: | |
| return self.scale.to(dtype) * self._visible_t(dtype) | |
| def _decode_rows(self, lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor]: | |
| return decode_state(self.state[lo:hi], self.C) | |
| def _refresh_t_cache(self, lo: int, hi: int, t: torch.Tensor) -> None: | |
| if self.cache_mode != "none": | |
| if not hasattr(self, "_t_cache"): # state for these rows is already written -> snapshot | |
| self._build_t_cache() | |
| elif self.cache_patch == "flip": | |
| # write only the elements whose visible ternary actually changed; the persistent | |
| # buffer already holds the correct value everywhere else (bit-identical to "full"). | |
| view = self._t_cache[lo:hi] | |
| tt = t.to(view.dtype) | |
| changed = view != tt | |
| if changed.any(): | |
| view[changed] = tt[changed] | |
| self.cache_patches.add_(int(changed.sum().item())) | |
| else: | |
| self._t_cache[lo:hi] = t.to(self._t_cache.dtype) | |
| def _write_rows(self, lo: int, hi: int, t: torch.Tensor, c: torch.Tensor) -> None: | |
| self.state[lo:hi].copy_(encode_state(t, c, self.C)) | |
| self._refresh_t_cache(lo, hi, t) | |
| def _update_tile(self, lo, hi, grad_w, t_i, c_i, s_i, proxy_gsq=None) -> None: | |
| if self.local_grad_clip > 0: | |
| row_norm = grad_w.norm(dim=1, keepdim=True).clamp_min(1e-30) | |
| grad_w = grad_w * (self.local_grad_clip / row_norm).clamp_max(1.0) | |
| # Learn one scale per output row; normalize by sqrt(fan_in) for width stability. | |
| grad_s = (grad_w * t_i.float()).sum(dim=1, keepdim=True) / math.sqrt(self.in_features) | |
| s_new = (s_i - self.lr_scale * grad_s).clamp_(1e-5, 10.0) | |
| update_signal = ( | |
| grad_w if self.pulse_mode == "direct" else ternary_gradient_unbiased(grad_w) | |
| ) | |
| # c stores a pending update in units of s/C; rebase it when the row scale changes. | |
| c_rebased = c_i.float() * (s_i / s_new) | |
| ticks = (-self._eff_lr() * update_signal) * (self.C / s_new) | |
| cc = stochastic_round(c_rebased + ticks) | |
| carry = torch.trunc(cc / self.C) | |
| remainder = cc - carry * self.C | |
| proposed_t = t_i.float() + carry | |
| new_t = proposed_t.clamp_(-1, 1) | |
| blocked = proposed_t != new_t | |
| remainder = torch.where( | |
| blocked, torch.sign(cc) * (self.C - 1), remainder | |
| ).clamp_(-(self.C - 1), self.C - 1) | |
| self._write_rows(lo, hi, new_t, remainder) | |
| self.scale[lo:hi].copy_(s_new) | |
| flips = int((new_t != t_i).sum().item()) | |
| self._dec_flips_step += flips # accumulate across tiles -> per-layer flip-rate | |
| self._dec_elems_step += new_t.numel() | |
| self.update_events.add_(int((cc != c_i).sum().item())) | |
| self.weight_flips.add_(flips) | |
| def apply_update_from_grad_w(self, grad_w: torch.Tensor) -> None: | |
| """Drive ONE full-matrix counter update from an externally-supplied grad_w [out,in]. | |
| For callers that compute grad_w outside this layer's own autograd Function -- e.g. the | |
| grouped Counter-MoE, where every expert's weight gradient comes from one grouped GEMM, so | |
| the per-expert forward Function is bypassed. Routes through the same fused/torch update path | |
| as the in-backward self-update (so the dynamics are identical), at full-matrix granularity. | |
| Honors training/update_enabled; does NOT apply decimation or the DDP grad_w all-reduce (the | |
| grouped path is single-process / dense-update by construction).""" | |
| if not (self.training and self.update_enabled): | |
| return | |
| lo, hi = 0, self.out_features | |
| t_i, c_i = self._decode_rows(lo, hi) | |
| if not self._fused_update(lo, hi, grad_w): | |
| self._update_tile(lo, hi, grad_w, t_i, c_i, self.scale[lo:hi]) | |
| def state_statistics(self) -> dict[str, float]: | |
| t, c = decode_state(self.state, self.C) | |
| return { | |
| "minus": float((t == -1).float().mean()), | |
| "zero": float((t == 0).float().mean()), | |
| "plus": float((t == 1).float().mean()), | |
| "counter_abs_mean": float(c.float().abs().mean()), | |
| "counter_edge": float((c.abs() == self.C - 1).float().mean()), | |
| "scale_mean": float(self.scale.mean()), | |
| } | |
| class RMSCounterLinear(CompactCounterLinear): | |
| """CompactCounterLinear + per-row RMS adaptive scaling (the cheap analogue of Adam's | |
| variance term). A per-output-row second moment v (O(out_features), negligible memory) | |
| normalizes the gradient before it drives the counter. This closes most of the gap to | |
| AdamW that vanilla counter-SGD leaves open.""" | |
| def __init__(self, *args, rms_beta: float = 0.9, rms_eps: float = 1e-3, | |
| use_rms: bool = True, rms_mode: str = "exact", | |
| scale_rebase: str = "eager", **kw) -> None: | |
| super().__init__(*args, **kw) | |
| self.rms_beta = float(rms_beta) | |
| self.rms_eps = float(rms_eps) | |
| self.use_rms = bool(use_rms) | |
| # rms_mode: "exact" uses the freshly-updated v for the denominator (the tick depends on a | |
| # row-stat of THIS step's grad -> two passes); "lagged" uses last step's v so the tick | |
| # needs no row-stat of the current grad -> one pass (the strict update-from-IO kernel). | |
| # scale_rebase: "eager" rebases the counter to s_new before ticking (needs s_new first); | |
| # "lazy" keeps a per-row calibration scale s_base and rebases at the next step's read, | |
| # so the tick uses the current scale only. lagged+lazy together are the one-pass update. | |
| # "proxy": post-LayerNorm E[r_o^2|D] ~ ||D_o||^2, so the RMS denominator comes from cheap | |
| # grad_out row-norms (x a scalar activation-energy) instead of the full ||G_o||^2 stat. | |
| assert rms_mode in ("exact", "lagged", "proxy") | |
| assert scale_rebase in ("eager", "lazy") | |
| self.rms_mode = rms_mode | |
| self.scale_rebase = scale_rebase | |
| self.register_buffer("v", torch.zeros((self.out_features, 1), dtype=torch.float32)) | |
| self.register_buffer("s_base", self.scale.clone()) # scale the counter is calibrated to | |
| def _update_tile(self, lo, hi, grad_w, t_i, c_i, s_i, proxy_gsq=None) -> None: | |
| if self.use_rms: | |
| # proxy mode takes the row second-moment from grad_out norms (passed in) instead of | |
| # the full ||G_o||^2 reduction; lagged uses the previous v, else the freshly-updated v. | |
| g_sq = proxy_gsq if (self.rms_mode == "proxy" and proxy_gsq is not None) \ | |
| else grad_w.pow(2).mean(dim=1, keepdim=True) | |
| if self.rms_mode == "lagged": | |
| denom = self.v[lo:hi].sqrt().clamp_min(self.rms_eps) # previous v | |
| self.v[lo:hi].mul_(self.rms_beta).add_(g_sq, alpha=1.0 - self.rms_beta) | |
| else: | |
| self.v[lo:hi].mul_(self.rms_beta).add_(g_sq, alpha=1.0 - self.rms_beta) | |
| denom = self.v[lo:hi].sqrt().clamp_min(self.rms_eps) # freshly-updated v | |
| grad_eff = grad_w / denom | |
| else: | |
| grad_eff = grad_w | |
| if self.local_grad_clip > 0: | |
| row_norm = grad_eff.norm(dim=1, keepdim=True).clamp_min(1e-30) | |
| grad_eff = grad_eff * (self.local_grad_clip / row_norm).clamp_max(1.0) | |
| # scale is still learned from the RAW gradient (its statistics are not normalised). | |
| grad_s = (grad_w * t_i.float()).sum(dim=1, keepdim=True) / math.sqrt(self.in_features) | |
| s_new = (s_i - self.lr_scale * grad_s).clamp_(1e-5, 10.0) | |
| update_signal = ( | |
| grad_eff if self.pulse_mode == "direct" else ternary_gradient_unbiased(grad_eff) | |
| ) | |
| if self.scale_rebase == "lazy": | |
| # counter is calibrated to s_base; bring it to the current scale s_i, tick in s_i units | |
| # (no dependence on s_new), and record that the stored counter is now calibrated to s_i. | |
| c_cur = stochastic_round(c_i.float() * (self.s_base[lo:hi] / s_i)) | |
| ticks = (-self._eff_lr() * update_signal) * (self.C / s_i) | |
| cc = stochastic_round(c_cur + ticks) | |
| self.s_base[lo:hi].copy_(s_i) | |
| return self._finish_update(lo, hi, cc, t_i, c_i, s_new) | |
| c_rebased = c_i.float() * (s_i / s_new) | |
| ticks = (-self._eff_lr() * update_signal) * (self.C / s_new) | |
| cc = stochastic_round(c_rebased + ticks) | |
| self._finish_update(lo, hi, cc, t_i, c_i, s_new) | |
| def _finish_update(self, lo, hi, cc, t_i, c_i, s_new) -> None: | |
| # carry/remainder -> ternary flip + residual, then write state + scale (shared by the | |
| # eager and the lazy-rebase paths). | |
| carry = torch.trunc(cc / self.C) | |
| remainder = cc - carry * self.C | |
| proposed_t = t_i.float() + carry | |
| new_t = proposed_t.clamp_(-1, 1) | |
| blocked = proposed_t != new_t | |
| remainder = torch.where( | |
| blocked, torch.sign(cc) * (self.C - 1), remainder | |
| ).clamp_(-(self.C - 1), self.C - 1) | |
| self._write_rows(lo, hi, new_t, remainder) | |
| self.scale[lo:hi].copy_(s_new) | |
| flips = int((new_t != t_i).sum().item()) | |
| self._dec_flips_step += flips # accumulate across tiles -> per-layer flip-rate | |
| self._dec_elems_step += new_t.numel() | |
| self.update_events.add_(int((cc != c_i).sum().item())) | |
| self.weight_flips.add_(flips) | |