Spaces:
Paused
Paused
File size: 2,188 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 | """Derived visible-weight compute cache (acceleration memo M5).
cache_mode keeps the visible ternary T in a GEMM-friendly dtype (fp16/int8) so forward/grad_x
never pay the 6-bit decode tax. T is the source of truth's *derived view*: forward through the
cache must equal forward through the decode (T = -1/0/1 is exact in fp16 and int8), and the cache
must track the state as visible weights flip during updates.
"""
import math
import pytest
import torch
from memory_native import RMSCounterLinear
from memory_native.counter import decode_state
from memory_native.packed import PackedRMSCounterLinear, unpack_codes
@pytest.mark.parametrize("cache_mode", ["fp16", "int8"])
def test_cache_forward_matches_decode(cache_mode):
torch.manual_seed(0)
ref = PackedRMSCounterLinear(64, 48, C=11)
cached = PackedRMSCounterLinear(64, 48, C=11, cache_mode=cache_mode)
cached.state.copy_(ref.state)
cached.scale.copy_(ref.scale)
cached._build_t_cache()
x = torch.randn(8, 5, 64)
with torch.no_grad():
assert torch.equal(ref(x), cached(x)) # T is exact in fp16/int8 -> bit-exact forward
@pytest.mark.parametrize("cache_mode", ["fp16", "int8"])
def test_cache_tracks_state_through_updates(cache_mode):
torch.manual_seed(0)
lay = PackedRMSCounterLinear(64, 48, C=11, cache_mode=cache_mode).train()
x = torch.randn(16, 64)
for _ in range(10):
lay(x).pow(2).mean().backward() # self-updates flip visible weights
t, _ = decode_state(unpack_codes(lay.state, 64), 11) # truth T from the packed state
assert torch.equal(lay._t_cache.to(torch.int16), t) # cache stayed in sync
def test_cache_layer_still_recovers_teacher():
torch.manual_seed(0)
n, N, C = 24, 256, 11
ts = 0.25
base = math.sqrt(3.0 / (2.0 * n))
tw = torch.randint(-1, 2, (n, n)).float()
x = torch.randn(N, n)
y = x @ (ts * tw).t()
lay = RMSCounterLinear(n, n, C=C, lr=0.02, lr_scale=2e-4, init_gain=ts / base,
cache_mode="int8").train()
for _ in range(500):
(lay(x) - y).pow(2).mean().backward()
with torch.no_grad():
assert (lay(x) - y).pow(2).mean().item() < 0.02
|