Spaces:
Paused
Paused
File size: 6,246 Bytes
6f75f79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """Packed 6-bit storage: the real sub-byte persistent footprint.
Each finite-state code needs only ceil(log2(63)) = 6 bits, so 4 codes pack into 3 bytes.
PackedRMSCounterLinear stores its state packed -- persistent state is 0.75 byte/weight, not
1.0 -- and unpacks at the storage boundary (forward weight, per-tile decode, per-tile write).
The update math is inherited unchanged from RMSCounterLinear via the storage hooks
(_dense_weight / _decode_rows / _write_rows).
This is pure PyTorch and matches the engine's packing bit-for-bit (kernels/compact_counter.cl
cc_pack4/cc_unpack4), so a counter model trained here exports to the same packed format.
"""
from __future__ import annotations
import torch
from .counter import RMSCounterLinear, decode_state, encode_state
__all__ = ["pack_codes", "unpack_codes", "PackedRMSCounterLinear"]
def pack_codes(codes: torch.Tensor) -> torch.Tensor:
"""codes uint8 [out, in] in [0,63] -> packed uint8 [out, (in//4)*3] (4 codes / 3 bytes)."""
out, in_ = codes.shape
assert in_ % 4 == 0, "in_features must be divisible by 4 for 6-bit packing"
c = codes.to(torch.int32).reshape(out, in_ // 4, 4)
c0, c1, c2, c3 = c[..., 0], c[..., 1], c[..., 2], c[..., 3]
p0 = (c0 & 0x3F) | ((c1 & 0x03) << 6)
p1 = ((c1 >> 2) & 0x0F) | ((c2 & 0x0F) << 4)
p2 = ((c2 >> 4) & 0x03) | ((c3 & 0x3F) << 2)
packed = torch.stack([p0, p1, p2], dim=-1) # [out, in//4, 3]
return (packed & 0xFF).to(torch.uint8).reshape(out, (in_ // 4) * 3)
def unpack_codes(packed: torch.Tensor, in_features: int) -> torch.Tensor:
"""packed uint8 [out, (in//4)*3] -> codes uint8 [out, in] in [0,63]."""
out = packed.shape[0]
gpr = in_features // 4
p = packed.to(torch.int32).reshape(out, gpr, 3)
p0, p1, p2 = p[..., 0], p[..., 1], p[..., 2]
c0 = p0 & 0x3F
c1 = ((p0 >> 6) | (p1 << 2)) & 0x3F
c2 = ((p1 >> 4) | (p2 << 4)) & 0x3F
c3 = (p2 >> 2) & 0x3F
codes = torch.stack([c0, c1, c2, c3], dim=-1) # [out, gpr, 4]
return codes.to(torch.uint8).reshape(out, in_features)
class PackedRMSCounterLinear(RMSCounterLinear):
"""RMSCounterLinear whose persistent `state` buffer is packed to 6 bits (0.75 B/weight).
Same learning dynamics as RMSCounterLinear; only the storage layout differs. Codes are
unpacked transiently for the GEMM and the per-tile update, then repacked, so the resident
state stays at 0.75 byte/weight.
"""
def __init__(self, *args, **kw) -> None:
super().__init__(*args, **kw) # builds unpacked state [out, in], scale, v
codes = self.state # uint8 [out, in]
del self._buffers["state"]
self.register_buffer("state", pack_codes(codes)) # [out, (in//4)*3]
self._sr_step = 0 # per-call seed for the kernel's deterministic stochastic rounding
def _fused_update(self, lo: int, hi: int, grad_w: torch.Tensor) -> bool:
"""One-launch Triton RMS+SR update for a packed row range.
Originally the fused kernel only fired for the full matrix (`lo==0, hi==out`). That made
the *fast* path and the *low-peak* path mutually exclusive: `tile_rows>0` used cuBLAS for a
small grad_w tile but then fell back to the slow torch transition. The kernel is row-local,
so a contiguous row slice of the packed state is sufficient.
This gives a practical middle ground between:
* full fast path: materialize one [out,in] grad_w and fuse the transition;
* strict from-IO: no grad_w but a hand-written GEMM (very slow on T4);
* tiled fast path: materialize only [tile_rows,in] grad_w, then fuse the transition.
The deterministic hash-SR seed is advanced per tile, so different row tiles do not reuse the
same rounding stream even though the Triton kernel sees local row indices 0..tile_rows-1.
"""
from .fused_update import HAS_TRITON, triton_counter_update
if not (HAS_TRITON and grad_w.is_cuda and self.use_rms
and self.pulse_mode == "direct" and self.local_grad_clip == 0
and self.rms_mode == "exact" and self.scale_rebase == "eager"
and not self.decimate_updates): # kernel doesn't report flip-rate; torch path does
return False
if grad_w.shape != (hi - lo, self.in_features):
return False
seed = self._sr_step
self._sr_step += 1
triton_counter_update(self.state[lo:hi], self.scale[lo:hi], self.v[lo:hi], grad_w.contiguous(),
C=self.C, lr=self.lr, lr_scale=self.lr_scale,
rms_beta=self.rms_beta, rms_eps=self.rms_eps, seed=seed)
# The kernel mutates the packed state directly, bypassing _write_rows/_refresh_t_cache.
# Refresh only the affected rows instead of rebuilding the whole derived T cache.
if self.cache_mode != "none":
t_new, _ = self._decode_rows(lo, hi)
self._refresh_t_cache(lo, hi, t_new)
return True
# storage hooks ------------------------------------------------------------
def _all_codes(self) -> torch.Tensor:
return unpack_codes(self.state, self.in_features)
# _dense_weight is inherited from the base: it routes through _visible_t (derived cache when
# cache_mode != "none", else _decode_rows -> the packed unpack below). No override needed.
def _decode_rows(self, lo: int, hi: int):
codes = unpack_codes(self.state[lo:hi], self.in_features)
return decode_state(codes, self.C)
def _write_rows(self, lo: int, hi: int, t: torch.Tensor, c: torch.Tensor) -> None:
self.state[lo:hi].copy_(pack_codes(encode_state(t, c, self.C)))
self._refresh_t_cache(lo, hi, t)
@torch.no_grad()
def state_statistics(self) -> dict[str, float]:
t, c = decode_state(self._all_codes(), 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()),
}
|