Spaces:
Runtime error
Runtime error
| """ | |
| 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)") |