SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mitchell Wortsman*1
Tim Dettmers*1
Luke Zettlemoyer<sup>12</sup>
Ari Morcos†²
Ali Farhadi†
Ludwig Schmidt†134
We introduce new methods for 1) accelerating and 2) stabilizing training for large language-vision models. 1) For acceleration, we introduce SwitchBack, a linear layer for int8 quantized training which provides a speed-up of  $13 - 25\%$  while matching the performance of bfloat16 training within 0.1 percentage points for the 1B parameter CLIP ViT-Huge—the largest int8 training to date. Our main focus is int8 as GPU support for float8 is rare, though we also analyze float8 training through simulation. While SwitchBack proves effective for float8, we show that standard techniques are also successful if the network is trained and initialized so that large feature magnitudes are discouraged, which we accomplish via layer-scale initialized with zeros. 2) For stability, we analyze loss spikes and find they consistently occur 1-8 iterations after the squared gradients become underestimated by their AdamW second moment estimator. As a result, we recommend an AdamW-Adafactor hybrid which avoids loss spikes when training a CLIP ViT-Huge model and outperforms gradient clipping at the scales we test.
Large models trained on large datasets have recently led to multiple breakthroughs in machine learning such as GPT-3 [5] and PaLM [11]. While many components are necessary for successful large-scale training, two critical elements are training speed and stability. To enable further progress, we must ensure that 1) training is fast—the model should be able to see a lot of data even if it is large, and 2) training is stable—large models should not suffer from loss spikes which degrade performance. We study these two directions in the context of contrastive language-image pre-training (CLIP) [46]. We examine CLIP-style models because of their importance in computer vision: CLIP-style models reach state-of-the-art performance on a wide range of image classification tasks [46, 64, 44, 7] and underlie image generation methods such as DALLE-2 [49] and Stable Diffusion [51]. Our contributions towards fast training and stable training are as follows.
Towards fast training, we introduce SwitchBack, a linear layer for quantized training with int8 precision which matches the performance of the bfloat16 [63] baseline within 0.1 percentage points for CLIP ViT-Huge—a larger model than considered in the original CLIP paper. Linear layers account for the majority of the compute in standard transformer models, usually more than  $90\%$ , comprising the key, query, value, and out projection of the attention blocks as well as the multilayer perceptron. We perform all linear layers in low-precision (int8) while retaining other layers, such as layer norms, in higher precision. With this setup, we observe end-to-end speedups between 13 and  $25\%$  for CLIP ViT-Huge training:  $25\%$  compared to a standard linear layer implemented using the PyTorch [43] autograd python module and  $13\%$  compared to the standard PyTorch layer which include background CUDA/C++ optimizations which are difficult to replicate for custom layers.
<sup>1</sup>University of Washington. <sup>2</sup>Meta AI Research, FAIR Team. <sup>3</sup>Allen Institute for AI. <sup>4</sup>LAION. *Equal contribution. †Equal senior contribution.
Figure 1: We introduce SwitchBack, a linear layer for low-precision training. (Left) SwitchBack for int8 training matches the zero-shot ImageNet [15] accuracy of standard bfloat16 training within 0.1 percentage point for CLIP ViT-Huge [46, 20] and outperforms LLM.int8() [17]. (Right) For float8 (fp8) training [40], a baseline which uses tensor-wise quantization diverges for large models while SwitchBack matches the baseline. In these large-model, small-data experiments, our focus is on comparing methods and not final model accuracy, so we use short runs which makes it feasible to run many experiments.
SwitchBack starts from the observation that quantization noise grows with the inner dimension in a matrix multiplication. For CLIP training, the weight gradient computation involves a large inner dimension because CLIP training requires a large batch size [44]. Hence SwitchBack uses 16 bit precision matrix multiplication for the weight gradient computation while using int8 multiplications for the forward pass and layer input gradient computations. This approach leads to large accuracy improvements compared to LLM.int8() [17] (Figure 1). We will provide open-source Triton [57] kernels for Switchback to enable future work on efficient quantization schemes.
Besides int8 training, we also study large-scale 8-bit float (fp8) [40] training. We do not have access to hardware that supports fp8 data types, which is currently more rare than int8, so we use an accurate simulation of fp8 computation. SwitchBack also outperforms straightforward 8-bit float (fp8) baselines because tensor-wise quantized baselines diverge at  $>420\mathrm{M}$  scale (Figure 1). However, we demonstrate that these methods can achieve high accuracy if the network is trained while keeping feature magnitudes small, which we accomplish via layer-scale [58] initialized with zeros.
Towards stable training, we find that loss spikes occur in CLIP training when the AdamW [37] second moment estimator becomes out-of-date in the patch embedding [20] layer. In particular, the learning signal changes so that the moving averages of squared gradients underestimates their true magnitude. Indeed, in the absence of stability interventions, we show that loss spikes can be predicted by examining this ratio of the squared gradients to their moving average. We therefore recommend an AdamW-AdaFactor [54] hybrid, which we refer to as StableAdamW as it removes instabilities at the scales we consider and outperforms gradient clipping. Concretely, StableAdamW is AdamW with the update clipping technique introduced in AdaFactor. Update clipping tracks the average ratio of the gradient square to the second moment estimator and lowers the learning rate when the ratio is large.
The remainder of this paper is organized as follows: Section 2 focuses on low-precision training while Section 3 stabilizes training by reducing loss spikes.
2 8-bit training
This section develops and compares methods for eight-bit training of language-vision transformer models. First, Section 2.1 discusses preliminaries and related work. Next, Section 2.2 introduces and tests SwitchBack, a linear layer for int8 and float8 training. Finally, Section 2.3 develops alternatives to SwitchBack which can be used for float8.
2.1 Preliminaries and related work
Neural networks today typically use 16-bit operations for training [39] in either the float16 or bfloat16 format [63]. Floating point formats use a subset of bits to represent the exponent while the remainder specifies the fraction (often referred to as the mantissa). The float16 format uses 5 bits for the exponent while bfloat16 uses 8 and therefore covers a larger range—float16 has a range of  $(5.96 \cdot 10^{-8}, 65504)$  while bfloat16 has a range of  $(10^{-38}, 3 \cdot 10^{38})$ . Most floating point formats also have denormalized numbers which allow for a “soft underflow” which gets exponentially closer to 0.0f for each additional
Algorithm 1: PyTorch code for SwitchBack
class SwitchBackMatmul(autograd.Function):
    @staticmethod
def forward(ctx, X, W):
    ctx.save_for_backward = X, W
    X_int8, state_X = row-wisequantize(X)
    W_int8, state_W = tensor-wisequantize(W)
    return matmul_int8_and_dequanitze(
        X_int8, W_int8.t(), state_X, state_W
    )
@staticmethod
def backward(ctx, G):
    X, W = ctx.save_for_backward
    G_rowwise = rowwisequantize(G)
    W_int8, state_W = tensor-wisequantize_transpose(W)
    X_gradien = matmul_int8_and_dequanitze(
        G_int8, W_int8.t(), state_X, state_W
    )
    W_gradien = matmul.fp16(G.t(), X)
    return X_gradien, W_gradien
Algorithm 2: StableAdamW  $(\{\alpha_{t}\}, \beta_{1}, \beta_{2}, \epsilon)$
$v_{0},u_{0} = 0$    
for  $t = 1$  to  $T$  do   
 $g_{t} = \nabla f(\theta_{t})$    
// apply correction term to debias EMA.   
 $\hat{\beta}_{1} = \beta_{1}\cdot \frac{1 - \beta_{1}^{t - 1}}{1 - \beta_{1}^{t}}$ $\hat{\beta}_{2} = \beta_{2}\cdot \frac{1 - \beta_{2}^{t - 1}}{1 - \beta_{2}^{t}}$    
// update moving averages   
 $v_{t} = \hat{\beta}_{1}v_{t - 1} + (1 - \hat{\beta}_{1})g_{t}$ $u_{t} = \hat{\beta}_{2}u_{t - 1} + (1 - \hat{\beta}_{2})g_{t}^{2}$    
// for implementation convenience   
// operations below occur independently   
for each tensor   
 $\mathrm{RMS}_t = \sqrt{\mathbb{E}[g_t^2 / u_t]}$    
// update parameters   
 $\eta_t = \alpha_t / \max (1,\mathrm{RMS}_t)$ $\theta_t = \theta_{t - 1} - \eta_t\lambda \theta_{t - 1} - \eta_t v_t / (\sqrt{u_t} +\epsilon)$
bit in the mantissa. To prevent underflows float16 mixed precision training [39] has been developed which works as follows. The loss of a mini-batch is multiplied by a loss scalar to scale the loss and following backpropagation gradients into the representable range of fp16. This loss scaling is undone by rescaling the weight gradients before the optimizer updates fp32 main weights with the fp16 gradients. In PyTorch [43], the loss scalar is initialized to 65536. Every time an Inf/NaN is encountered, the update is skipped and the loss scalar is halved. If no Inf/NaN are encountered for 2k iterations, the scalar is doubled.
When the loss scalar becomes too low in float16 training the loss slowly diverges. This was observed by Cherti et al. [9] when training ViT-Huge CLIP models and remedied by switching to bfloat16. Another instance of float16 creating issues at scale was the training of OPT [73] and BLOOM models [52]. Indeed, many obstacles faced during the OPT project could have been alleviated by using bfloat16 [72]. Similarly, all float16 training runs for BLOOM ended in divergence, only after using bfloat16 was the training stable. However, fast bfloat16 support is only available on TPUs, or GPUs developed with or after the NVIDIA Ampere series (2021 or later).
While 16 bit training is the standard today, hardware support for 8 bit operations are becoming more common. Hopper GPUs support float8 (fp8) [40] and Ampere GPUs support int8. However, it is currently (2023) very difficult to attain Hopper GPUs. Moreover, while int8 and int4 are used for inference [17, 65, 16], and there is earlier work exploring 8 bit training for convnets [61, 78, 10], these formats are not commonly used for training transformer models at scale. The CLIP ViT-Huge models we train have 1B parameters including the image and text towers which is  $40\mathrm{x}$  larger than a standard ResNet-50 (23M) [28], and quantization is more challenging for large tensors [17]. Additional related work on quantization of large scale models (larger than BERT-large) and low-precision training and be found in Appendix A.
2.2 SwitchBack
2.2.1 Method
Overview. A linear layer consists of three matrix multiplications—one in the forward pass to compute outputs and two in the backwards pass to compute gradients for the input and weights. Our SwitchBack layer uses 8 bit precision for the first two matrix multiplies but switches back to higher precision for the weight gradient.
We compute the weight gradient in higher precision because this matrix multiplication involves dot products between vectors which have a length of batch size times sequence length. As CLIP training requires large batch sizes [46, 44], this inner dimension of batch size times sequence length is much larger than for the other matrix multiplies. As we show in Appendix D, variance due to quantization
increases with the inner dimension of the matrix multiply. This modification is what differentiates SwitchBack from LLM.int8(), allowing SwitchBack to match the bfloat16 baseline (Figure 1).
Notation. A standard linear layer is comprised of inputs  $X \in \mathbb{R}^{b \times n}$ , weights  $W \in \mathbb{R}^{m \times n}$ , and outputs  $Y \in \mathbb{R}^{b \times m}$ . In the forward pass, outputs are computed as  $Y = XW^{\top}$ . In the backwards pass the layer receives gradients of the loss with respect to  $Y$ , which we denote  $\dot{Y}$ . Then, gradients to inputs  $\dot{X}$  are computed via  $\dot{X} = \dot{Y}W$  while gradients to the weights  $\dot{W}$  are computed via  $\dot{W} = \dot{Y}^{\top}X$ . For linear layers in a transformer [60],  $b$  is batch size times sequence length, while  $n$  and  $m$  are small multiples of the embedding dimension.
Quantization. For the matrix multiplies in 8 bit precision we use quantization. There are a multiple quantization techniques to choose from and we will release code for all these alternatives. However, we find the best trade-off of simplicity and performance is from using i) row-wise quantization [31] for the inputs and gradients and ii) tensor-wise quantization for the weights. Additional information on quantization methods is provided by Dettmers et al. [17] but we summarize below. Using int8 as an example, which can represent integers from  $-127$  to  $127$ , we now define row-wise and tensor wise quantization. For a matrix  $X$  with rows  $x_{1},\ldots,x_{b}$ , row-wise quantization  $Q_{\mathrm{row}}$  and tensor-wise quantization  $Q_{\mathrm{tensor}}$  are given respectively by
$$
Q _ {\text {r o w}} \left(\left[ \begin{array}{c} x _ {1} \\ \vdots \\ x _ {n} \end{array} \right]\right) = \text {r o u n d} \left(\left[ \begin{array}{c} \frac {1 2 7}{\text {a b s m a x} \left(x _ {1}\right)} \cdot x _ {1} \\ \vdots \\ \frac {1 2 7}{\text {a b s m a x} \left(x _ {b}\right)} \cdot x _ {b} \end{array} \right]\right), Q _ {\text {t e n s o r}} (X) = \text {r o u n d} \left(\frac {1 2 7}{\text {a b s m a x} (X)} \cdot X\right) \tag {1}
$$
where  $\text{absmax}$  is the maximum of the absolute value.
Importantly, when applying  $Q_{\mathrm{row}}$  we also save the row-wise absolute maximums so that we can use them later for dequantization. We refer to this as the quantization state, or state, for short, so  $\mathrm{state}_{\mathrm{row}}(X) = [\mathrm{absmax}(x_1), \dots, \mathrm{absmax}(x_b)]^\top \in \mathbb{R}^{b \times 1}$ . Equivalently, for tensor-wise quantization we only need to store the tensor-wise absolute maximum so  $\mathrm{state}_{\mathrm{tensor}}(X) = \mathrm{absmax}(X) \in \mathbb{R}$ .
Since only the matrix multiply occurs in int8 precision we need to dequantize the outputs back to the original floating point precision. The forward pass with quantization and dequantization becomes
$$
\frac {\operatorname {s t a t e} _ {\text {t e n s o r}} (W)}{1 2 7 ^ {2}} \cdot \operatorname {s t a t e} _ {\text {r o w}} (X) * \underbrace {Q _ {\text {r o w}} (X) Q _ {\text {t e n s o r}} (W) ^ {\top}} _ {\text {i n t 8 m a t m u l}} \tag {2}
$$
where  $*$  denotes elementwise-multiplication, which in this case is broadcasted so that row  $i$  of the matrix  $Q_{\mathrm{row}}(X)Q_{\mathrm{tensor}}(W)^{\top}$  is multiplied by element  $i$  of  $\mathrm{state}_{\mathrm{row}}(X)$ .
As mentioned previously, we use row-wise quantization for the inputs and gradients and tensor-wise quantization for the weights. We find that using row-wise quantization for both matrices increases complexity at a negligible or no performance increase. As such, we use this simpler approach.
The last detail in our algorithm is hardware specific. NVIDIA GPUs, which we use in this work, do not implement the int8/float8 operation  $AB$  for matrices  $A$  and  $B$  and only  $AB^T$  is implemented. As such, it is necessary to transpose the weight matrix in the backward pass. To reduce the overhead of transposition and quantization we fuse both operations, meaning we load the required data once from slow DRAM into fast SRAM/shared memory and then perform both operations in this cached memory – this is critical for achieving speedups. We call this operation tensor-wisequantize_transpose, which is a fused tensor-wise quantize and transpose operation. Putting the pieces together, the result is Algorithm 1.
Variants. While Algorithm 1 is the most straightforward version of SwitchBack, we also present two alternative versions—SwitchBackM and SwitchBackQ—and will release triton [57] implementations for all three. Appendix C contains pseudocode. SwitchBackM (Algorithm 3) is a memory efficient version of SwitchBack which only saves 8 bit tensors for the backwards pass—we recommend its use when memory is limited. The small downside of SwitchBackM is that it requires an additional dequantize operation during the backwards pass which increases the runtime overhead. For CLIP ViT-Huge we observed only a negligible accuracy differences between SwitchBack and SwitchBackM. In addition, we present SwitchBackQ (Algorithm 4) which uses row-wise and column-wise quantization for the weights instead of tensor-wise. While we did not observe this to improve accuracy at the scales we consider, it's possible that it will perform better than SwitchBack at larger scale.
float8. While the explanation so far has used int8 as an example, the code for SwitchBack and float8 (fp8) is nearly identical. The only modification is that operations such as round  $(127x / \mathrm{absmax}(x))$  are replaced by float8cast  $(x / \mathrm{absmax}(x))$  where we simulate float8cast by rounding to the exact values of the float8 data type. This simulation improves on the simulation of [40] which only clips the input tensors into the representable range of the float8 data type, but not the exact values of the float8 data type. This simulation theoretically matches float8 training, but we are unable to perform real float8 training because we lack the hardware that supports float8 arithmetic. As such, we perform arithmetic in 16-bit with exact float8 values. For our int8 experiments we conduct the multiplications in int8 using A100 GPUs—we perform real int8 training without any simulation.
2.2.2 Experimental setup
To evaluate SwitchBack we train CLIP [46] visual transformer [20] models on LAION-2B [53]. Typically CLIP training, especially at ViT-Huge scale, is prohibitively expensive. Our goal is not high final accuracy but rather to contrast different methods for low-precision training. To enable running multiple experiments, we therefore only train for a small number of samples seen—380 million images—and use patch-dropout 0.5 [35]. We note that the experiment is still very expensive, corresponding to roughly 300 epochs of ImageNet training in terms of samples seen, or approximately 2.9e20 FLOPs per training run. After training on LAION-2B we evaluate the models zero-shot on ImageNet [15] using the 80 prompt templates from CLIP [46].
We use batch size 16384 (per-gpu batch size of 256) and train for a total of  $20\mathrm{k}$  iterations. The first 5k iterations are linear warmup while the remaining  $15\mathrm{k}$  are cosine decay. Training and evaluation are conducted with the OpenCLIP library [29] with learning rate 2e-3, weight decay 0.2, and batch-size 16384 using the optimizer described in Section 3.5.
2.2.3 Results
We test two main questions: (1) can we replicate 16-bit performance with SwitchBack and (2) can we get speedups. To test (1) we train CLIP models with SwitchBack across multiple scales with both int8 and float8 precision (Figure 1). To test (2) we profile operations in an individual linear layer and also measure end-to-end training speed. Loss curves for the training runs in Figure 1 are shown in Appendix Figure 9.
Accuracy. We find that SwitchBack can match standard 16-bit training performance and outperform baselines for both a) int8 precision and b) float8 precision.
For our int8 experiments (Figure 1, right), we contrast the performance of i) the standard baseline which uses mixed-precision bfloat16, ii) the matrix multiplication kernels from LLM.int8() [17], which is equivalent to SwitchBackQ (Algorithm 4) if the weight gradient multiplication was also performed in int8 using row- and column-wise quantization, and iii) SwitchBack. SwitchBack has a negligible accuracy drop of 0.1 percentage points compared to the bfloat16 baseline for CLIP ViT-Huge. In contrast, there is a drop of 5.9 percentage points when training with LLM.int8(). Section D details our hypothesis for why LLM.int8() fails to replicate 16-bit performance for CLIP training.
For our simulated float8 training experiments (Figure 1, right), we contrast the performance of i) the standard baseline which uses mixed-precision bfloat16, ii) a baseline which uses tensor-wise quantization for all matrices, that is the weights, inputs, and gradients, and iii) SwitchBack. SwitchBack has a negligible accuracy drop of 0.1 percentage points from the bfloat16 baseline for CLIP ViT-Huge. In contrast, training diverges for the baseline that uses tensor-wise quantization for all matrices.
Speed. By writing custom triton kernels [57] we achieve end-to-end speedups from  $13 - 25\%$  for CLIP ViT-Huge training.  $25\%$  compared to a standard linear layer implemented using the PyTorch [43] autograd python module and  $13\%$  compared to the standard PyTorch layer which include optimizations that are difficult to replicate for custom layers. Moreover, we find that the overhead due to quantization operations decreases with scale and is  $\sim 10\%$  for CLIP ViT-Huge. A detailed analysis of our speed-ups is in Appendix B.
2.3 Float8 training by reducing feature magnitude
We find that SwitchBack is necessary for high accuracy int8 training. However, this section develops other interventions which enable float8 training without SwitchBack. We show that high accuracy can be achieved via float8 training with tensor-wise quantization for the inputs, weights, and gradients, so
Figure 2: (Left) Training CLIP ViT-Large models with simulated fp8 precision using tensor-wise quantization for the inputs, weights, and gradients. All methods we try diverge except for using zero-init layerscale [58], which multiplies the output of each self-attention or mlp block with a learnable vector initialized to zero. (Right) Examining feature magnitudes (i.e., the average absolute value of the output for transformer block  $k$ ) for CLIP ViT-Huge at the beginning (init) and end of training. This suggests why zero-init layer scale enables float8 training—zero-init layer scale prevents high feature magnitudes which may cause issues for low precision training [17]. Without the intervention, the average feature magnitude becomes large for later blocks.
long as the network is initialized and trained in a way which discourages large feature magnitudes. We accomplish via layer-scale [58] initialized to zero.
We use the bitsandbytes library [18] to simulate float8 training using the fp8 types from Micikevicius et al. [40]. We use tensor-wise quantization for the inputs, weights, and gradients, so that all operations occur in simulated float8. In our simulation, we represent each value only with the exact values representable by float8, but we perform computations in float16 precision. We believe that tensor-wise quantization approximates the removal of quantize operations entirely. This is because, as we show in Appendix C.2 (Figure 11), the maximum of these tensors tends to evolve smoothly. Consequently, using a moving average for a maximum which is divided directly in the matmul is similar to tensor-wise quantization.
Layer-scale, introduced by Touvron et al. [58], scales each self-attention and MLP block output hidden state by a learnable vector of shape embed_dim. A pre-norm transformer block with layer-scale tensors  $\gamma_{1}$  and  $\gamma_{2}$  is defined as
$$
x _ {k} ^ {\prime} = x _ {k} + \gamma_ {1} * \text {s e l f ＿ a t t e n t i o n} \left(\operatorname {n o r m} _ {1} \left(x _ {k}\right)\right), \quad x _ {k + 1} = x _ {k} ^ {\prime} + \gamma_ {2} * \operatorname {m l p} \left(\operatorname {n o r m} _ {2} \left(x _ {k} ^ {\prime}\right)\right), \tag {3}
$$
where  $*$  is broadcasted elementwise multiplication.
Typically, layers are initialized so that they approximately preserve the variance of their inputs, and inputs have approximately unit variance [26, 27]. However, when combined with residual connections this can lead to higher norms in deeper networks.