Check gradients and custom backward kernels#

The problem#

You wrote a custom torch.autograd.Function with its own backward, or a fused backward kernel, and need to know its gradients are right. torch.autograd.gradcheck compares with finite differences, which needs float64 and is approximate. Comparing with PyTorch's own autograd in float32 runs into the same rounding questions as the forward.

Why the usual approach falls short#

Finite differences check the formula, not the rounding of a float32 or bfloat16 gradient kernel. Autograd's gradient is one correct answer among many, and a hand-picked tolerance around it fails or passes for the wrong reasons, just as in the forward.

The solution#

Each operation has a backward with the same guarantees as the forward. Check every gradient your function returns with covers():

Code
import math
import random
import torch
import exactref as er

# Seeded, so every run uses the same values
rng = random.Random(0)

# A tensor of the given shape, filled with random values
def randn(*shape):
    # One value per element: normal, with mean 0 and standard deviation 1
    values = [rng.gauss(0, 1) for _ in range(math.prod(shape))]
    # torch.tensor makes float32, PyTorch's default for Python floats
    return torch.tensor(values).reshape(shape)

# The function under test: A @ B with its own backward
class MyMatmul(torch.autograd.Function):
    @staticmethod
    def forward(ctx, A, B):
        # Keep the inputs for the backward
        ctx.save_for_backward(A, B)
        return A @ B

    @staticmethod
    def backward(ctx, grad):
        A, B = ctx.saved_tensors
        # The gradients for A and B, from the output's gradient grad
        return grad @ B.T, A.T @ grad

# Random inputs A and B, and an output gradient
A, B, grad = randn(16, 32), randn(32, 8), randn(16, 8)
# Copies that track gradients; A and B themselves stay as they are
A_, B_ = A.clone().requires_grad_(), B.clone().requires_grad_()
# Run the function and its backward, which fills A_.grad and B_.grad
MyMatmul.apply(A_, B_).backward(grad)

# ExactRef's gradients of matmul(A, B) for the same output gradient
g = er.matmul.backward(A, B, grad)
# Check each gradient against its own result
print("grad_A covered:", g["grad_A"].covers(A_.grad).all())
print("grad_B covered:", g["grad_B"].covers(B_.grad).all())
Output
grad_A covered: True
grad_B covered: True

A backward that quietly computes grad_A in float16 is caught, on the gradient where the bug is:

Code
# MyMatmul with a buggy backward
class HalfGradient(MyMatmul):
    @staticmethod
    def backward(ctx, grad):
        A, B = ctx.saved_tensors
        # grad_A computed from float16 copies, then converted back to float32
        return (grad.half() @ B.T.half()).float(), A.T @ grad

# Fresh copies, and the buggy backward
A_, B_ = A.clone().requires_grad_(), B.clone().requires_grad_()
HalfGradient.apply(A_, B_).backward(grad)

# The same ExactRef gradients g judge the new ones
print("grad_A covered:", g["grad_A"].covers(A_.grad).all())
print("grad_B covered:", g["grad_B"].covers(B_.grad).all())
Output
grad_A covered: False
grad_B covered: True

Attention's backward works the same way, with grad_query, grad_key and grad_value. For a fused backward kernel that is given the forward's output and logsumexp, use the matching attention algorithm's backward_covers().