File size: 4,375 Bytes
4397e12 | 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 | """Correctness tests for the model. Run: source env.sh && $TA_PY -m pytest -q tests/"""
import math
import torch
import torch.nn.functional as F
from tiny_agent.model import ModelConfig, TinyAgentLM, dense_reference_attention_mask, make_block_mask
from tiny_agent.optim import build_optimizers
DEV = "xpu" if torch.xpu.is_available() else "cpu"
def small_cfg(**kw):
base = dict(vocab_size=512, d_model=128, n_layers=4, n_heads=4, head_dim=32, n_kv_heads=2,
kv_group=2, swa_window=8, max_seq_len=256, engram_layers=(1,), engram_rows_per_head=1000,
engram_heads=2, engram_head_dim=16)
base.update(kw)
return ModelConfig(**base)
def random_model(cfg):
torch.manual_seed(0)
m = TinyAgentLM(cfg).to(DEV)
# zero-inits make outputs trivially constant; randomize everything for the tests
with torch.no_grad():
for p in m.parameters():
p.normal_(0, 0.05)
return m
def docs(B, T, cuts=(37, 80)):
d = torch.zeros(B, T, dtype=torch.long)
for c in cuts:
d[:, c:] += 1
return d.to(DEV)
def test_flex_matches_dense_reference():
torch.manual_seed(0)
B, H, Hk, T, D, W = 2, 4, 2, 128, 32, 8
q = torch.randn(B, H, T, D, device=DEV)
k = torch.randn(B, Hk, 2 * T, D, device=DEV)
v = torch.randn(B, Hk, 2 * T, D, device=DEV)
doc = docs(B, T)
from torch.nn.attention.flex_attention import flex_attention
y = flex_attention(q, k, v, block_mask=make_block_mask(doc, W), enable_gqa=True)
m = dense_reference_attention_mask(doc, W)
ref = F.scaled_dot_product_attention(q, k.repeat_interleave(H // Hk, 1), v.repeat_interleave(H // Hk, 1), attn_mask=m)
assert torch.allclose(y, ref, atol=2e-4, rtol=1e-3), (y - ref).abs().max()
def test_mask_semantics():
doc = docs(1, 64, cuts=(20,))
m = dense_reference_attention_mask(doc, 8)[0, 0]
T = 64
assert m[30, 25] and not m[30, 31] # global causal
assert not m[30, 10] # other document
assert m[30, T + 25] and not m[30, T + 22] # local window 8: 23..30 visible
assert not m[30, T + 31] # local causal
def test_causal_and_doc_isolation():
cfg = small_cfg()
m = random_model(cfg).eval()
B, T = 2, 128
idx = torch.randint(0, cfg.vocab_size, (B, T), device=DEV)
doc = docs(B, T)
with torch.no_grad():
a = m(idx, doc)
idx2 = idx.clone()
idx2[:, 60] = (idx2[:, 60] + 1) % cfg.vocab_size
b = m(idx2, doc)
assert torch.allclose(a[:, :60], b[:, :60], atol=1e-4) # nothing before t=60 changes
assert not torch.allclose(a[:, 60:80], b[:, 60:80]) # same doc after t changes
assert torch.allclose(a[:, 80:], b[:, 80:], atol=1e-4) # later document is isolated
def test_engram_addresses_deterministic_and_in_range():
cfg = small_cfg()
m = random_model(cfg)
e = m.engram_modules()[0]
ids = torch.randint(0, cfg.vocab_size, (2, 50), device=DEV)
a1, a2 = e.addresses(ids), e.addresses(ids)
assert torch.equal(a1, a2)
assert a1.min() >= 0 and a1.max() < e.table.num_embeddings
# bigram head address at t depends only on tokens t-1, t
ids2 = ids.clone()
ids2[:, 10] += 1
b = e.addresses(ids2)
n_heads = cfg.engram_heads
assert torch.equal(a1[:, :10], b[:, :10]) and torch.equal(a1[:, 12:, :n_heads], b[:, 12:, :n_heads])
def test_kv_reuse_groups():
cfg = small_cfg(n_layers=8, kv_group=4)
m = TinyAgentLM(cfg)
modes = [b.attn.mode for b in m.blocks]
assert modes == ["full", "reuse", "reuse", "reuse"] * 2
assert not hasattr(m.blocks[1].attn, "w_kv_global")
def test_overfit_one_batch():
cfg = small_cfg()
torch.manual_seed(0)
m = TinyAgentLM(cfg).to(DEV)
opts = build_optimizers(m, lr=1e-2)
idx = torch.randint(0, cfg.vocab_size, (4, 65), device=DEV)
doc = torch.zeros(4, 64, dtype=torch.long, device=DEV)
losses = []
for _ in range(60):
logits = m(idx[:, :-1], doc)
loss = F.cross_entropy(logits.reshape(-1, cfg.vocab_size), idx[:, 1:].reshape(-1))
loss.backward()
for o in opts:
o.step()
o.zero_grad(set_to_none=True)
losses.append(loss.item())
assert losses[0] > math.log(cfg.vocab_size) - 0.5
assert losses[-1] < 0.5 * losses[0], losses[::10]
|