ZeroPoint v0 research prototype: new from-scratch NumPy framework, memory-slot architecture, audited Transformer comparison
5c6c847 verified | """框架自检：用数值差分验证每个算子 / 层的反向传播，并跑两个小模型确认能学会。 | |
| 运行: python3 tests/test_zerocg.py | |
| """ | |
| import os | |
| import sys | |
| import numpy as np | |
| sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")) | |
| import zerocg as zc # noqa: E402 | |
| from zerocg import functional as F # noqa: E402 | |
| from zerocg import nn # noqa: E402 | |
| zc.set_default_dtype(np.float64) # 数值梯度检查需要 float64 的精度 | |
| rng = np.random.default_rng(0) | |
| def grad_check(fn, *shapes, tol=1e-5, positive=False): | |
| """fn(*tensors) -> Tensor。比较自动求导和中心差分的梯度（float64）。""" | |
| arrays = [rng.standard_normal(s) for s in shapes] | |
| if positive: | |
| arrays = [np.abs(a) + 0.5 for a in arrays] | |
| ts = [zc.Tensor(a.copy(), requires_grad=True) for a in arrays] | |
| w = np.asarray(rng.standard_normal(fn(*ts).shape)) # 随机投影成标量，检查完整雅可比 | |
| (fn(*ts) * zc.Tensor(w)).sum().backward() | |
| eps = 1e-6 | |
| for t, a in zip(ts, arrays): | |
| num = np.zeros_like(a) | |
| for i in np.ndindex(a.shape): | |
| for sign in (1, -1): | |
| b = a.copy() | |
| b[i] += sign * eps | |
| inputs = [zc.Tensor(b) if x is t else zc.Tensor(x.data) for x in ts] | |
| with zc.no_grad(): | |
| num[i] += sign * (fn(*inputs).data * w).sum() | |
| num[i] /= 2 * eps | |
| err = np.abs(num - t.grad).max() / max(1.0, np.abs(num).max()) | |
| assert err < tol, f"梯度误差 {err:.2e}" | |
| def module_grad_check(module, x): | |
| for p in module.parameters(): | |
| p.data = p.data.astype(np.float64) | |
| params = module.parameters() | |
| # 直接对参数逐个做差分 | |
| xt = zc.Tensor(x) | |
| w = np.asarray(rng.standard_normal(module(xt).shape)) | |
| for p in params: | |
| p.grad = None | |
| (module(xt) * zc.Tensor(w)).sum().backward() | |
| eps = 1e-6 | |
| for p in params: | |
| num = np.zeros_like(p.data) | |
| for i in np.ndindex(p.shape): | |
| old = p.data[i] | |
| p.data[i] = old + eps | |
| with zc.no_grad(): | |
| up = (module(xt).data * w).sum() | |
| p.data[i] = old - eps | |
| with zc.no_grad(): | |
| down = (module(xt).data * w).sum() | |
| p.data[i] = old | |
| num[i] = (up - down) / (2 * eps) | |
| err = np.abs(num - p.grad).max() / max(1.0, np.abs(num).max()) | |
| assert err < 1e-5, f"参数梯度误差 {err:.2e}" | |
| # ---------------- 算子 ---------------- | |
| def test_ops(): | |
| grad_check(lambda a, b: a + b, (3, 4), (4,)) | |
| grad_check(lambda a, b: a - b, (3, 4), (3, 1)) | |
| grad_check(lambda a, b: a * b, (2, 3), (2, 3)) | |
| grad_check(lambda a, b: a / b, (2, 3), (3,), positive=True) | |
| grad_check(lambda a: a ** 3, (4,)) | |
| grad_check(lambda a: a.sqrt(), (4,), positive=True) | |
| grad_check(lambda a: -a + 2.0 - a * 3.0, (4,)) | |
| grad_check(lambda a: 1.0 / a, (4,), positive=True) | |
| grad_check(lambda a: a.exp(), (3, 2)) | |
| grad_check(lambda a: a.log(), (3, 2), positive=True) | |
| grad_check(lambda a: a.tanh(), (5,)) | |
| grad_check(lambda a: a.relu(), (5,)) | |
| grad_check(lambda a: a.masked_fill(np.eye(3, dtype=bool), -5.0), (3, 3)) | |
| def test_reductions_and_shapes(): | |
| grad_check(lambda a: a.sum(), (3, 4)) | |
| grad_check(lambda a: a.sum(axis=1), (3, 4)) | |
| grad_check(lambda a: a.mean(axis=0, keepdims=True), (3, 4)) | |
| grad_check(lambda a: a.max(axis=-1), (3, 4)) | |
| grad_check(lambda a: a.reshape(4, 3), (3, 4)) | |
| grad_check(lambda a: a.transpose(2, 0, 1), (2, 3, 4)) | |
| grad_check(lambda a: a[1:, ::2], (3, 4)) | |
| grad_check(lambda a: a[np.array([0, 2, 0])], (3, 4)) # 重复索引 | |
| grad_check(lambda a, b: zc.concat([a, b, a], axis=1), (2, 3), (2, 1)) | |
| def test_matmul(): | |
| grad_check(lambda a, b: a @ b, (3, 4), (4, 5)) | |
| grad_check(lambda a, b: a @ b, (2, 3, 4), (4, 5)) # 批量 @ 共享权重 | |
| grad_check(lambda a, b: a @ b, (2, 2, 3, 4), (2, 2, 4, 3)) | |
| def test_functional(): | |
| grad_check(lambda a: F.softmax(a, axis=-1), (3, 5)) | |
| grad_check(lambda a: F.log_softmax(a, axis=-1), (3, 5)) | |
| grad_check(lambda a: F.gelu(a), (6,)) | |
| grad_check(lambda a: F.sigmoid(a), (6,)) | |
| grad_check(lambda a: F.cross_entropy(a, np.array([1, 0, 4])), (3, 5)) | |
| # 和直接的数学定义对比数值 | |
| x = rng.standard_normal((4, 6)) | |
| t = np.array([0, 5, 2, 2]) | |
| ref = -np.log(np.exp(x) / np.exp(x).sum(1, keepdims=True))[np.arange(4), t].mean() | |
| assert abs(F.cross_entropy(zc.Tensor(x), t).item() - ref) < 1e-10 | |
| def test_layers(): | |
| module_grad_check(nn.Linear(4, 3), rng.standard_normal((2, 4))) | |
| module_grad_check(nn.LayerNorm(5), rng.standard_normal((3, 5))) | |
| module_grad_check(nn.Sequential(nn.Linear(3, 4), nn.GELU(), nn.Linear(4, 2)), rng.standard_normal((2, 3))) | |
| module_grad_check(nn.CausalSelfAttention(8, 2), rng.standard_normal((2, 4, 8))) | |
| emb = nn.Embedding(5, 3) | |
| out = emb(np.array([[1, 1, 4]])) | |
| out.sum().backward() | |
| assert emb.weight.grad[1].tolist() == [2, 2, 2] and emb.weight.grad[0].tolist() == [0, 0, 0] | |
| def test_attention_is_causal(): | |
| attn = nn.CausalSelfAttention(8, 2).eval() | |
| x = rng.standard_normal((1, 5, 8)).astype(np.float32) | |
| y1 = attn(zc.Tensor(x)).data | |
| x[0, 4] += 10 # 改最后一个位置，前面的输出不应变化 | |
| y2 = attn(zc.Tensor(x)).data | |
| assert np.allclose(y1[0, :4], y2[0, :4]) and not np.allclose(y1[0, 4], y2[0, 4]) | |
| def test_count_flops(): | |
| with zc.count_flops() as c: | |
| zc.randn(2, 3, 4) @ zc.randn(4, 5) | |
| assert c.flops == 2 * (2 * 3 * 5) * 4 | |
| def test_default_dtype(): | |
| zc.set_default_dtype(np.float32) | |
| try: | |
| x = zc.Tensor(np.ones(3)) # float64 数组进来也变成 float32 | |
| assert x.dtype == np.float32 and (x * 0.5).dtype == np.float32 | |
| assert nn.Linear(2, 2).weight.dtype == np.float32 | |
| assert zc.Tensor(np.arange(3)).dtype.kind == "i" # 整数下标不动 | |
| assert np.isfinite(F.sigmoid(zc.Tensor([-1000.0, 0.0, 1000.0])).data).all() | |
| finally: | |
| zc.set_default_dtype(np.float64) | |
| def test_state_dict_roundtrip(tmp="/tmp"): | |
| path = os.path.join(os.environ.get("TMPDIR", tmp), "zerocg_test.npz") | |
| a = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4)) | |
| a.save(path) | |
| b = nn.Sequential(nn.Linear(3, 4), nn.LayerNorm(4)).load(path) | |
| x = zc.randn(2, 3) | |
| assert np.allclose(a(x).data, b(x).data) | |
| assert [n for n, _ in a.named_parameters()] == ["layers.0.weight", "layers.0.bias", "layers.1.weight", "layers.1.bias"] | |
| def test_optimizer_resume(): | |
| """训练 10 步 == 训练 5 步、存档、新建模型和优化器、读档、再训练 5 步。""" | |
| X, y = rng.standard_normal((8, 3)), np.array([0, 1, 1, 0, 1, 0, 0, 1]) | |
| def make(): | |
| np.random.seed(7) | |
| m = nn.Sequential(nn.Linear(3, 4), nn.Tanh(), nn.Linear(4, 2)) | |
| return m, zc.optim.AdamW(m.parameters(), lr=0.05, weight_decay=0.01) | |
| def train(m, opt, n): | |
| for _ in range(n): | |
| loss = F.cross_entropy(m(zc.Tensor(X)), y) | |
| opt.zero_grad() | |
| loss.backward() | |
| opt.step() | |
| a, opt_a = make() | |
| train(a, opt_a, 10) | |
| b, opt_b = make() | |
| train(b, opt_b, 5) | |
| model_state, opt_state = b.state_dict(), opt_b.state_dict() | |
| c, opt_c = make() | |
| c.load_state_dict(model_state) | |
| opt_c.load_state_dict(opt_state) | |
| train(c, opt_c, 5) | |
| assert all(np.array_equal(p, q) for p, q in zip(a.state_dict().values(), c.state_dict().values())) | |
| # ---------------- 端到端学习 ---------------- | |
| def test_learns_xor(): | |
| np.random.seed(0) | |
| X = np.array([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=np.float32) | |
| y = np.array([0, 1, 1, 0]) | |
| model = nn.Sequential(nn.Linear(2, 16), nn.Tanh(), nn.Linear(16, 2)) | |
| opt = zc.optim.AdamW(model.parameters(), lr=0.05) | |
| for _ in range(300): | |
| loss = F.cross_entropy(model(zc.Tensor(X)), y) | |
| opt.zero_grad() | |
| loss.backward() | |
| opt.step() | |
| assert (model(zc.Tensor(X)).data.argmax(1) == y).all(), f"XOR 没学会, loss={loss.item()}" | |
| def test_attention_memorizes_sequence(): | |
| """一层注意力 + 嵌入，记住一个短序列的下一个字符。""" | |
| np.random.seed(0) | |
| seq = np.array([0, 1, 2, 3, 2, 1, 0, 3]) | |
| x, y = seq[None, :-1], seq[1:] | |
| class Tiny(nn.Module): | |
| def __init__(self): | |
| self.tok = nn.Embedding(4, 16) | |
| self.pos = nn.Embedding(8, 16) | |
| self.attn = nn.CausalSelfAttention(16, 2) | |
| self.ln = nn.LayerNorm(16) | |
| self.head = nn.Linear(16, 4) | |
| def forward(self, idx): | |
| T = idx.shape[1] | |
| h = self.tok(idx) + self.pos(np.arange(T)) | |
| h = h + self.attn(self.ln(h)) | |
| return self.head(h) | |
| model = Tiny() | |
| opt = zc.optim.AdamW(model.parameters(), lr=0.02) | |
| for _ in range(200): | |
| loss = F.cross_entropy(model(x).reshape(-1, 4), y) | |
| opt.zero_grad() | |
| loss.backward() | |
| zc.optim.clip_grad_norm(model.parameters(), 1.0) | |
| opt.step() | |
| assert loss.item() < 0.05, f"loss={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: # noqa: BLE001 | |
| 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) | |