tests/test_gradcheck.py

1import numpy as np2from gradcheck import numerical_grad, check_params3from transformer import ModelConfig, init_params, loss_and_grads456def tiny_setup():7    cfg = ModelConfig(vocab_size=7, context_len=8, emb_dim=4,8                      num_heads=2, ff_dim=8, num_layers=2)9    rng = np.random.default_rng(0)10    params = init_params(cfg, rng)              # float6411    ids = rng.integers(7, size=5)12    targets = rng.integers(7, size=5)13    return cfg, params, ids, targets, rng141516def test_numerical_grad_on_quadratic():17    x = np.array([1., -2., 3.])18    g = numerical_grad(lambda v: float((v ** 2).sum()), x)   # d/dx sum(x^2) = 2x19    assert np.allclose(g, 2 * x, atol=1e-8)202122def test_every_tensor_passes():23    cfg, params, ids, targets, rng = tiny_setup()24    loss_fn = lambda p: loss_and_grads(p, ids, targets, cfg)[0]25    _, grads = loss_and_grads(params, ids, targets, cfg)26    worst = check_params(loss_fn, params, grads, rng)27    for name, err in worst.items():28        if name.endswith(".bk"):29            continue                       # true gradient is 0, checked below30        assert err < 1e-4, (name, err)313233def test_key_bias_gradient_is_zero():34    # adding the same vector to every key shifts each score row by a constant,35    # softmax ignores constants, so dL/dbk must be 036    cfg, params, ids, targets, rng = tiny_setup()37    _, grads = loss_and_grads(params, ids, targets, cfg)38    for l in range(cfg.num_layers):39        assert np.abs(grads[f"b{l}.bk"]).max() < 1e-12