bfloat16

A computer stores a real number as a sign, an exponent (how big) and a mantissa (the significant digits), see floating point. The usual 32-bit format spends 8 bits on the exponent and 23 on the mantissa: about 7 significant digits.

bfloat16 (“brain float”, from Google Brain) keeps the same 8-bit exponent but only 7 bits of mantissa: about 3 significant digits.

$\text{value} = (-1)^{s} \times 1.m \times 2^{e - 127}$

Because the exponent is the same, a bfloat16 can hold any number a 32-bit float can, from tiny to huge, just less precisely. That matters for neural networks: their weights and signals span a wide range, but they rarely need more than a few digits, so the network’s big matrix products can be done in bfloat16 with little effect on learning. Half the bits also means half the memory traffic, and processors with special units for it (NVIDIA’s tensor cores, Intel’s AMX, AMD’s AVX-512 BF16 instructions) do several times as many multiplications a second.

The usual way to use it is mixed precision: the matrix products run in bfloat16, while the weights, the sums of the loss and the updates stay in 32-bit, where rounding would add up.

In this project

In Two fighters and vision, the learner on the 96-vCPU EC2 machine ran the network’s passes in bfloat16 (the --bf16 option; the losses stay 32-bit). The update alone took 0.48 s instead of 0.88 s at 12 threads on an AMD EPYC, and a training iteration went from 1.74 to 1.43 s. It changes the numbers slightly, and it hasn’t yet been checked that a policy trained this way learns as well, so the post reports it as a separate, marked option.