| """ZeroPoint 架构自检。运行: python3 tests/test_zeropoint.py""" |
| import itertools |
| import json |
| import os |
| import sys |
|
|
| import numpy as np |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, os.path.join(HERE, "..")) |
| sys.path.insert(0, HERE) |
| import zerocg as zc |
| from models import ZeroX |
| from models.zeropoint import CHUNK, MemoryMixer |
| from test_zerocg import module_grad_check |
|
|
| zc.set_default_dtype(np.float64) |
| rng = np.random.default_rng(0) |
|
|
| |
| OPTIONS = [dict(zip(("gate", "write", "init", "learn_lambda", "pos", "ema"), combo)) |
| for combo in itertools.product((True, False), ("softmax", "sigmoid"), ("hier", "uniform"), |
| (True, False), ("none", "sin"), (False, True))] |
|
|
|
|
| def test_parallel_equals_recurrent(): |
| """训练用的分块并行形式和推理用的递推形式，在所有开关组合下都必须一致。 |
| 长度 300：跨好几个块，也远超旧版并行形式 120 字的上限。""" |
| idx = rng.integers(0, 11, size=(2, 300)) |
| worst = 0.0 |
| for opts in OPTIONS: |
| np.random.seed(1) |
| model = ZeroX(vocab_size=11, n_embd=16, n_head=2, n_slot=5, n_layer=2, **opts).eval() |
| worst = max(worst, np.abs(model(idx).data - model.forward_recurrent(idx)).max()) |
| assert worst < 1e-8, worst |
|
|
|
|
| def test_matches_golden_reference(): |
| """和 tests/golden/zeropoint_ref.json 对答案：防止以后改代码时并行或递推形式悄悄变了。""" |
| with open(os.path.join(HERE, "golden", "zeropoint_ref.json"), encoding="utf-8") as f: |
| ref = json.load(f) |
| cfg = {k[7:]: int(v) for k, v in ref.items() if k.startswith("config:")} |
| model = ZeroX(**cfg).eval() |
| model.load_state_dict({k[6:]: np.array(v, dtype=np.float64) for k, v in ref.items() |
| if k.startswith("param:")}) |
| idx, logits = np.array(ref["idx"]), np.array(ref["logits"]) |
| assert np.abs(model(idx).data - logits).max() < 1e-10 |
| assert np.abs(model.forward_recurrent(idx) - logits).max() < 1e-8 |
|
|
|
|
| def test_causal_across_chunks(): |
| np.random.seed(2) |
| model = ZeroX(vocab_size=7, n_embd=16, n_head=2, n_slot=4).eval() |
| idx = rng.integers(0, 7, size=(1, 3 * CHUNK)) |
| y1 = model(idx).data |
| p = CHUNK + 5 |
| idx[0, p] = (idx[0, p] + 1) % 7 |
| y2 = model(idx).data |
| assert np.allclose(y1[0, :p], y2[0, :p]) and not np.allclose(y1[0, p:], y2[0, p:]) |
|
|
|
|
| def test_ema_state_is_bounded(): |
| """ema=True：写入恒为常数 v 时，槽会收敛到 v 本身，不会随长度无限变大（ema=False 时会涨到约 v/(1-λ)）。""" |
| np.random.seed(6) |
| for ema, bounded in ((True, True), (False, False)): |
| mixer = MemoryMixer(8, 1, 4, ema=ema) |
| S = mixer.init_state(1) * 0 |
| x = np.ones((1, 8)) |
| for _ in range(3000): |
| _, S = mixer.step(x, S) |
| v = np.abs(x @ mixer.value.weight.data + mixer.value.bias.data).max() |
| assert (np.abs(S).max() <= v * 1.0001) == bounded, (ema, np.abs(S).max(), v) |
|
|
|
|
| def test_mixer_gradients(): |
| np.random.seed(3) |
| module_grad_check(MemoryMixer(8, 2, 3), rng.standard_normal((2, 6, 8))) |
|
|
|
|
| def test_mixer_gradients_across_chunks(): |
| """梯度要能穿过块与块之间传递的状态，也要传到可学习的初始状态 s0。""" |
| np.random.seed(3) |
| mixer = MemoryMixer(4, 1, 2) |
| module_grad_check(mixer, rng.standard_normal((1, CHUNK + 3, 4))) |
| assert np.abs(mixer.s0.grad).max() > 0 |
|
|
|
|
| def test_initial_state_is_learned(): |
| model = ZeroX(vocab_size=5, n_embd=8, n_head=2, n_slot=3, n_layer=1) |
| assert np.allclose(model.init_state(3)[0][2], model.blocks[0].mix.s0.data) |
| assert "blocks.0.mix.s0" in dict(model.named_parameters()) |
|
|
|
|
| def test_retention_init_is_hierarchical(): |
| lam = MemoryMixer(16, 2, 6, max_half_life=64).retention().data |
| half_life = np.log(0.5) / np.log(lam[0]) |
| assert np.all(np.diff(lam[0]) > 0) and abs(half_life[0] - 1.5) < 1e-3 and abs(half_life[-1] - 64) < 1e-2 |
|
|
|
|
| def test_usage_recorded(): |
| model = ZeroX(vocab_size=5, n_embd=8, n_head=2, n_slot=3, n_layer=1) |
| model(rng.integers(0, 5, size=(2, 10))) |
| read, write = model.blocks[0].mix.usage |
| assert read.shape == write.shape == (2, 3) and np.allclose(read.sum(-1), 1) |
|
|
|
|
| def test_state_is_constant_size(): |
| model = ZeroX(vocab_size=5, n_embd=16, n_head=2, n_slot=4, n_layer=3) |
| states = model.init_state(1) |
| for t in range(300): |
| _, states = model.step(np.array([t % 5]), states, t) |
| assert sum(s.size for s in states) == model.state_floats() == 3 * 2 * 4 * 8 |
| assert all(np.isfinite(s).all() for s in states) |
|
|
|
|
| def test_float32_long_sequence_is_finite(): |
| """float32 下 1000 个字的分块并行不溢出，并且和递推形式一致。""" |
| zc.set_default_dtype(np.float32) |
| try: |
| np.random.seed(5) |
| model = ZeroX(vocab_size=9, n_embd=16, n_head=2, n_slot=6).eval() |
| idx = rng.integers(0, 9, size=(1, 1000)) |
| par = model(idx).data |
| assert par.dtype == np.float32 and np.isfinite(par).all() |
| assert np.abs(par - model.forward_recurrent(idx)).max() < 1e-3 |
| finally: |
| zc.set_default_dtype(np.float64) |
|
|
|
|
| def test_learns_to_copy_from_far_back(): |
| """记忆任务：开头一个字，中间 40 个干扰字（跨块），最后复述开头那个字。""" |
| np.random.seed(4) |
| V, gap = 6, 40 |
| model = ZeroX(vocab_size=V + 1, n_embd=32, n_head=2, n_slot=8, n_layer=1) |
| opt = zc.optim.AdamW(model.parameters(), lr=0.01) |
| for _ in range(300): |
| first = np.random.randint(0, V, 32) |
| idx = np.concatenate([first[:, None], np.full((32, gap), V)], axis=1) |
| loss = zc.functional.cross_entropy(model(idx)[np.arange(32), gap], first) |
| opt.zero_grad() |
| loss.backward() |
| opt.step() |
| assert loss.item() < 0.05, loss.item() |
|
|
|
|
| if __name__ == "__main__": |
| tests = [(k, v) for k, v in list(globals().items()) if k.startswith("test_")] |
| failed = 0 |
| for name, fn in tests: |
| try: |
| fn() |
| print(f" ok {name}") |
| except Exception as e: |
| failed += 1 |
| print(f" FAIL {name}: {type(e).__name__}: {e}") |
| print(f"\n{len(tests) - failed}/{len(tests)} 通过") |
| sys.exit(1 if failed else 0) |
|
|
