File size: 7,697 Bytes
6e802a4 | 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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | # test_cached_decode.py: The cached-state decode matches the recompute path.
import os
import sys
from pathlib import Path
import pytest
import torch
REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO))
sys.path.insert(0, str(REPO / "scripts"))
from model import AttentionMixer
class _Params:
def __init__(self, seqlen_offset=0):
self.seqlen_offset = seqlen_offset
self.key_value_memory_dict = {}
def _mixer(d=128):
torch.manual_seed(0)
return AttentionMixer(d_model=d, head_dim=64, device="cpu", dtype=torch.float32,
layer_idx=3).eval()
def test_uncached_forward_is_unchanged_and_prefill_matches_it():
m = _mixer()
x = torch.randn(1, 11, 128)
with torch.no_grad():
ref = m(x)
pre = m(x, inference_params=_Params(0))
assert torch.allclose(ref, pre, atol=1e-6)
def test_stepwise_decode_equals_full_forward():
m = _mixer()
x = torch.randn(1, 13, 128)
with torch.no_grad():
ref = m(x)
p = _Params(0)
outs = [m(x[:, :5], inference_params=p)]
p.seqlen_offset = 5
for i in range(5, 13):
outs.append(m(x[:, i:i + 1], inference_params=p))
p.seqlen_offset = i + 1
step = torch.cat(outs, dim=1)
assert torch.allclose(ref, step, atol=1e-5), (ref - step).abs().max()
assert p.key_value_memory_dict[3][0].shape[2] == 13
def test_multi_position_step_onto_cache_is_causal_and_correct():
m = _mixer()
x = torch.randn(1, 10, 128)
with torch.no_grad():
ref = m(x)
p = _Params(0)
a = m(x[:, :4], inference_params=p); p.seqlen_offset = 4
b = m(x[:, 4:], inference_params=p)
assert torch.allclose(ref, torch.cat([a, b], dim=1), atol=1e-5)
def test_layer_idx_is_required_for_the_cache():
m = AttentionMixer(d_model=128, head_dim=64, device="cpu", dtype=torch.float32)
with pytest.raises(RuntimeError):
m(torch.randn(1, 3, 128), inference_params=_Params(0))
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA (mamba_ssm kernels)")
class TestDepthGenerateIdentityGPU:
@pytest.fixture(scope="class")
def tdc(self):
import test_depth_checkpoint as t
return t
def _model(self, hybrid, tdc):
from model import MambaCoTModel
torch.manual_seed(0)
m = MambaCoTModel(model_name="state-spaces/mamba2-130m", device="cuda",
dtype=torch.float32, mtp_num_heads=0)
if hybrid:
m.enable_hybrid_attention(top_k=2)
m.enable_depth_module(num_levels=tdc.NUM_LEVELS, codebook_size=64,
d_depth=64, depth_layers=1, depth_arch="gru")
return m.eval()
def _prefix(self, m):
torch.manual_seed(1)
text = torch.randint(5, 1000, (1, 9))
return torch.cat([text, torch.tensor([[m.token_registry.audio_start_id]])], dim=1).cuda()
def _run(self, m, tdc, cached, sg=None, n=24):
torch.manual_seed(2)
with torch.no_grad():
return tdc.depth_generate(
m, self._prefix(m), m.token_registry, max_frames=n, min_frames=n,
allow_stop=False, feedback="all", cross_frame_decode=False,
device="cuda", cached=cached, silence_guard=sg)
@pytest.mark.parametrize("hybrid", [False, True], ids=["pure", "hybrid"])
def test_cached_equals_recompute_greedy(self, tdc, hybrid):
m = self._model(hybrid, tdc)
a = self._run(m, tdc, cached=False)
b = self._run(m, tdc, cached=True)
assert a is not None and b is not None
assert a == b, f"first divergence at frame {next(i for i in range(len(a[0])) if any(a[k][i] != b[k][i] for k in range(len(a))))}"
def test_rewind_reprefill_matches_recompute(self, tdc):
m = self._model(False, tdc)
base = self._run(m, tdc, cached=False)
pause = sorted(set(base[0][:6]))
mk = lambda: tdc.SilenceGuard(codes=pause, budget=2, stop_after=99,
retries=1, retry_temp=0.0)
g0, g1 = mk(), mk()
a = self._run(m, tdc, cached=False, sg=g0)
b = self._run(m, tdc, cached=True, sg=g1)
if g0.stats().get("rewinds", 0) == 0:
pytest.skip("guard did not fire on this init; rewind path not exercised")
assert g0.stats() == g1.stats()
assert a == b
def test_cached_plan_injection_matches_recompute_teacher_forced(self, tdc):
m = self._model(False, tdc)
m.enable_plan_injection()
m.eval()
torch.manual_seed(3)
torch.nn.init.normal_(m.plan_injector.proj.weight, std=0.05)
torch.nn.init.normal_(m.plan_injector.proj.bias, std=0.05)
reg = m.token_registry
prefix = self._prefix(m)
P, n, NL = prefix.shape[1], 24, tdc.NUM_LEVELS
sizes = m.plan_injector.bin_sizes
g = torch.Generator().manual_seed(4)
bins = torch.stack([torch.randint(0, s, (3,), generator=g) for s in sizes],
dim=-1).unsqueeze(0).cuda()
sched = tdc.build_decode_plan_schedule([6, 9, 9], prefix_len=P, max_frames=n,
mode="pack", device="cuda")
backbone = m.backbone.backbone
emb = backbone.embedding
gen = torch.Generator(device="cuda").manual_seed(7)
frames = [torch.randint(0, 64, (NL,), generator=gen, device="cuda") for _ in range(n)]
def ids_codes(t):
ids = torch.cat([prefix, torch.full((1, t), reg.audio_start_id,
dtype=torch.long, device="cuda")], 1)
codes = torch.full((1, P + t, NL), -1, dtype=torch.long, device="cuda")
for i in range(t):
codes[0, P + i] = frames[i]
return ids, codes
def stage(ids, codes, lo, plan):
emb.set_codes(codes[:, lo:])
if plan:
emb.set_plan(m.plan_injector(bins, sched[:, :ids.shape[1]])[:, lo:])
def run(plan):
re, ca = [], []
with torch.no_grad():
for t in range(n):
ids, codes = ids_codes(t)
stage(ids, codes, 0, plan)
re.append(backbone(ids)[:, -1, :].float().clone())
ip = tdc._fresh_inference_cache(P + n + 8); n_fed = 0
for t in range(n):
ids, codes = ids_codes(t)
lo = n_fed
stage(ids, codes, lo, plan)
ip.seqlen_offset = n_fed
h = backbone(ids[:, lo:], inference_params=ip)
n_fed = ids.shape[1]
ca.append(h[:, -1, :].float().clone())
lm = sum(int(m.backbone.lm_head(a).argmax()) == int(m.backbone.lm_head(b).argmax())
for a, b in zip(re, ca))
dp = sum(list(m.depth_module.generate_frame(a)) == list(m.depth_module.generate_frame(b))
for a, b in zip(re, ca))
rel = [float((a - b).norm() / (a.norm() + 1e-9)) for a, b in zip(re, ca)]
return re, rel, lm, dp
with_plan, rel, lm, dp = run(True)
assert max(rel) < 5e-3, f"plan path: hidden rel {max(rel):.2e} (per frame {rel})"
assert lm == n and dp == n, f"plan path: lm-head {lm}/{n}, depth {dp}/{n}"
no_plan, rel0, _, _ = run(False)
moved = max(float((a - b).norm() / (a.norm() + 1e-9)) for a, b in zip(with_plan, no_plan))
assert moved > 1e-2, f"the randomised injector moved the hidden state by only {moved:.2e}"
assert max(rel0) < 5e-3
|