File size: 3,353 Bytes
60e7bb5 | 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 | import torch
from grimboost.config import DrafterConfig
from grimboost.model.transformer import GrimBoostLM
def tiny_config() -> DrafterConfig:
return DrafterConfig(
vocab_size=256,
hidden_size=64,
num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=2,
head_dim=16,
intermediate_size=128,
max_position_embeddings=64,
rope_theta=10_000.0,
)
def make_model() -> GrimBoostLM:
torch.manual_seed(1)
return GrimBoostLM(tiny_config()).eval()
def test_incremental_equals_full():
model = make_model()
ids = torch.randint(0, 256, (1, 12))
with torch.no_grad():
full = model(ids)
cache = model.new_cache(max_len=64, dtype=torch.float32)
_ = model(ids[:, :-1], past_key_values=cache, use_cache=True)
step = model(ids[:, -1:], past_key_values=cache, use_cache=True)
assert cache.length == 12
assert torch.allclose(step[:, -1], full[:, -1], atol=1e-5)
def test_multi_token_block_after_cache():
model = make_model()
ids = torch.randint(0, 256, (1, 14))
with torch.no_grad():
full = model(ids)
cache = model.new_cache(max_len=64, dtype=torch.float32)
_ = model(ids[:, :8], past_key_values=cache, use_cache=True)
block = model(ids[:, 8:], past_key_values=cache, use_cache=True)
assert cache.length == 14
assert torch.allclose(block, full[:, 8:], atol=1e-5)
def test_newest_suffix_rollback_is_exact():
model = make_model()
ids = torch.randint(0, 256, (1, 14))
alt = torch.randint(0, 256, (1, 3))
with torch.no_grad():
cache = model.new_cache(max_len=64, dtype=torch.float32)
_ = model(ids[:, :11], past_key_values=cache, use_cache=True)
_ = model(ids[:, 11:], past_key_values=cache, use_cache=True)
assert cache.length == 14
cache.rollback(3)
assert cache.length == 11
block = model(alt, past_key_values=cache, use_cache=True)
reference = model(torch.cat([ids[:, :11], alt], dim=1))
assert torch.allclose(block[:, -1], reference[:, -1], atol=1e-5)
def test_rebuild_with_positions_from_zero():
model = make_model()
ids = torch.randint(0, 256, (1, 10))
with torch.no_grad():
full = model(ids)
cache = model.new_cache(max_len=64, dtype=torch.float32)
rebuilt = model(ids, past_key_values=cache, use_cache=True)
assert cache.length == 10
assert torch.allclose(rebuilt, full, atol=1e-5)
def test_left_crop_is_not_equivalent():
"""Documents why the engine never left-crops a full-attention cache."""
model = make_model()
ids = torch.randint(0, 256, (1, 12))
with torch.no_grad():
full_suffix = model(ids[:, 5:], position_ids=torch.arange(5, 12).unsqueeze(0))
cache = model.new_cache(max_len=64, dtype=torch.float32)
_ = model(ids[:, :-1], past_key_values=cache, use_cache=True)
for layer in range(model.config.num_hidden_layers):
cache.k[layer] = cache.k[layer][:, :, 5:, :]
cache.v[layer] = cache.v[layer][:, :, 5:, :]
cache.length -= 5
cropped = model(ids[:, -1:], past_key_values=cache, use_cache=True,
position_ids=torch.tensor([[11]]))
assert not torch.allclose(cropped[:, -1], full_suffix[:, -1], atol=1e-4)
|