Spaces:
Paused
Paused
| """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 | |
| 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 | |
| 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 | |