> Mixed precision training. While the weights and the activations are quantized to low precision, the gradients and the optimizer states are stored in high precision to ensure training stability and accuracy. Following the previous work [LSL+21], we maintain a latent weight in a high-precision format for the learnable parameters to accumulate the parameter updates. The latent weights are binarized on the fly during the forward pass and never used for the inference process.
In this case there are two areas to optimize for: training efficiency and inference efficiency.
If I understand correctly, it stores the weights, gradients and second-moment estimates in FP32 like every other mixed-precision training (the Gopher paper has details on why storing them in FP32 is important), and quantized weights are used in forward pass. What I'm not sure is whether latent weights are used in backward pass, and my instinction is that the "Straight-through estimator" requires high-precision latent weights so they may still be needed. Training FLOPS can be roughly estimated as 6 FLOP per parameter per token, where 2 is forward pass, 2 is gradient computation and 2 is gradient accumulation (see https://medium.com/@dzmitrybahdanau/the-flops-calculus-of-la...). If only forward pass is quantized, this means only 1/3 of all FLOPS are optimized (and even then it has to be accumulated in FP32). So I'm skeptical of the gains in training efficiency here, and I can't find the numbers (how much energy or how much time is used for training, compared to regular FP16 mixed precision training? The papers boast inference energy savings which makes me even more skeptical of training energy savings)
For quantization efficiency, while QAT can certainly avoid the quantization step, PTQ methods are very cheap (usually <24 hours on RTX 4090 for Llama-2-70b) so I consider the cost of the quantization step negligible. There is not much difference in inference efficiency gains as PTQ and QAT can quantize to the same format. For final accuracy, unfortunately there is a lack of comparison between QAT and PTQ of fp16 models, and PTQ has the advantage of not requiring access to the original dataset, so I think it's very hard to make a fair comparison here but it's also likely the only area where QAT has actual gains compared to best PTQ methods.