SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 7: (Left) Measuring the  $\%$  of time occupied by quantize operations for a SwitchBack linear layer, which is usually less than  $20\%$  and decreases with dim. (Right) Benchmarking speedups for end-to-end CLIP training on a single node (with 4 A100 GPUs, per-GPU batch size 256, and gradient checkpointing) for various model sizes when replacing all linear operations in the transformer with SwitchBack (i.e., key, query, value, and out projections as well as the MLP). speedups reported over i) a custom linear layer implemented with torch.autograd (Algorithm 5), which matches our implementation of SwitchBack that uses torch.autograd, and ii) using the standard PyTorch nn.Linear which includes additional background C++/CUDA optimizations which we do not replicate. LLM.int8() [17] does not provide speed-ups over the torch.autograd or nn.Linear baseline at this scale—we compare the speed of SwitchBack and LLM.int8() in Figure 8.
Figure 8: Benchmarking speedups of SwitchBack compared to LLM.int8() [17] for end-to-end CLIP training on a single node (with 4 A100 GPUs, per-GPU batch size 128, and gradient checkpointing) for various model sizes when replacing all linear operations in the transformer (i.e., key, query, value, and out projections as well as the MLP).
size while we use triton's autotune feature which provides fine-grained optimized kernels for matrices of any size. Our kernels are easy to modify as they are written in Triton [57], and the code to run the benchmarks and produce Figure 6 is open sourced. In doing so, we invite the community to further improve the kernels and provide a benchmark for measuring this progress. Due to computational constraints we have not tested  $\dim > 4096$  and it's possible the kernels require additional tuning to perform well at that scale.
One downside of SwitchBack is that it requires quantize operations. However, it is already evident from Figure 6 that quantize operations occupy a small amount of time compared to matmuls. This is highlighted by Figure 7 (left) which displays the fraction of time occupied by quantize operations relative to matmuls for SwitchBack linear layers. Quantize operations occupy at most  $25\%$  of the time, this fraction decreases to around  $10\%$  or below for large dim.
We now conduct end-to-end speed tests for CLIP training on a single node with 4x A100 GPUs (Figure 7, right). This is in contrast with the speedup measurements so far in this which have measured individual layers independently. We benchmark speedups relative to using i) a baseline linear layer which we implement in PyTorch with torch.autograd.Linear (Algorithm 5) and ii) the PyTorch optimized linear layer nn.Linear. In both cases the speedups increase when going from CLIP ViT-Base to CLIP ViT-Huge. However, there is an additional  $\sim 12.5\%$  speedup when comparing SwitchBack to the baseline linear layer which uses torch.autograd. We believe this comparison is fair because SwitchBack is also implemented using torch.autograd, while the standard PyTorch nn.Linear layer has additional C++ and CUDA optimizations that we do not implement. We hope to collaborate with the PyTorch team to realize the additional  $\sim 12.5\%$  speedup. Finally, we note that the kernels from LLM.int8() [17] do not provide speedups over fp16 at the scale we consider.
C Additional code and figures
C.1 Additional Code
This Section provides additional pseudocode:
- Algorithm 3 is the memory efficient variant of SwitchBack.
- Algorithm 4 is the variant of SwitchBack which uses row- and column-wise quantization for the weights. For SwitchBackQ, the forward pass is given by
$$
\frac {1}{1 2 7 ^ {2}} \operatorname {s t a t e} _ {\text {r o w}} (X) \operatorname {s t a t e} _ {\text {r o w}} (W) ^ {\top} * \underbrace {Q _ {\text {r o w}} (X) Q _ {\text {r o w}} (W) ^ {\top}} _ {\text {i n t 8 m a t m u l}} \tag {4}
$$
where  $*$  is an elementwise product. Again, we append _transpose to a function in Algorithm 4 to mean that the operation is fused with a transpose.
- Algorithm 5 is a standard linear layer implemented with torch.autograd.
Algorithm 3 Memory efficient SwitchBackM
class SwitchBackMMmatmul(autograd.Function):
    @staticmethod
    def forward(ctx, X, W):
        # X [b, n] inputs
        # W [n, m] weights
        X_int8, state_X = row-wisequantize(X)
        del X
        W_int8, state_W = tensor-wisequantize(W)
        # save tensors in ctx
        ctx.save = X_int8, state_X, W_int8, state_W
        # Return output
        return matmul_int8_and_dequanitze(
            X_int8, W_int8.t(), state_X, state_W
        )
    @staticmethod
    def backward(ctx, G):
        # G [b, m] gradient to output
        # Recover tensors from ctx
        X_int8, state_X, W_int8, state_W = ctx.save
        X = dequantize_row-wise(X_int8, state_X)
        del X_int8
        W_gradient = matmul.fp16(G.t(), X)
        del X
        G_int8 = row-wisequantize(G)
        del G
        W_int8 = W_int8.t().contiguous()
        # Use 8bit matmul only for X_gradient
        X_gradient = matmul_int8_and_dequanitze(
            G_int8, W_int8.t(), state_X, state_W
        )
        return X_gradient, W_gradient
class SwitchBackMLinear(nn.Linear):
    def forward(self, X):
        return SwitchBackMMmatmul.apply(X, self.weight)
Algorithm 4 SwitchBack with row-wise and column-wise quantization for the weights SwitchBackQ
class SwitchBackQMatmul(autograd.Function):
    @staticmethod
def forward(ctx, X, W):
    # X [b, n] inputs
    # W [n, m] weights
    # save tensors in ctx
    ctx.save_for_backward = X, W
    X_int8, state_X = row-wisequantize(X)
    W_int8, state_W = row-wisequantize(W)
    # Return output
    return matmul_int8_and_dequanitze(
        X_int8, W_int8.t(), state_X, state_W
    )
@staticmethod
def backward(ctx, G):
    # G [b, m] gradient to output
    # Recover tensors from ctx
    X, W = ctx.save_for_backward
    G_rowwise = rowwisequantize(G)
    W_int8, state_W = column-wisequantize(transpose(W))
    # Use 8bit matmul only for X-gradient
    X_gradient = matmul_int8_and_dequanitze(
        G_int8, W_int8.t(), state_X, state_W
    )
    W_gradient = matmul_fp16(G.t(), X)
    return X_gradient, W_gradient
class SwitchBackQLinear(nn.Linear):
    def forward(self, X):
        return SwitchBackQMatmul.apply(X, self.weight)
Algorithm 5 A standard linear layer implemented with torch.autograd
class StandardLinearMatmul(autograd.Function):
    @staticmethod
def forward(ctx, X, W):
    # X [b, n] inputs
    # W [n, m] weights
    # save tensors in ctx
    ctx.save_for_backward = X, W
    # Return output
    return torch/matmul(X, W.t())
@staticmethod
def backward(ctx, G):
    # G [b, m] gradient to output
    # Recover tensors from ctx
    X, W = ctx.save_for_backward
    X_grad = torch/matmul(G, W)
    W_grad = torch/matmul(G.t(), X)
    return X_grad, W_grad
class StandardLinear(nn.Linear):
    def forward(self, X):
        return StandardLinearMatmul.apply(X, self.weight)
C.2 Additional Figures
This section presents additional figures.
- Figure 9 shows the loss curves throughout training for the models in Figure 1.
- Figure 10 presents a more fine-grained version of Figure 6.
- Figure 11 shows the mean and max for the gradient and activation (i.e., feature) throughout training.
- Figures 12 and 13 respectively examine the effect of batch size and learning rate on loss spikes. They find that loss spikes increase with high batch size and learning rate, but can be mitigated by reducing AdamW  $\beta_{2}$ . However, reducing  $\beta_{2}$  by too much can slow learning.
- Figure 14 shows that using a schedule for  $\beta_{2}$  of the form  $1 - \text{iteration}^{-\lambda}$  does not improve accuracy.
JLM.int8()baseline
Figure 9: Loss curves for the CLIP ViT-Base and CLIP ViT-Huge models evaluated in Figure 1. The left two plots display results for int8 training while the right two plots display results for float8 (fp8) training.
ViT-Base model
float16 baseline
—fp8 tensor-wise baseline
ViT-Huge model
SwitchBack fp8
Figure 10: A more fine-grained version of Figure 6.
Figure 11: The mean and max for (left) the gradient to the MLP weight and (right) the output of a transformer block throughout training. Different rows correspond to different choice of model size and layer scale.
Figure 12: Loss spikes increase with batch size for fixed learning rate and model size. Reducing AdamW  $\beta_{2}$  from its default in PyTorch of 0.999 mitigates loss spikes. Reducing  $\beta_{2}$  too much slows training.
Figure 13: Loss spikes increase with learning rate for fixed batch size and model size. Reducing AdamW  $\beta_{2}$  from its default in PyTorch of 0.999 mitigates loss spikes. Reducing  $\beta_{2}$  too much slows training.
Figure 14: We try a schedule for  $\beta_{2}$  which is used in AdaFactor [54] and PaLM [11] and refer to the experiment as  $\beta_{2}$  warmup. This means that  $\beta_{2}$  at iteration  $k$  is  $1 - \text{iteration}^{-\lambda}$ . In this Figure we try  $\lambda = 0.45, 0.5, 0.65$  and show on the  $x$ -axis  $\beta_{2}$  at the final iteration. This  $\beta_{2}$  warm-up does not improve accuracy in our setting.
Consider a matrix multiplication  $UV$  for  $U \in \mathbb{R}^{n \times k}$  and  $V \in \mathbb{R}^{k \times m}$ . This matmul consists of computing inner products between vectors of length  $k$ .
This section shows that error due to quantization increases with  $k$ . This suggests why SwitchBack may achieve high accuracy, as we avoid quantizing matmuls for which  $k$  is very large. For the weight gradient computation, which we leave in high precision,  $k$  is batch size times sequence length, which is often  $\approx 32000$  in our experiments. For the other operations which comprise a matmul,  $k$  is less than  $4 \cdot \text{embed\_dim}$  which is  $\leq 8000$  in our experiments. These dimensions are standard for CLIP training experiments [46, 9].
D.1 Analyzing variance due to quantization for inner products
This section measures the variance due to quantization for the inner product between  $u$  and  $v$ . Let  $u, v$  be vectors of length  $k$  vectors with each element drawn i.i.d. from a distribution with mean 0. Let  $u_i$  have variance  $\sigma_u^2$  and  $v_i$  have variance  $\sigma_v^2$ .
Next, let  $\hat{u}$  and  $\hat{v}$  be the quantized versions of  $u$  and  $v$ , respectively. We model quantization error as  $\hat{u}_i = u_i + \epsilon_i$  and  $\hat{v}_i = v_i + \xi_i$  where  $\epsilon_i, \xi_i$  are i.i.d. mean centered random variables with variance  $\sigma_q^2$ .
The aim of this section is to show that variance due to quantization grows with  $k$ . Our analysis is conservative because we do not assume the variance of  $\epsilon_{i}, \xi_{i}$  increase with  $k$ , though in practice we believe they would as the absmax of  $u$  and  $v$  increases with  $k$ .
We first examine the variance of  $\hat{u}_i\hat{v}_i$ . By using that all random variables are mean centered, this variance is given by,
$$
\begin{array}{l} \operatorname {V a r} \left(\hat {u} _ {i} \hat {v} _ {i}\right) = \mathbb {E} \left[ \left(\hat {u} _ {i} \hat {v} _ {i}\right) ^ {2} \right] (5) \\ = \mathbb {E} \left[ \left(\left(u _ {i} + \epsilon_ {i}\right) \cdot \left(v _ {i} + \xi_ {i}\right)\right) ^ {2} \right] (6) \\ = \mathbb {E} \left[ \left(u _ {i} v _ {i} + \epsilon_ {i} v _ {i} + \xi_ {i} u _ {i} + \epsilon_ {i} \xi_ {i}\right) ^ {2} \right] (7) \\ = \mathbb {E} \left[ u _ {i} ^ {2} v _ {i} ^ {2} + \epsilon_ {i} ^ {2} v _ {i} ^ {2} + \xi_ {i} ^ {2} u _ {i} ^ {2} + \epsilon_ {i} ^ {2} \xi_ {i} ^ {2} \right] (8) \\ = \operatorname {V a r} \left(u _ {i} v _ {i}\right) + \sigma_ {q} ^ {2} \left(\sigma_ {u} ^ {2} + \sigma_ {v} ^ {2} + \sigma_ {q} ^ {2}\right). (9) \\ \end{array}
$$
Next, we use linearity of variance for independent random variables to calculate  $\mathrm{Var}(\langle \hat{u},\hat{v}\rangle)$ . This is given by,
$$
\begin{array}{l} \operatorname {V a r} \left(\langle \hat {u}, \hat {v} \rangle\right) = \sum_ {i = 1} ^ {k} \operatorname {V a r} \left(\hat {u} _ {i} \hat {v} _ {i}\right) (10) \\ = \sum_ {i = 1} ^ {k} \operatorname {V a r} \left(u _ {i} v _ {i}\right) + \sum_ {i = 1} ^ {k} \sigma_ {q} ^ {2} \left(\sigma_ {u} ^ {2} + \sigma_ {v} ^ {2} + \sigma_ {q} ^ {2}\right) (11) \\ = \operatorname {V a r} \left(\langle u, v \rangle\right) + k \cdot \sigma_ {q} ^ {2} \left(\sigma_ {u} ^ {2} + \sigma_ {v} ^ {2} + \sigma_ {q} ^ {2}\right). (12) \\ \end{array}
$$
D.2 Takeaways
We have shown that for inner products with length  $k$  vectors, variance due to quantization increases with  $k$ . This means the variance of output units/features due to quantization increases with  $k$  which can be thought of making the outputs more noisy. Noise compounds throughout the network and will eventually drown out useful signal—for large  $k$  the network features or gradient will no longer lead to effective learning.
D.3 Why LLM.int8() fails: LLMs vs CLIP models
This Section details our hypothesis for why SwitchBack outperforms LLM.int8() for CLIP training, which is conditioned on the analysis in Section D.1 being a good model for training.
From our analysis we have shown that the variance in the output features increases with the size of the inner products of a quantized matrix multiplication compared to the full precision matrix multiplication. As such, we may have different failure modes for transformers pretrained on text, such as GPT-3 [5] or LLaMA [59], compared to CLIP models [46].
Pretrained large language models (LLMs) tend to have larger weight matrices relative to their batch sizes when compared to CLIP models. CLIP models perform best when the batch size is large [46, 44, 9]. As a consequence, LLMs and CLIP models have their most noisy operations for different matrix multiplications. LLMs are most noisy in the forward pass  $XW^{T}$  and during layer-to-layer back propagation  $\dot{Y}_k W_k = \dot{X}_{k-1}$  where inner product dimension are large, for example, they are 32768 and 8192 for the output projection of LLaMA 65B, 32768 and 8192. While the weight gradient inner product size is determined by the per-GPU batch size, which is 2048 for LLaMA [59] (4M tokens per full batch distributed across 2048 GPUs). As such, if the quantization produces the same variance in quantization errors, then the weight gradient in LLM int8 training is between 4x and 16x less noisy if the analysis in Section D.1 is a good model for training.
For CLIP training with ViT-Huge, we have a batch size of 65536 per GPU (256x images of size 224x224 inputs with patch size 14x14, leading to 16x16 patches for each images, resulting in 65536 patches per GPU). The dimensions for the weight matrices are  $1280 \times 5120$ . As such, analogous to above for the LLaMA LLM, the weight gradient in CLIP models is between 51.2x to 12.8x more noisy compared to the forward and layer-to-layer backpropagation operations if the analysis in Section D.1 is a good model for training. Notice that the CLIP weight gradient is twice as noisy compared to the most noisy LLaMA 65B operations if we assume that all quantization operations have the same error variance.
As such, low-precision LLM training and CLIP requires high-precision quantization routines for different parts of the training.
This also gives the reason why we believe LLM.int8() fails despite replicating inference performance – the weight gradient in CLIP training is a highly noisy operation which might not give enough signal to SGD to converge to a local minimum.
E RMS Spikes precede Loss Spikes
This section further elaborate on the predictive relationship between an RMS spike in the embedding layer and a loss spike as in Figure 4.
We define a heuristic to characterize loss and RMS spikes which we use for analysis. We determined these heuristics by checking if they qualitatively coincided with what appeared to be a loss spike. We display results in this Section so that the reader can also evaluate if these heuristics appear reasonable.
We define RMS spikes events as  $\{t: \mathrm{RMS}_t \geq 2.3\}$  while loss spike events are defined as the set of  $t$  where loss at time  $t$  exceeds the running mean by 3.2 times the running standard deviation. Finally, we ignore the first 1000 iterations when learning rate is low.
We also deduplicate the RMS and loss spikes iterations as follows: multiple spikes over a short time interval of 10 iterations are only counted as one spike and start at the earliest time. Moreover, we only count a loss spike if there are multiple deviations in an interval of 10, which indicates that loss has meaningfully spiked.
Our results are as follows:
- Figure 15 observes that out of 15 total loss spikes for ViT-Huge across different  $\beta_{2}$ , 14 out of 15 come 1-8 iterations after an RMS spike in the patch embedding layer (module.conv1.weight). With only 76 total RMS spike events, the probability that a loss spike follows 1-8 iterations after an RMS spike by chance is  $< 1\%$ .
Figure 16 repeats this analysis for ViT-Large, wherein 13 out of 15 loss spikes follow an RMS spike by 1-8 iterations. The probability that a loss spike follows an RMS spike by chance is  $1.0\%$ .
- Figure 17 zooms in on Figure 15 to show additional detail.
- Figures 18 and 19 examine the cases where loss spikes fail to be detected in Figures 15 and 16, finding them to mainly be issues with the heuristic identifying loss spikes, i.e., false positive loss spikes.
- Finally, Figure 20 repeats Figure 15 but examines the RMS of a random layer in the middle of the transformer—not the patch embedding layer. In this case, none of the loss spikes follow RMS spikes.
Figure 15: Observing a predictive relation between RMS spikes and loss spikes. For CLIP ViT-Huge and multiple  $\beta_{2}$  values, we use heuristics (Appendix E) to automatically identify loss spikes which we use for analysis. Out of 15 total loss spikes, 14 follow an RMS spike in the patch embedding layer (RMS > 2.3) by 1-8 iterations. We show loss spikes which are identified by our heuristic. We use green if they follow an RMS spike and otherwise use red. An RMS Spike indicates that the second moment estimator is out of date (see Section 3.4 and Shazeer and Stern [54]). The chance that a loss spike follows 1-8 iterations after an RMS spike by chance is  $< 1\%$ .
Figure 16: Observing a predictive relation between RMS spikes and loss spikes. For CLIP ViT-Large and multiple  $\beta_{2}$  values, we use heuristics (Appendix E) to automatically identify loss spikes which we use for analysis. Out of 15 total loss spikes, 13 follow an RMS spike in the patch embedding layer (RMS > 2.3) by 1-8 iterations. We show loss spikes which are identified by our heuristic. We use green if they follow an RMS spike and otherwise use red. An RMS Spike indicates that the second moment estimator is out of date (see Section 3.4 and Shazeer and Stern [54]). The chance that a loss spike follows 1-8 iterations after an RMS spike by chance is  $1.0\%$ .
Figure 17: Zooming in on a section of Figure 15.
Figure 18: Examining the "failure" case in Figure 15. We believe this is not really a failure as the non-predicted red loss spike does not really appear to be a spike at all. However, adjusting our heuristic led to the issue of true spikes not being identified.
Figure 19: Examining the "failure" cases in Figure 16. We believe these to primarily issues with our heuristic, but adjusting our heuristic led to other issues such as true spikes not being identified.
Figure 20: This figure repeats part of Figure 15 but examines the RMS of a random layer in the middle of the transformer—blocks.20.attn.in Projekt—not the patch embedding layer. RMS spikes no longer precede loss spikes.
F StableAdamW continued
F.1 Q&A
This Section asks and answers a series of questions the reader may have concerning Section 3.5.
- First, why not just use AdaFactor? The answer is that the community has moved away from AdaFactor [11] as they find that AdaFactor under-performs AdamW at scale [47]. We believe this is likely due to the factored moments, and not other features such as update-clipping. The goal of this work is to advocate using a hybrid. We tried porting other features from AdaFactor to AdamW such as the  $\beta_{2}$  schedule but did not find them to help (Figure 14). Moreover, while PaLM uses an AdaFactor-AdamW hybrid, we believe they don't use update clipping.
- Another question is, why not use an optimizer such as Lion [7] which does not divide updates by any value, and is therefore immune to the stuck-in-the-past scenario. We believe this may be a promising path forward. However, while we observe that Lion outperforms AdamW at small scale, Lion still slightly under-performs AdamW for CLIP ViT-Huge scale in our experiments.
- A final question is, why consider  $g_{t}^{2}$  in the numerator for computing  $\mathrm{RMS}_{t}$  and not  $v_{t}^{2}$ ? We also tried  $v_{t}^{2}$  and found the performance worse.
F.2 Implementation considerations
To prevent divide by 0 issues when computing  $\mathsf{RMS}_t$  we compute  $\mathsf{RMS}_t = \sqrt{\mathbb{E}[g_t^2 / \text{maximum}(u_t, \epsilon^2)]}$  where  $\epsilon$  is the AdamW hyperparameter for which we use 1e-6 and maximum is an elementwise maximum. This is instead of  $\mathsf{RMS}_t = \sqrt{\mathbb{E}[g_t^2 / u_t]}$ .
G Loss spikes and the loss scalar
This Section ties the low precision training results 2 with our investigation into stability. Overall we find that loss spikes can co-occur with large activations and gradients. Large activations and gradients may cause issues during low precision training due to a more limited representible range. Therefore, reducing loss spikes is an important step for successful low precision training.
Figure 21: Avoiding loss spikes is helpful for low precision training. As shown in this figure, loss spikes can coincide with with activation spikes and gradient spikes. Large activations/gradients can cause issues during low precision training due to a more limited representible range [17].
Supporting data is illustrated by Figure 21, in which an RMS spike precedes a loss spikes which coincides with spikes in the activations (i.e., features) and gradients. As we've previously seen (Figure 2), high feature magnitudes can pose challenges for low-precision training. Moreover, the spikes in the gradient are so large that Inf/NaN values occur, which results in the loss scalar [40] dropping many times. There are a few takeaways from this observation. First, reducing loss spikes is an important step to enabling low-precision training. Second, spikes in gradient magnitude can be transient and therefore we may be adjusting the loss scalar too often—if using the PyTorch default loss scalar, thousands of iterations would be required before the loss scalar recovered to its value before this event. Finally, the layers highlighted in this figure are the main layers where Inf/NaN are encountered. Concretely, while we only track every tenth block, we never observe any Inf/NaN for any transformer block greater than 0. However, with the PyTorch default loss scalar an Inf/NaN in a single layer will skip the update for the whole network.
This motivates the loss scalar that we use in our experiments when one is required (except for in Figure 21). We use a loss scalar which i) checks for Inf/NaN at the individual tensor level and skips the update at the tensor level—not globally, and ii) remains fixed at its initial value.
This scalar allows fp16 mixed precision training for CLIP models at ViT-Huge scale where previously the scalar became too low and training diverged [9]. We also believe an adaptive block-wise scalar as in Ramesh et al. [48] would remedy this issue. One interesting remark is that often when we observe an Inf/NaN, it is in the patch embedding layer. Therefore, in the case where Inf/NaN's happen frequently it recovers the stability solution of Chen et al. [8] which is to freeze the embedding layer. As a final remark, we note that loss spikes do not always cause the loss scalar to drop, and emphasize the loss scalar can drop for various other reasons than spikes. Figure 21 is just an existence example that loss spikes can result in activation spikes and Inf/NaN gradients.