Spaces:
Runtime error
Runtime error
File size: 2,795 Bytes
cc15101 | 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 | """
Phase E self-checks. From the project root:
python -m ml.model.test_model
"""
import math
import torch
from .config import AbyssConfig
from .norm import AbyssRMSNorm
from .mlp import AbyssSwiGLU
from .block import AbyssBlock
from .model import TheAbyssModel
def tiny():
# small enough to run on CPU in a blink
return AbyssConfig(vocab_size=256, block_size=32, n_layer=2, n_head=4, n_embd=64)
def test_rmsnorm():
out = AbyssRMSNorm(64)(torch.randn(2, 16, 64) * 5 + 3)
out_rms = out.pow(2).mean(-1).sqrt()
assert torch.allclose(out_rms, torch.ones_like(out_rms), atol=1e-4), out_rms
print("rmsnorm ok: outputs have unit RMS")
def test_swiglu_shape():
cfg = tiny()
x = torch.randn(2, 16, cfg.n_embd)
y = AbyssBlock(cfg)(x)
assert y.shape == x.shape
assert not torch.allclose(y, x) # the block actuallt did something
print("block ok:", tuple(y.shape))
def test_block_shape():
cfg = tiny()
x = torch.randn(2, 16, cfg.n_embd)
y = AbyssBlock(cfg)(x)
assert y.shape == x.shape
assert not torch.allclose(y, x)
print("block ok:", tuple(y.shape))
def test_forward_and_loss():
cfg = tiny()
idx = torch.randint(0, cfg.vocab_size, (2, 16))
targets = torch.randint(0, cfg.vocab_size, (2, 16))
logits, loss = TheAbyssModel(cfg)(idx, targets)
assert logits.shape == (2, 16, cfg.vocab_size), logits.shape
assert loss.ndim == 0 and torch.isfinite(loss)
# an untrained model should sit near ln(vocab): uniform guessing
expected = math.log(cfg.vocab_size)
assert abs(loss.item() - expected) < 0.7, (loss.items(), expected)
print(f"forward ok: loss {loss.item():.3f} ~ ln(V) {expected:.3f}")
def test_backward():
cfg = tiny()
model = TheAbyssModel(cfg)
idx = torch.randint(0, cfg.vocab_size, (2, 16))
_, loss = model(idx, idx)
loss.backward()
assert all(p.grad is not None for p in model.parameters()), "a param got no gradient"
print("backward ok: every parameter has a gradient")
def test_block_size_gaurd():
cfg = tiny()
model = TheAbyssModel(cfg)
try:
model(torch.randint(0, cfg.vocab_size, (1, cfg.block_size + 1)))
except AssertionError:
print("guard ok: rejects sequences longer than block_size")
return
raise SystemExit("block_size guard did not fire")
if __name__ == "__main__":
test_rmsnorm()
test_swiglu_shape()
test_block_shape()
test_forward_and_loss()
test_backward()
test_block_size_gaurd()
full = TheAbyssModel(AbyssConfig())
print(f"all E tests passed - full model has {full.num_params():,} params "
f"({full.num_params(non_embedding=True):,} non-embedding)") |