Check low-precision kernels#

The problem#

You are checking a kernel that works in bfloat16, float16 or an 8-bit float format. Its outputs are far from a float64 reference, as they should be, and you cannot tell a correct low-precision result from a wrong one. Some kernels accumulate in float32 and some in the low-precision format itself; float8 formats saturate or give NaN past their largest value; GPUs round float32 inputs to TF32 before a product.

Why the usual approach falls short#

The right tolerance depends on where the kernel rounds and in which format. A tolerance loose enough for a bfloat16 accumulator accepts real bugs in a kernel that accumulates in float32, and one tight enough for float32 accumulation fails every correct bfloat16 accumulator.

The solution#

Tell ExactRef the format the kernel computes in. By default, for inputs narrower than float32, it assumes float32 accumulation:

Code
import torch
import exactref as er

# 300 ones in bfloat16
x = torch.ones(300, dtype=torch.bfloat16)

# A kernel that keeps its sum in bfloat16: each addition rounds to bfloat16
acc = torch.tensor(0.0, dtype=torch.bfloat16)
for value in x:
    acc = acc + value
print("bfloat16 sum:", acc.item())

# ExactRef's reference for the same sum
r = er.sum(x)
# The second argument is the compute format; float32 is the default here
print("covered, float32 accumulation:", r.covers(acc).all())
print("covered, bfloat16 accumulation:", r.covers(acc, "bfloat16").all())

# PyTorch's own sum passes the float32 check
print("x.sum() covered:", r.covers(x.sum()).all())
Output
bfloat16 sum: 256.0
covered, float32 accumulation: False
covered, bfloat16 accumulation: True
x.sum() covered: True

The same output is wrong for a float32 accumulator and correct for a bfloat16 one, so the check needs to know which kernel it is judging.

float8#

ExactRef imports float8 tensors exactly, and its formats follow PyTorch's conversions with the right overflow mode: "saturate" for float8_e4m3fn, whose PyTorch conversion turns large values and infinities into the largest value, and "nan" for the fnuz formats.

Code
# Two values past both formats' largest finite value, an infinity, and 0.3
big = torch.tensor([1000.0, float("inf"), -500.0, 0.3])

# PyTorch's conversion to float8_e4m3fn
fp8 = big.to(torch.float8_e4m3fn)
# ExactRef's conversion of the same values, saturating
saturated = er.asarray(big, "float8_e4m3fn", overflow="saturate")
# identical() is True when dtype, shape and every encoding match.
# er.asarray(fp8) imports PyTorch's float8 values exactly.
print("e4m3fn identical:", saturated.identical(er.asarray(fp8)))

# The same for float8_e4m3fnuz, with NaN on overflow
fnuz = big.to(torch.float8_e4m3fnuz)
with_nan = er.asarray(big, "float8_e4m3fnuz", overflow="nan")
print("e4m3fnuz identical:", with_nan.identical(er.asarray(fnuz)))
Output
e4m3fn identical: True
e4m3fnuz identical: True

A float8 matrix product usually produces a wider output. Ask for it with dtype:

Code
# A 1x4 and a 4x1 matrix in float8_e4m3fn
a8 = torch.tensor([[0.5, 1.25, -2.0, 3.0]]).to(torch.float8_e4m3fn)
b8 = torch.tensor([[1.0], [0.75], [0.5], [-0.25]]).to(torch.float8_e4m3fn)

# The product, with a float32 output
r8 = er.matmul(a8, b8, dtype="float32")
print("rounded:", r8.rounded.tolist())
# A float32 product of the same values; float() converts them exactly
print("covered:", r8.covers(a8.float() @ b8.float()).all())
Output
rounded: [[-0.3125]]
covered: True

TF32#

A GPU in TF32 mode rounds each float32 input to TF32 (10 fraction bits) before the product. ExactRef has no TF32 format; TF32 values are float32 values, so round the inputs the way the hardware does and pass those as float32. The reference and tolerance then apply to the rounded inputs.