125 lines
4.2 KiB
Python
125 lines
4.2 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from microgpt.model import GPT, GPTConfig, apply_rope, rope_tables
|
|
|
|
CFG = GPTConfig(vocab_size=50, block_size=32, n_layer=2, n_head=4, n_embd=32)
|
|
|
|
|
|
@pytest.fixture()
|
|
def model():
|
|
torch.manual_seed(0)
|
|
return GPT(CFG).eval()
|
|
|
|
|
|
def test_output_shapes(model):
|
|
idx = torch.randint(0, CFG.vocab_size, (3, 10))
|
|
logits, loss, cache = model(idx)
|
|
assert logits.shape == (3, 10, CFG.vocab_size)
|
|
assert loss is None
|
|
assert len(cache) == CFG.n_layer and cache[0][0].shape == (3, CFG.n_head, 10, 8)
|
|
|
|
|
|
def test_weights_are_tied(model):
|
|
assert model.lm_head.weight is model.tok_emb.weight
|
|
|
|
|
|
def test_no_future_leakage(model):
|
|
"""Changing a later token must not change logits at earlier positions."""
|
|
a = torch.randint(0, CFG.vocab_size, (1, 12))
|
|
b = a.clone()
|
|
b[0, 8:] = (b[0, 8:] + 1) % CFG.vocab_size
|
|
with torch.no_grad():
|
|
la, lb = model(a)[0], model(b)[0]
|
|
assert torch.allclose(la[:, :8], lb[:, :8], atol=1e-6)
|
|
assert not torch.allclose(la[:, 8:], lb[:, 8:], atol=1e-6)
|
|
|
|
|
|
def test_kv_cache_matches_full_forward(model):
|
|
"""Prefill + one-token-at-a-time decoding == a single full forward pass."""
|
|
idx = torch.randint(0, CFG.vocab_size, (2, 20))
|
|
with torch.no_grad():
|
|
full, _, _ = model(idx)
|
|
logits, _, cache = model(idx[:, :7])
|
|
steps = [logits]
|
|
for t in range(7, 20):
|
|
step, _, cache = model(idx[:, t : t + 1], cache=cache)
|
|
steps.append(step)
|
|
assert torch.allclose(full, torch.cat(steps, dim=1), atol=1e-5)
|
|
|
|
|
|
def test_cache_chunked_prefill_matches_full(model):
|
|
"""Multi-token continuation exercises the non-trivial causal-mask offset."""
|
|
idx = torch.randint(0, CFG.vocab_size, (1, 18))
|
|
with torch.no_grad():
|
|
full, _, _ = model(idx)
|
|
l1, _, cache = model(idx[:, :6])
|
|
l2, _, _ = model(idx[:, 6:], cache=cache)
|
|
assert torch.allclose(full, torch.cat([l1, l2], dim=1), atol=1e-5)
|
|
|
|
|
|
def test_greedy_generation_identical_with_and_without_cache(model):
|
|
prompt = torch.randint(0, CFG.vocab_size, (1, 5))
|
|
a = model.generate(prompt, 20, temperature=0, use_cache=True)
|
|
b = model.generate(prompt, 20, temperature=0, use_cache=False)
|
|
assert torch.equal(a, b)
|
|
|
|
|
|
def test_generation_past_block_size(model):
|
|
prompt = torch.randint(0, CFG.vocab_size, (1, 5))
|
|
out = model.generate(prompt, CFG.block_size * 2, temperature=1.0, top_k=10)
|
|
assert out.shape == (1, 5 + CFG.block_size * 2)
|
|
assert int(out.max()) < CFG.vocab_size
|
|
|
|
|
|
def test_sequence_longer_than_block_size_is_rejected(model):
|
|
with pytest.raises(ValueError):
|
|
model(torch.zeros(1, CFG.block_size + 1, dtype=torch.long))
|
|
|
|
|
|
def test_seeded_sampling_is_reproducible(model):
|
|
prompt = torch.randint(0, CFG.vocab_size, (1, 4))
|
|
runs = [
|
|
model.generate(prompt, 15, top_p=0.9, generator=torch.Generator().manual_seed(7))
|
|
for _ in range(2)
|
|
]
|
|
assert torch.equal(runs[0], runs[1])
|
|
|
|
|
|
def test_rope_is_norm_preserving():
|
|
cos, sin = rope_tables(16, 40, 10000.0)
|
|
x = torch.randn(2, 3, 40, 16)
|
|
y = apply_rope(x, cos, sin)
|
|
assert torch.allclose(x.norm(dim=-1), y.norm(dim=-1), atol=1e-5)
|
|
|
|
|
|
def test_rope_scores_depend_only_on_relative_position():
|
|
"""<rope(q, m), rope(k, n)> must equal <rope(q, m+s), rope(k, n+s)>."""
|
|
cos, sin = rope_tables(16, 64, 10000.0)
|
|
q, k = torch.randn(1, 1, 1, 16), torch.randn(1, 1, 1, 16)
|
|
|
|
def score(m: int, n: int) -> torch.Tensor:
|
|
qm = apply_rope(q, cos[m : m + 1], sin[m : m + 1])
|
|
kn = apply_rope(k, cos[n : n + 1], sin[n : n + 1])
|
|
return (qm * kn).sum()
|
|
|
|
assert torch.allclose(score(5, 2), score(25, 22), atol=1e-4)
|
|
assert torch.allclose(score(0, 9), score(30, 39), atol=1e-4)
|
|
assert not torch.allclose(score(5, 2), score(5, 3), atol=1e-4)
|
|
|
|
|
|
def test_can_overfit_a_tiny_batch():
|
|
torch.manual_seed(0)
|
|
m = GPT(CFG)
|
|
x = torch.randint(0, CFG.vocab_size, (4, 16))
|
|
y = torch.roll(x, -1, dims=1)
|
|
opt = torch.optim.AdamW(m.parameters(), lr=3e-3)
|
|
first = None
|
|
for _ in range(80):
|
|
loss = m(x, y)[1]
|
|
first = first if first is not None else loss.item()
|
|
opt.zero_grad()
|
|
loss.backward()
|
|
opt.step()
|
|
assert loss.item() < 0.3 * first
|