Gradient Checking

The project has no autograd. Every backward function is written by hand, so the tests need another way to tell whether a gradient is right. They estimate each gradient from the loss alone, using only forward evaluations, and compare the two.

Central differences

What is it?

Why does the model need it?

This project has no autograd, so every backward function is written by hand and could be wrong. Central differences use only forward evaluations, so they give an estimate that does not depend on the backward code being checked.

Mathematics

The default step is h = 1e-5 (gradcheck.py:3). check_params has its own h argument with the same default, and the tests leave both at their defaults.

What does each symbol mean?

Tensor shapes

g starts as np.zeros_like(x), so it has the shape of x, whatever that is. The only caller, test_numerical_grad_on_quadratic, passes a vector of 3 numbers. No test runs numerical_grad on a model parameter. That is the job of check_params, below.

3def numerical_grad(f, x, h = 1e-5):4    g = np.zeros_like(x)5    for idx in np.ndindex(*x.shape):6        old = x[idx]7        x[idx] = old + h8        fp = f(x)9        x[idx] = old - h10        fm = f(x)11        x[idx] = old12        g[idx] = (fp - fm) / (2 * h)1314    return g

Where is it used?

Source

gradcheck.py, lines 3-14, numerical_grad

PDF reference

Section 4.4, Eq. 37, p. 8

Used by

Not called by the pipeline.

Tested by: tests.test_gradcheck.test_numerical_grad_on_quadratic

• Tests it: tests/test_gradcheck.py, line 18 (test_numerical_grad_on_quadratic), for Central difference for one entry

Status

Implemented

Relative error

What is it?

Why does the model need it?

In general, gradients can differ in size by many orders of magnitude, so a single absolute limit would suit some entries and not others. A relative measure allows one limit for all of them. The eps guard keeps the division defined when both numbers are 0.

Mathematics

The epsilon guard. If a and b are both 0, the denominator would be 0 without eps. With it, the result is 0 / 1e-12, which is 0. This happens in practice: at any entry the loss does not depend on, both the gradient and the central difference are exactly 0.

The guard only helps below 1e-12. If both numbers are near 0, but one is above eps, the result is close to 1. The last row of the example shows this, and it is why the test skips one kind of parameter, as described under What the tests assert.

What does each symbol mean?

Tensor shapes

Both arguments are scalars, and so is the result. check_params calls it once for each sampled entry.

16def relative_error(a, b, eps = 1e-12):17    return abs(a - b) / max(abs(a), abs(b), eps)

Where is it used?

Its only caller is check_params, which only a test calls, so the pipeline never reaches it. No test calls it directly.

Source

gradcheck.py, lines 16-17, relative_error

PDF reference

Section 4.4, Eq. 38, p. 8

Used by

Called by: gradcheck.check_params

• Calls it: gradcheck.py, line 31 (check_params), for Relative error with the epsilon guard

Status

Implemented

Checking every parameter

What is it?

Why does the model need it?

It checks the whole forward and backward pass against the loss, for every parameter. Sampling keeps the test quick, because each sampled entry costs two evaluations of the loss. Reporting by name says which tensor disagrees.

Mathematics

For each name, the function draws samples_per_param random entries, 5 by default. Each entry takes one random integer per axis (gradcheck.py:24). It nudges that entry in place, runs the loss twice, restores the entry, and keeps the largest relative error so far.

What does each symbol mean?

Tensor shapes

grads[name] has the shape of the parameter with that name, because loss_and_grads builds grads with np.zeros_like. Each idx has one integer per axis of that shape.

19def check_params(loss_fn, params, grads, rng, samples_per_param = 5, h = 1e-5):20    worst: dict[str, float] = {}21    for name, W in params.items():22        w = 0.023        for _ in range(samples_per_param):24            idx = tuple(int(rng.integers(s)) for s in W.shape)25            old = W[idx]26            W[idx] = old + h27            fp = loss_fn(params)28            W[idx] = old - h29            fm = loss_fn(params)30            W[idx] = old31            w = max(w, relative_error(grads[name][idx], (fp - fm) / (2 * h)))32        worst[name] = w33    return worst

Where is it used?

Source

gradcheck.py, lines 19-33, check_params

PDF reference

Section 4.4, p. 9

Used by

Not called by the pipeline.

Tested by: tests.test_gradcheck.test_every_tensor_passes

• Tests it: tests/test_gradcheck.py, line 26 (test_every_tensor_passes), for Worst relative error for one parameter name

Status

Implemented

What the tests assert

tests/test_gradcheck.py builds a small model, smaller than the one in config.py, and passes check_params a function that returns only the loss:

6def 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, rng

22def 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)

The file has three limits, and these are all of them:

The last one also applies the default relative tolerance of np.allclose.

Why .bk is skipped. The key bias adds the same amount to every score in a row, and softmax ignores that, so its true gradient is 0. Both the backward code and the central difference then give only round-off, and as the Relative error section shows, the ratio of two such numbers says nothing useful. So a separate test checks the backward value directly:

39        assert np.abs(grads[f"b{l}.bk"]).max() < 1e-12

Limits of these checks.

• Each name gets only 5 random entries, so a mistake in a single entry of a large tensor can go unnoticed.
• The test runs one setup with one seed, on the tiny model, not at the sizes in config.py.
• The test test_numerical_grad_on_quadratic checks numerical_grad itself. check_params repeats the same formula in its own code and does not call numerical_grad.

Reading a failure

A failed assertion prints the name and the error, because of the second part of assert err < 1e-4, (name, err). The name is a key of the parameter dictionary, and it tells you which tensor to look at:

The pattern of names says more than one name does. Backward runs from the head to the embeddings, so a wrong gradient that is passed along affects everything computed after it:

• One name fails. Look at the line that fills that tensor. In the second column below, only b1.wq is above the limit.
• Many names fail together. Go through the names in backward order, and look between the last one that passes and the first one that fails. In the third column, the value block_backward returns for block 1 is wrong. Block 1 still passes, because its gradients are filled before it returns. Block 0 and both embedding tables fail.

The .bk rows hold round-off, as explained above, so their values are not a failure.

The assertion stops at the first failure. test_every_tensor_passes loops over worst in the order init_params created the names, and the first failure ends the test. With the third column, the report names only tok_emb, the first name in the dictionary, even though the mistake is elsewhere. To see the pattern, print the whole dictionary that check_params returns.

A check can only point to where a gradient disagrees with the loss. It cannot say whether the mistake is in the forward or the backward code, because the loss comes from the forward code, and the check treats it as correct.