Training with larger number of parameters while keeping fast iterations is an increasingly adopted strategy and trend for developing better performing Deep Neural Network (DNN) models. This necessitates increased memory footprint and computational requirements for training. Here we introduce a novel methodology for training deep neural networks using 8-bit floating point (FP8) numbers. Reduced bit precision allows for a larger effective memory and increased computational speed. We name this method Shifted and Squeezed FP8 (S2FP8). We show that, unlike previous 8-bit precision training methods, the proposed method works out-of-the-box for representative models: ResNet-50, Transformer and NCF. The method can maintain model accuracy without requiring fine-tuning loss scaling parameters or keeping certain layers in single precision. We introduce two learnable statistics of the DNN tensors -shifted and squeezed factors that are used to optimally adjust the range of the tensors in 8-bits, thus minimizing the loss in information due to quantization.
There has been a great deal of interest and research into reduced numerical precision of weights, activations and gradients of neural networks. If, for example, 16 bit floating point can be used instead of 32 bit floating point, then the memory bandwidth is halved along with significant gains in computational performance.
In this work the authors propose an 8-bit floating point format (denoted S2FP8) for tensors. In general, computing activations and gradients with such low precision at training time, has generally proved challenging without a number of tricks such as scaling the loss of each minibatch to within a reasonable range. Such ``tricks'' can be difficult to tune for each problem.
The key idea here is that for each tensor of 8-bit numbers, two 32 bit floating point statistics are recorded as well. These determine (in log-space) a scale and an offset for the 8-bit numbers (eq 1). This means that in this format tensors of significantly different scales can be well-represented (although larger scales necessarily implies low precision).
Matrix multiplications are done in FP32 precision and then converted in S2FP8 format. This requires an additional step to accumulate the summary statistics of each tensor in order to convert from FP32 to S2FP8 (the mean and the max of the tensor elements in log space).
The weights of the network are stored in FP32 and the gradients and activations are computed in S2FP8 and used to update the weights.
They test this approach in ResNet, a small transformer and a MLP for collaborative filtering. They find it reaches similar performance to FP32 where standard FP8 format has worse performance or results in Nan's.
Improvements in computational efficiency, both at training and inference, are active areas of research and this work contributes a novel approach using summary statistics. However, there are a several ways this work could be improved.
1. There is no comparisons with bfloat16, which is becoming a widely used approach to lower precision and is gaining significant hardware support [1].
2. Discussion and analysis regarding the need to gather summary statistics after each matrix multiplication (or other tensor operation). It is claimed that this brings minimal HW complexity, but this doesn't seem well justified. For a large tensor, this additional reduction to compute statistics may be expensive (in memory bandwidth and computation), particularly since this is done with FP32.
3. Even with the current implementation on a GPU, it should be possible with kernel fusions to gain some significant savings in memory bandwidth (and therefore computational speed), but there is no attempt anywhere to show any runtime benefit on current hardware.
Minor issues:
Some captions are very terse and the figures would benefit from a clearer explanation (e.g. figure 6).
[1] https://en.wikipedia.org/wiki/Bfloat16_floating-point_format