外观
llm/test_llm.py
配套源码，运行方法见同目录 README。返回实验总览。
python
"""Run: python3 -m unittest discover -s labs/llm -p 'test_*.py' -v"""
import math
from pathlib import Path
import tempfile
import unittest
import numpy as np
from bpe import ByteBPE, CharacterTokenizer, merge_pair
from model import Config, TinyTransformer, AdamW, cross_entropy, make_batch, generate
from algorithms import (online_attention, exact_deduplicate, fit_power_law,
group_advantages, masked_sft_loss, clipped_policy_objective, dpo_loss)
class BPETests(unittest.TestCase):
def test_utf8_round_trip(self):
tokenizer = ByteBPE.train(["blue blue 蓝色", "red red 红色"], 20)
for text in ["", "蓝色 fox 🦊", "\n\t unseen café ", "blue blue"]:
self.assertEqual(tokenizer.decode(tokenizer.encode(text)), text)
def test_non_overlapping_merges(self):
self.assertEqual(merge_pair([1, 1, 1, 1, 1], (1, 1), 256), [256, 256, 1])
def test_determinism_and_document_boundary(self):
self.assertEqual(ByteBPE.train(["ab", "ac"], 1).merges, [(97, 98)])
self.assertEqual(ByteBPE.train(["a", "b"], 1).merges, [])
self.assertEqual(ByteBPE.train(["a b"], 1).merges, [])
def test_serialization(self):
tokenizer = ByteBPE.train(["abc abc abc"], 5)
with tempfile.TemporaryDirectory() as directory:
file = Path(directory) / "bpe.json"
tokenizer.save(file)
self.assertEqual(ByteBPE.load(file).encode("abc!"), tokenizer.encode("abc!"))
class TransformerTests(unittest.TestCase):
def setUp(self):
self.model = TinyTransformer(Config(7, dim=8, heads=2, hidden=12, layers=2, context=5))
self.x = np.array([[1, 2, 1, 3, 4], [2, 1, 3, 2, 5]])
self.y = np.array([[2, 1, 3, 4, 6], [1, 3, 2, 5, 6]])
def test_future_cannot_change_past(self):
changed = self.x.copy()
changed[:, 3:] = 6
original_logits = self.model.forward(self.x)[0]
changed_logits = self.model.forward(changed)[0]
np.testing.assert_allclose(original_logits[:, :3], changed_logits[:, :3], atol=1e-12)
self.assertGreater(np.max(np.abs(original_logits[:, 3:] - changed_logits[:, 3:])), 1e-5)
def test_shifted_batch(self):
x, y = make_batch(np.arange(10), 20, 4, np.random.default_rng(1))
np.testing.assert_array_equal(y, x + 1)
np.testing.assert_array_equal(x[:, 1:], y[:, :-1])
def test_uniform_and_extreme_loss(self):
loss, grad = cross_entropy(np.zeros((2, 5, 7)), self.y)
self.assertAlmostEqual(loss, math.log(7))
np.testing.assert_allclose(grad.sum(axis=-1), 0, atol=1e-15)
logits = np.array([[[1000., -1000.]]])
value, derivative = cross_entropy(logits, np.array([[1]]))
self.assertAlmostEqual(value, 2000.)
self.assertTrue(np.all(np.isfinite(derivative)))
def test_every_parameter_family_gradient(self):
_, grads = self.model.loss(self.x, self.y, backward=True)
rng = np.random.default_rng(8)
eps = 1e-5
# Four selected scalar entries in every parameter tensor, across two
# blocks; repeated embeddings test accumulation and residual paths.
for name, value in self.model.params.items():
for flat_index in rng.choice(value.size, min(4, value.size), replace=False):
index = np.unravel_index(flat_index, value.shape)
old = value[index]
value[index] = old + eps
plus = self.model.loss(self.x, self.y)
value[index] = old - eps
minus = self.model.loss(self.x, self.y)
value[index] = old
numeric = (plus - minus) / (2 * eps)
self.assertAlmostEqual(numeric, grads[name][index], delta=2e-5,
msg=f"{name}{index}: numerical={numeric}, analytic={grads[name][index]}")
def test_can_overfit_small_batch(self):
optimizer = AdamW(self.model.params, lr=.005, decay=0)
initial = self.model.loss(self.x, self.y)
for _ in range(60):
_, grads = self.model.loss(self.x, self.y, backward=True)
optimizer.step(grads)
self.assertLess(self.model.loss(self.x, self.y), initial * .3)
def test_checkpoint_and_generation(self):
tokenizer = CharacterTokenizer("abcdefg")
with tempfile.TemporaryDirectory() as directory:
file = Path(directory) / "model.npz"
self.model.save(file, tokenizer)
restored, tok = TinyTransformer.load(file)
np.testing.assert_array_equal(self.model.forward(self.x)[0], restored.forward(self.x)[0])
self.assertEqual(tok.encode("abc"), [0, 1, 2])
first = generate(restored, [1, 2], 12, seed=4)
second = generate(restored, [1, 2], 12, seed=4)
self.assertEqual(first, second)
self.assertEqual(len(first), 14)
def test_invalid_config(self):
with self.assertRaises(ValueError):
Config(7, dim=9, heads=3)
with self.assertRaises(ValueError):
generate(self.model, [], 5)
class AlgorithmTests(unittest.TestCase):
def test_sft_response_only(self):
logits = np.zeros((1, 4, 3))
labels = np.array([[0, 1, 2, 0]])
mask = np.array([[0, 0, 1, 1]])
loss, grad = masked_sft_loss(logits, labels, mask)
self.assertAlmostEqual(loss, math.log(3))
np.testing.assert_array_equal(grad[:, :2], 0)
changed = logits.copy()
changed[:, :2, 0] = 999
self.assertAlmostEqual(masked_sft_loss(changed, labels, mask)[0], loss)
with self.assertRaises(ValueError):
masked_sft_loss(logits, labels, np.zeros_like(mask))
def test_online_attention_matches_dense(self):
from model import causal_attention
rng = np.random.default_rng(21)
q, k, v = (rng.normal(size=(13, 8)) for _ in range(3))
dense = causal_attention(q[None, None], k[None, None], v[None, None])[0][0, 0]
for block in (1, 3, 8, 32):
np.testing.assert_allclose(online_attention(q, k, v, block), dense, rtol=1e-12, atol=1e-12)
def test_data_and_scaling(self):
self.assertEqual(exact_deduplicate(["a b", " a\n b ", "", "c"]), ["a b", "c"])
x = np.array([1., 2., 4., 8., 16.])
a, alpha = fit_power_law(x, 1.2 + 3 * x**(-.4), 1.2)
self.assertAlmostEqual(a, 3.)
self.assertAlmostEqual(alpha, .4)
with self.assertRaises(ValueError):
fit_power_law(x, np.ones(5), 1.)
def test_alignment_sign_and_mask(self):
advantages = group_advantages([[1, 1, 1], [0, 1, 0]])
np.testing.assert_array_equal(advantages[0], 0)
self.assertGreater(advantages[1, 1], 0)
self.assertAlmostEqual(dpo_loss(-5, -5, -5, -5), math.log(2))
self.assertLess(dpo_loss(-4, -6, -5, -5), math.log(2))
value = clipped_policy_objective([math.log(2), 9], [0, 0], [1, -100], [1, 0])
self.assertAlmostEqual(value, 1.2)
if __name__ == "__main__":
unittest.main()