File size: 3,365 Bytes
76b78ee | 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 | """Unit tests for multi-token prediction (Meta MTP) support.
Run: .venv/bin/python tests/test_mtp.py
"""
import subprocess
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import torch
import torch.nn.functional as F
from model.config import TinyLiquidConfig, CONFIGS
from model.tiny_liquid import TinyLiquid
ROOT = Path(__file__).resolve().parents[1]
def make_model(vocab=512, mtp=2):
cfg = TinyLiquidConfig(vocab_size=vocab, mtp_heads=mtp, **CONFIGS["micro6m"])
return TinyLiquid(cfg)
def test_forward_mtp_shapes():
m = make_model()
ids = torch.randint(0, 512, (2, 16))
logits, aux = m.forward_mtp(ids)
assert tuple(logits.shape) == (2, 16, 512)
assert len(aux) == 2 and all(tuple(a.shape) == (2, 16, 512) for a in aux)
# main forward unchanged
assert tuple(m(ids).shape) == (2, 16, 512)
def test_mtp_loss_backward():
m = make_model()
ids = torch.randint(0, 512, (2, 24))
logits, aux = m.forward_mtp(ids)
loss = F.cross_entropy(logits.reshape(-1, 512), ids.reshape(-1))
for k, a in enumerate(aux):
off = k + 2
loss = loss + 0.1 * F.cross_entropy(
a[:, :-off].reshape(-1, 512), ids[:, off:].reshape(-1))
loss.backward()
assert m.mtp_heads[0][0].weight.grad is not None
assert m.tok_emb.weight.grad is not None
assert torch.isfinite(loss)
def test_mtp_checkpoint_roundtrip():
m = make_model()
sd = {"config": m.cfg.__dict__, "model": m.state_dict()}
m2 = TinyLiquid(TinyLiquidConfig(**sd["config"]))
missing, unexpected = m2.load_state_dict(sd["model"], strict=True)
assert not missing and not unexpected
def test_train_lm_mtp_smoke(tmp=None):
tmp = Path(tmp or (ROOT / "data" / "_mtp_smoke"))
tmp.mkdir(parents=True, exist_ok=True)
import numpy as np
rng = np.random.default_rng(0)
(tmp / "train.bin").write_bytes(rng.integers(1, 500, size=20000, dtype=np.uint16).tobytes())
(tmp / "valid.bin").write_bytes(rng.integers(1, 500, size=5000, dtype=np.uint16).tobytes())
ckpt = tmp / "ckpt"
cmd = [
sys.executable, "-u", "train/train_lm.py",
"--data", str(tmp / "train.bin"), "--val", str(tmp / "valid.bin"),
"--tok", "data/tokenizer.json", "--config", "micro6m",
"--ckpt", str(ckpt), "--batch", "2", "--seq", "32",
"--lr", "1e-4", "--warmup", "0", "--steps", "3",
"--eval-every", "2", "--save-every", "2", "--threads", "2",
"--mtp", "2", "--log-every", "1",
]
env = {"PYTHONPATH": str(ROOT)}
r = subprocess.run(cmd, capture_output=True, text=True, cwd=ROOT, env=env,
timeout=300)
assert r.returncode == 0, r.stderr[-1500:]
saved = sorted((ckpt).glob("*.pt"))
assert saved, "no checkpoint saved"
import torch as T
sd = T.load(saved[-1], map_location="cpu", weights_only=False)
assert sd["config"]["mtp_heads"] == 2
assert any("mtp_heads" in k for k in sd["model"])
for f in ["train.bin", "valid.bin"]:
(tmp / f).unlink()
for f in ckpt.glob("*.pt"):
f.unlink()
ckpt.rmdir()
tmp.rmdir()
if __name__ == "__main__":
fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
for fn in fns:
fn()
print(f"ok {fn.__name__}")
print(f"\n{len(fns)} mtp tests passed")
|