Stable and low-precision training for large-scale vision-language models
Authors: Mitchell Wortsman
Tim Dettmers
Luke Zettlemoyer
Ari S. Morcos
Ali Farhadi
Ludwig Schmidt
Venue: NeurIPS 2023 poster (2023)
Keywords: CLIP
int8
stability
TL;DR: We introduce new methods for 1) accelerating and 2) stabilizing training for large language-vision models.

Abstract: 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 under-estimated 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.

Stable and low-precision training for large-scale vision-language models
Mitchell Wortsman"!
Tim Dettmers"!
Luke Zettlemoyer
Ari Morcos
Ali Farhadi
Ludwig Schmidt 1131
Abstract
We introduce new methods for 1) accelerating and 2) stabilizing training for large language-vin visiin 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 6 training within 0.1 percentage points for the IB 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 under- estimated by their AdamW second moment estimator. As result, recommend an Adam W-A Wd-affactoo hybrid which avoids loss spikes when training a CLIP ViT--uge model and outperforms gradient clipping at the scales test.
Introduction
Large models trained on large datasets have recently led to multiple breakthroughs in machine learning such as GPT-3 SS and PaLM [T]. 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 fast- tte model should be able see a lot of data even it is large, and 2) training is stable stable-large models should not suffer from loss spikes which degrade performance. We study these two directions the context of contrastive language-image pre-training (CLIP) [46) We examine CL.IP-style models because of their importance in computer vision: CL.IP-style models reach state-of-the-ard performance on a wide range of image classification tasks [46]
64
44
and underlie image generation methods such as DALLE-2 [49 and Stable Diffuusion contributions towards fast training and stable training are as follows.
5
Our
Towards fast training, we introduce SwitchBack, linear layer for quantized training with int8 precision which matches the performance of the bfloat 663 bascline within 0.1 percentage points for CLIP ViT-Huge-aa larger model than considered in the original CLIP paper. I. 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 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 and 25% for CLIP ViT-Huge training: 25% compared to standard linear layer implemented using the PyTorch 43 autograd python module and 13% compared to the standard PyTorc layer which include background CUDA/C+ optimizations which are difficult to replicate for custom layers.
University of Washington. M Research, FAIR Team, 3 Allen Institute AL. 'LAION. "Equal contribution. Equal senior contribution.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Zero-shot ImageNet accuracy
Zero-shat ImageNet, accuracy
50
45
b0 baseline
blooat16 baseline
Mintt) baseline
fpä tensar-wise baseline SwitchBack fp8l
ViT-Large
ViT-Huge
VIT-Large
Model size
Model size
Figure 1: We introduce SwitchBack, lincar layer for low low-precision training. (Left) SwitchBack for int8 training matches the zero-shot ImageNet 15 accuracy of standard 6 training within 0.1 percentage point for CLIP ViT-Huge (46 20) and outperforms LLM.int8() [17]. (Right) For float8 (fp8) training [40), bascline which uses tensor-wise quantization diverges for large models while SwitchBack matches the baseline. In these large-model, amall-dabe experiments, our focus is comparing methods and not final model accuracy, so short runs which makes it feasible 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 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 large accuracy improvements compared to LLM.int8( [17] (Figure D. 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 accurate simulation of fp8 computation. SwitchBack also outperforms straightforward 8-bit float (fp8) baselines because tensor- wise quantized baselines diverge at >420M scale (Figure [ However, demonstrate that these methods 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, find that loss spikes occur in CLIP training when the Adam [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 their moving average. We therefore recommend an AdamW W-AdaFactor [54] hybrid, which we refer to as dtableAdam as it removes instabilities at the scales we consider and outperforms gradient clipping. Concretely, Stable/damW 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 large.
The remainder of this paper is organized follows: Section focuses low- low-precision training while Section stabilizes training by reducing loss spikes.
2 8-bit training
This section develops and compares methods for eight-bit training languge- vision transformer models. First, Section 2.1 discusses preliminaries and related work. Next, Section 2.2 2 introduces and tests SwitchBack, 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 [3 in either the float16 or bfloat16 format [63]. Floating point formats use subset of bits to represent the exponent while the remainder specifies the fraction (often referred to as the mantissa). The f format uses bits for the exponent while bfloat 6 uses 8 and therefore covers a larger Lange-floatl 6 has range of (5.96 10- 65504) while 6 has range of (10-38 388 3- 1038). 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
Algorithm 2: StableAdamW
Un, 2Ap = 0
class SwitchBacikhatmul (autograd.Punction)) sataticnethod
for = 1 to T do
def (orx,rd X, w): ctx. save_ttr___aaaaeeee _o__back - X,
= V (0()
// apply correction term debias EMA.
inins, state_X row-viee, sow-viee_quantizee
81 B1 1-8
R_ant8, state_N tensor tensor-wise_ quantizae
1-
B3 B2
1-8
return natmul, cattal_intt____ wtw.iiittnnnndeneeeaeenee und_ddequnntae K_int,, V. in, state_I,
T-B5
state_k
update moving averages
0t = 1 +1 (1--31), 31)g
@staticnethod
backvard(ctz, Q):
W₁ 1224411 + 32)92
X, V ctx.save_for caarr__fo__aacaaaae
// for implementation convenience
G_rowwise Sovrise_quantizefe
I/ operations below occur independently for each tensor RMS: / [gf /u
N_int8, state_N tensor- kise_quantize. (rrnpppeeee
X_gradient = natmul int8 and_ dequaniatt G_int8, V_ ()) state_X, state_k
// update parameters
0/ max (1, RMS:)
R_gradient nataul (p16(..0))
0. = 0c-1 - T/e n0,0.-1 - m (vrr +()
x_gradient, N.gradiont
bit in the mantissa. To prevent underflows 16 mixed precision training 339 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. Everytime 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 bfloat 6. Another instance of 1 creating issues at scale was training of OPT 73) and BLOOM models [22] Indeed, many obstacles faced during the OPT project could have been alleviated by using 6 [2]. Similarly, all float16 training runs for BLOOM ended divergence, only after using bfloat16 6 was the training stable. However, fast bfloat16 1 support is only available on TPUs, 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 [6], and there earlier work exploring bit training for convnets [6] 78 [0] these formats are not commonly used for training transformer models at scale. The CLIP ViT-Huge models we train have IB parameters including the image and towers which 40x larger than a standard [17]. Additional related work on quantization of large scale models (larger than BERT-large) and low-precision training and be found in Appendix A
ResNet-50 (23M) (28), and quantization is more challenging for large tensors
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 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.
compute the weight gradient in higher precision because this matrix multiplication involves dot products between vectors which have length of batch size times sequence length. As CLIP training requires large batch sizes 46 this inner dimension of batch size times sequence length much larger than for the other matrix multiplies. As show in Appendix D variance due quantization

---

increases with the inner dimension of the matrix multiply. This modification is what differentiates SwitchBack from LLM. allowing SwitchBack to match the b6 baseline (Figure D
Notation. A standard linear layer is comprised of inputs X E R weights W E RRmxx, and outputs Y E RRxx._ In the forward pass, outputs computed Y XW In the backwards pass the layer receives gradients of the loss with respect to Y, which denote Y. Then, gradients to inputs X are computed via X YW while gradients the weights W computed via W Y X. For linear layers in a transformer [60], 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 [3] 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 example, which can represent integers from - 127 to 127, now define row-wise and tensor wise quantization. For a matrix X with rows Ii, Tb. row-wise quantization Qrow and tensor-wise quantization Qiensoe are given respectively by
127 absmax (21)
21
127 absmax (X) X
Qrow
round
Quensor (X) round
(1)
127
absmax (16)
where absmax is the maximum of the absolute value.
Importantly, when applying Qrow we also save the row-wise absolute maximums so that we can use them later for dequantization. We refer this as the quantization state, or state, for short, so state (X) = (absmax(x1)) ..., absmax(2.)
E Rbx1_ Equivalently, for censor-wise quantization we only need store the tensor-wise absolute maximum state (X) absmax (X) R.
Since only the matrix multiply occurs in int8 precision we need to dequantize the outputs back the original floating point precision. The forward pass with quantization and dequantization becomes
ee (W)
state (X) * Qrow (X) Quensor (W)
1272
int8 matmul
where * denotes cetpenenttiie-mlliippiiiee
which in this case is broadcasted so that row of the matrix Qrow (X) Qremsor (W) T is multiplied by element of state (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 negligible or performance increase. As such, 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" is implemented. As such, it necessary to transpose the weight matrix in the backward pass. To reduce 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 operation in this cached memory this is critical for achieving speedups. We call this operation tensor-wiet setsse_____eeeeeeeeeeee .uuantize_trasspossee which is fused tensor-wise quantize and transpose operation. Putting the pieces together, the result is Algorithm D
overhead of
Variants. While Algorithm| is the most straightforward version of SwitchBack, we also present two alternative versions-SSwichhBccee SwitchBBakkMM and SwitchBackQ and will release triton (57) implementations for all three. Appendix [ pscudocode. SwitchBacKM (Algorithm 3) is 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 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 SwitcBBackQ (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 that operations such as round (177 /277 are replaced by float8cast(1:/ (2/ /abbmax(aax))) 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 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 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 [6]].
We use batch size 16384 (per-gpu batch size of 256) and train for total of 20k iterations. The first 5k iterations are linear warmup while the remaining 15k 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 D 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 [ Jrre 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 I right), we contrast the performance of i) the standard baseline which uses muxed-precision biloat 16, ii) the matrix multiplication kernels from LLM.int80 [17). which is equivalent to SwitchBackQ (Algorithm 4 if the weight gradient multiplication was also performed in int8 using row- and column- quantization, and iii) SwitchBack. SwitchBack has a negligible accuracy drop of 0.1 percentage points compared 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.
the bfloat16 baseline for CLIP
For our simulated float8 training experiments (Figure [ right), we contrast the performance of i) the standard baseline which uses mixed-precisiog bfloat1 ii) a baseline which uses tensor- wise quantization for all matrices, that the weights, inputs, and gradients, and iii) SwitchBack. SwitchBack has a negligible accuracy drop of 0.1 percentage points from the bfloat 6 baseline for CLIP Vig- 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 optimiza- tions that are difficult to replicate for custom layers. Moreover, we find that the overhead due to quantization operations decreases with scale and 100 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
5

---

ViT-Large model
10
Average feature magnituds
standadd
Stanciadd frai
2ero-irit layer scale at
scale final
Loss
2500 5000 7500
10000 2500 15000 17500 20000
Iterations
laye-om
10.00 suatt 25 111 tell lenacr grad cip
fpr tenser layorscale
Transformer black index
Figure 2: (Left) Training CLIP Vir-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 mip block with learnable vector initialized to zero. (Right) Examining feature magniuuees (i.c., the average absolute value of the output for transformer block k) for CLIP ViT-Huge at the beginning (init) and end of training. This suggest why zero-init layer scale enables float8 training zero-init layer scale prevents high feature magnitudes which may cause issues for low precision training 177. Without the intervention, the average feature magnitude becomes large for later blocks.
long the network is initialized and trained in way which discourages large feature magnitudes. We accomplish via layer-scale (58) initialized to zero.
We use the bitsandbytes library to simulate float8 training using the fp8 types from Micikevicius et al. [40] We use tensor-wis wise quantization for the inputs, weights, and gradients, that all operations occur in simulated float8. In our simulation, we represent each value only with the exact values representable by float8, but perform computations in float16 precision. We believe that tensor-wise quantization approximates the removal of quantize operations entirely. This is because, we show in Appendix (C.2 (Figure [T), the maximum of these tensors tends evolve smoothly. Consequently, using a moving average for a maximum which is divided directly in the matmul is similar tensor-wise quantization.
Layer-scall scale, introduced by Touvron et al. [58], scales each self-attention and MLP block output hidden state by learnable vector of shape embed dim. pre-norm transformer block with layer-scale tensors ? and 72 is defined
1 = Us + 71 * calf_aattetionnnoneee (rx)),, Fk41 = 2's + 72 mlp((orma
where * is broadcasted elementwise multiplication.
Typically, layers are initialized so that they approximately preserve the variance 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.
Consequently, researchers have proposed initialization and scaling schemes which remedy this issue [17 19.. Layer-scale with initialization 0 is example of one such scheme-at initialization the transformer an identity function. While 71, 7/ are typically initialized as vectors of 10 or 10 use 0 for simplicity.
Figure [ (right) demonstrates that the layer-scale intervention is successful at controlling the average magnitude output. Without the intervention, the average feature magnitude E[abs(xx)] becomes high for later blocks. Previous work [7 has shown that large feature magnitudes result in issues for low precision training.
Results for simulated fp88 training are shown in Figure (left) for ViT-Large. We find that all fp8 runs diverge except for when we use layer-scale initialized to zero. Concretely, Figure compares i) the baseline which uses bfloat16 training, ii) using fp8 with tensor-wise quantization and no further modifications, which slowly diverges, iii) adding gradient clipping to ii). which also diverges, iv) adding KQ layernorm [14] ii), which also diverges, and v) using zero-init layerscale, which trains without diverging. While there is difference still between fp8 and bfloat 6 training. this is primarily because of layerscale. Moreover, we believe that with hyperparameter tuning layerscale would match standard training in terms of accuracy.

---

Base
Large model
modell
5000
10000
15000
20000
5000
10300
15000
20000
5000
10000
15000
20000
Iteration
Iteration
Iteration
beta2
botbz
betaZ
Figure 3: Loss spikes increase with model size for fixed learning rate and batch size. Reducing AdamW B2 from its default in Poroh 0.999 mitigates loss spikes. Reducing 02 too much slows training.
3 Stability
We now switch focus from accelerating learning by reducing precision to addressing instabilities which can arise during training. Section reviews preliminaries and related work while Section 3.2 details the experimental setup. Next, Section 3.3 examines trends for training instability, finding loss spikes to increase with model scale but decrease with lower Adam W 32- Then, Section 3.4 finds that loss spikes arise in setting due to out-of-date Adam W second moment estimator leading Section 3.33 to adopt and tests fix developed in the context of AdaFactor [54]. Appendix Section connects this section stability with the previous section low-precision training.
3.1 Preliminaries and related work
Loss spikes emerge when scaling up models 8 may slow learning, or even destabilize training completely. Various solutions have been proposed. including freezing the embedding layer [8]. adding additional layer normalization [4 25 or
25
J4
68
70
54
73,.
These instabilities
reparametrizing the weights 688
In our work we investigate instabilities which arise during CLIP training. Unlike the instabilities observed in [14 68) which lead a slow divergence, we study fast loss spikes. Our results indicate that these spikes arise when the second moment estimator is out of date for early layers.
While analysis and methods build directly Shazeer and Stern [54] (AdaFactor), there important differences. contrast with Shazeer and Stern [54], who only observe instabilities without warmup, observe instabilities despite a long warmup period. Moreover, in contrast with Shazeer and Stern [54] we find that an out-of-date second moment estimator is primarily an issue for the (patch) embedding layer, and measure how well loss spikes predicted by this event. Finally, note that researchers have moved away from AdaFactor in its original formulation for large-scale training 47 11 69), finding AdaFactor to under-perform A (47). We believe this is due the
factored second moment or absence of first moment. This is why our focus is Adam W 37 which is the de facto standard optimizer for transformers.
After the initial version of this paper we became aware Cohen et 12 which offers a general and principled treatment of fast loss spikes, and which we recommend to readers. Moreover, we direct the readers attention to the concurrent work of
3.2 Experimental setup
As in Section 2, we train ViT CLIP models on LAION (53) using OpenCLIP 29 and evaluate them zero-shot on ImageNet. Since we are not interested in final performance and instead interested in studying instability-even for very large models-w- use short run which allows us to conduct multiple experiments. Concretly, we use patch-dropout 0.5 335 and 20k iterations. The first 5k iterations are linear warmup while the remainder are cosine decay (36]. We follow the CLIP paper [46] in that i) we do not use gradient clipping unless otherwise mentioned, though we do clip the logit_scale parameter, and ii) we add a layer-norm after the patch embedding and before the main transformer. Unless otherwise mentioned, experiments use batch size 16384 (per-gpu batch size of 256), learning rate 2e-3 and weight decay 0.2. We initially tried adding a layer-norm before the patch embedding as in (34], but removed this as we found it to hurt performance at CLIP ViT-Huge scale.
7

---

RMs spike preceeding loss
bets2
beto2
7.5
5.0
2500
2600
2700
2800
2900
3000
cisiel.corvv. miight
RMS
2500
2600
2700
2800
2900
3000
Iteration
Figure 4: The learning signal can change so that the AdamW second moment estimator is out-of-date and underestimates the squared gradients g. This can be detected the aggregate quantity RMS: = / [ot /uu]] far from 1. This figure observes predictive relationship between the event of an RMS spike and loss spike- we observe spike in RMS, 1-8 iterations before loss spike. For lower 32, RMS, does not deviate far from 1. This result looks RMS, for the patch embedding layer only. This predictive relationship is further examined Figures to 20 of Appendix
3.3 Loss spikes increase with model size, batch size, and learning rate
We begin our studying of loss spikes by observing how their presence varies when changing model size, batch size, and learning rate. The following sections build on these observations- particular the finding that lowering the Adam B3 hyperparameter removes spikes entirely.
We find that loss spikes increase when increasing model size, batch size, or learning rate. The first result is shown in wigle the second two analagous results are Appendix Figures Importantly, these figures show that loss spikes can be avoided by reducing the 32 hyperparameter for in AdamW. On the other hand, if (3 is reduced too much then learning slowed which results in worse performance [50).
3.4
On and out-of-date second moment estimator
Based on the observation in the previous section that lowering 82 reduces spikes, this section traces the cause loss spikes to an out-of-date second moment estimator the patch embedding layer.
Overview. Adaptive optimizers such AdaGrad [22], Adam (33. or AdaFactor (54) scale the
update differently for each individual parameter. This is often conceptualized per-parameter learning rate. For instance, in Adam/Adam exponential moving average of squared gradients (see the code for Adam in Algorithm 2 ignoring
per-parameter updates are scaled by the inverse root of the
for now the modifications in pink which discuss in Section 3.5).
This adaptivity can be a very useful tool for accelerating training, but can also cause issues when the learning signal changes. Concretely, exponential moving averages become out of date causing updates to be scaled by a value that is too large. This issue is discussed in Section of Shazeer and Stern (54), and summarize below.
As in Algorithm 2 let u, = {u. denote the exponential moving average (EMA) of squared gradients of = {gr.5 f for neural network parameters 0 E R". Ignoring the bias correction term]] at each iteration u, is updated as 3214-1 + (1 32)9î where 3g is referred to as the decay for the EMA. Then, the update is scaled by 1/ (vur + e), where is a small value added numerical stability, Often the ratio / (v + <) is thought of as signal-to-noise ratio of the gradient over time.
However, this method can break down when the learning signal changes and u, ceases to be a good estimator for the running average of of- Consider the case where the gradient magnitudes have been historically very small for some parameters so 1/ (vu, e) is large for those parameters. If, then, at iteration / those parameters suddenly receive larger gradient signal the update can be catastrophically big. We refer to the scenario the stuck-in biick-ii-thhe in-the- past scenario.
Overall, B2 is too small then convergence may be slowed [50). Bg is too large then u, can become out-of-date and no longer good estimator for 9² resulting in per-parameter scaling that is too large.
In practice the EMA is debiased with correction term. Algorithm 2 follows AdaFactor section 7.1 in applying the correction term to 3r, B2. Adam is often written with the correction term applied to ur but they equivalent (54).

---

10
Hir-Huee 0.99
Zero-shot imagellet accurac
58.0
MT--Huge
default
clipping
update cipping (= StableldamW)
57.0
Loos
56.0
$5.5
grad clipping
-
up date clisping (- StabickdamWi
2500 5000 7500 10000120011000105500000000000000 50001750020000
0.5
0.8
0.9
0.95
0.98
0.99
0.995
Iteration
Figure 5: Adding update clipping to AdamW mitigates loss spikes and outperforms other interventions such as gradient clipping with norm Code for the AdamW-AdaFactor hybrid we recommend Adam + update clipping is in Algorithm The left plot shows loss curves for 32 = 0.99 while the right displays accuracy ablating over $2.
Measurement. We now discuss measurement of the aforementioned stuck-in-the-Pe -in--itttte-pattttt scenario and search for predictive relationship between this event and loss spike. We follow Shazeer and Stern (54] and measure the following nottmeen-square quantity, RMS: / EE [g /ua). u4 a good estimator for of then the aggregate quantity RMS, will be around 1. The stuck-in- mntk-in-the-past scenario described above corresponds to an RMS, >> 1.
As illustrated in Figure 3 we observe instability for high 32 in our experiments even though have 5k iterations of warm-up. While Shazeer and Stern [54] first recognize the out-of-date second moment estimator issue, in their experimental setting they only observe instability without warm-up.
We now aim to establish a predictive relationship between the stuck-in-the-paat scenario and loss spikes. We present initial results in Figure 4 where examine RMS, for the the visual transformer patch embedding layer, parameters in visual.convl..wigge... only. This figure illustrates few important findings: i) loss spikes tend to follow 1-8 iterations after an RMS spike, ii) loss spikes slow learning as recovery time is required, and iii). RMS, stays around 1 for lower (3..
This means that the expectation is computed over
As this is just one example, we further elaborate on the predictive relationship between an RMS spike in the embedding layer in Section Batiiinn/@thnee through Figures 15 [6 17 18 19 2nd For analysis purposes, we define a heuristic to characterize loss and RMS spikes in
weight We then show that 28 of 30 detected loss spikes follow RMS spike by 1-8 iterations, while the probability that
loss spike follows an RMS spike by chance is only 1%. Moreover, we find that the same predictive relationship does not exist for the RMS in other transformer layers.
3.5 StableAdam W: Adam W with update clipping from AdaFactor
This Section develops and tests StableAdamW (Algorithm 2 an AdamW-Adafactor hybrid.
To stabilize training, the AdaFactor optimizer divides the learning rate for iteration t by 1/ m)xx(RMM,, 1) They refer to this as update clipping. The effect to slow training when d is no longer a good estimator for g²-
3.4, our stability issues be traced to an out-of-date w, which is what led Shazeer and Stern [44] to update clipping, even though their stability issues are also solved with warm-up. Therefore, we port update clipping the standard AdamW optimizer with d and refer to the resulting Adam Adaffacctor hybrid as StableAdamm (Algorithm 2). A modification we make compute and divide learning rate by max(RMS, 1) independently each tensor, which is for implementation convenience. This means that the expectation will be computed independently for each layer to produce different RMS,.
As discussed in Section
We now test how StamleAdamW compares with other stability interventions such as gradient clipping or lowering 32. These results, presented in Figure (5 find that AtamleAdamW (i.e., AdamW + update clipping) outperforms these aforementioned interventions for CLIP ViT-Huge. While gradient clipping and update clipping both remove instability, update clipping performs better in terms of zero-shot ImageNet accuracy. With update or gradient clipping, higher B3 such as 0.99 tends to
"They actually introduce hyperparameter and use 1/ mammmaxXMMS/4 /d, 1). but recommend setting 1. 'We clip at global norm 1. We observed instability when trying instead of 1. We did not tune this further. but note that 1.0 is standard in, e.g., PaLM (I1). and Scaling Vision Transformers 66

---

perform better. Appendix F provides further commentary and implementation considerations for StableAdamW.
Limitations, broader impacts, and conclusion
We believe the main limitation of our work is that is exhauususiiiee For instance, we only simulate float8 training and our experiments focus solely CLIP-style training. In terms of broader impact, our work may enable additional CLIP models, whose broader impact is examined extensively by Section 7 of Radford et al. [46]. Finally, we believe that our findings on accelerating and stabilizing large muulti-modll model training will be broadly useful to the community.
Acknowledgements
For insightful discussions we thank Romain Beaumont, Yair Carmon, Mehdi Cherti, Brian Che- ung, Alex Fang, Gabriel Ilharco, Jenia Jitsev, LAION, Sarah Pratt, Christoph Schuhmann, Ross Whightman, and Sho Yaida. We thank Emad Mostaque and stability.ai ai for compute resources.
This work is in part supported by NSF IIS 1652052, IIS 17303166, DARPA N66001-19-2-4031, DARPA W911NF-15-1-054 and gifts from Allen Institute for Artificial Intelligence.
References
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. In Uncertainty in Artificial Intelligence, pages 1352-1361. PMIR, 2021.
[2] Haoli Bai, Wei Zhang, Lu Hou, Lifeng Shang, Jing Jin, Xin Jiang, Qun Liu, Michael R. Lyu, and Irwin King. Binarybert: Pushing the limit of bert quantization. ArXiv, abs/2012.15701 2021.
[3] Charlic Blake, Douglas Oir, and Carlo Luschi. Unit scaling: Out-of-thc-box low-precision training. arXiv preprint Xiv:2303. / /257, 2023.
Andy Brock, Soham De, Samuel L. Smith, and Karen Simonyan. High-performance large-scale image recognition without normalization. International Conference on Machine Learning, pages 1059-1071. PMI.R, 2021.
[5] Tom Brown, Benjamin Mann. Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voo Voss, et al. Language models are few-shot leamers. In Advances in Neural Information Processing Systems (NeurIPS), 2020. /https://arxiv www//ogg///////ooo hwngg/abs/20055. 14165
[6] Léopold Cambier, Anahita Bhiwandiwalla, Ting Gong, Oguz H. Elibol, Mchran Nekuii, and Hanlin Tang. Shifted and squeezed 8-bit floating point format for low-precision training of deep neural networks. In 8th International Conference Learning Representations, ICLR 2020, Addix Ababa, Ethiopia, April 26-30. 2020. OpenReview.net, 2020. URL htt.:/ttt//ppeeeeeiiieeeee
net//forum?id-BxxeeAAee.
[7] Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang. Yao Liu, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsich, al. Symbolic discovery of optimization algorithms. arXiv preprint arXiv:2302. 06675, 2023.
[8] Xinlei Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers. Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 9640-9649, 2021.
[9] Mehdi Cherti, Romain Beaumont, Ross Wightman, Mitchell Wortsman, Gabriel Ilharco, Cade Gordon, Christoph Schuhmann, Ludwig Schmidt, and Jenia Jitsev. Reproducible scaling laws for contrastive language-image learning. arXiv preprint 2rXiv:2212.07143, 2022.
[10] Minsik Cho, Keivan Vahid, Saurabh Adya, and Mohammad Rastegari. Dkm: Differentiable k-means clustering layer for neural network compression. arXiv preprint
1110108..6599 12659, 2021.
[11] Aakanksha Chowdhery, Sharan Narang, Jacoh Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint 2rXiv:2204.02311 2022.
[12] Jeremy M Cohen, Behrooz Ghorbani, Shankar Krishnan, Naman Agarwal, Sourabh Medapati, Michal Badura, Daniel Suo, David Cardoze, Zachary Nado, George E Dahl, et al. Adaptive gradient methods at the edge of stability. arXiv preprint axXiv:2207.14484 2022.
10

---

[13] Matthieu Courbariaux, Yoshua Bengio, and Jean-Pierre David. Binaryconnect: Training deep neural networks with binary weights during propagations. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee. Masashi Sugiyama, and Roman Garnett, editors, Advunces in Neural Information Prucessing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, newrips http./errpp.ccppppee//////////// wwtpcc/pp/pp./...//e
Quebec, Canada, pages 3131-3131, 2015. URL wtw.httss:///////
5b066912cd Abstract.html
[14] Mostafa Dehghani, Josip Djolonga, Basil Mustafa. Piotr Padlewski, Jonathan Heek, Justin Gilmer, Andreas Steiner, Mathilde Caron, Robert Geirhos, Ibrahim Alabdulmohsin, et al. Scaling vision transformers to 22 billion parameters. arXiv preprint arXiv:2302. 05442, 2023.
[15] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hicrarchical image database. In Conference Computer Vision and Pattern Recognition, 2009. htttps://ieeexplors
leee. .rrg http//dcmmmt//222644444
[16] Tim Dettmers and Luke Zettlemoyer. The case for 4-bit precision: k-bit inference scaling laws, arXiv preprint 2r2iv:22/2.09772 2022.
[17] Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm, int8 O:: 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.0755: 2022.
[18] Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantiza- tion. 9th International Conference on Learning Representations, ICLR, 2022.
[19] Emily Dinan, Sho Yaida, and Susan Zhang. Effective theory of transformers at initialization. arXiv preprint arXiv:2304.02044 2023.
[20] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogmition at scale. In International Conference on Learning Representations (ICLR), 2021. https://arxiv org/ 2010
11929
[21] Mario Drumond, Tao Lin, Martin Jaggi, and Babak Falsafi. Training dans with hybrid block floating point. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Con- ference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Mon- tréal, Canada, pages 451-461, 2018. URL http:::/proceeeeeee www.igs.eeriis... www..ps.cc/pappeeer.eeee..ceeeee 2018/hash/ (Bsaecaff666cc1ceee33e3cccc333
wwt.rs-Abstttt..ti html
[22] John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
[23] Angela Fan, Pierre Stock, Benjamin Graham, Edouard Grave. Rémi Gribonval, Herve Jegou, and Armand Joulin. Training with quantization noise for extreme model compression. arXiv preprint arXiv:2004.07320 2020.
[24] Flias Frantar, Salch Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323 2022.
[25] Justin Gilmer, Andrea Schioppa, and Jeremy Cohen. Intriguing properties of transformer training instabili- ties. To appear.
[26] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249-256. JMLR Workshop and Conference Proceedings. 2010
[27] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. Proceedings of the IEEE international conference on computer vision, pages 1034,1 2015.
[28] Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition (CVPR), 2016. ptttp:://rrxv..or
1512. 03385
[29] Gabriel Ilharco. Mitchell Wortsman, Ross Wightman. Cade Gordon, Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi. Ali Farhadi, and Ludwig Schmidt. Openclip. July 2021. URL httt:ttt:: 1rg/ 514377777 If you use this software, please cite it as below.

---

[30] Benoit Jacob. Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. Quantization and training of neural networks for efficient integer- arithmetic-only inference. arxiv e-prints, art. arXiv preprint arXiv:1712005877, 2017.
[31] Daya Khudia, Jianyu Huang. Protonu Basu, Summer Deng. Haixin Liu, Jongsoo Park, and Mikhail Smelyanskiy. Fhgemm: Enabling high-performance low-precision deep learning inference. arXiv preprint arXiv:2101.05611 2021.
[32] Schoon Kim, Amir Gholami, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. I-bert: Integer-only bert quantization. International conference on machine learning, pages 5506-5518, PMLR, 2021.
[33] Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2014, https://arxiv.orr http:/iww..rg/bb/////oa/// 6980
[34] Manoj Kumar, Mostafa Dehghani, and Neil Houlsby. Dual patchnorm. arXiv preprint arXiv:2302.01322 2023.
[35] Yanghao Li, Haoqi Fan, Ronghang Hu, Christoph Feichtenhofer, and Kaiming He. Scaling language image pre- vve-trainig via masking. arXiv preprint arXiv:2112.00794, 2022.
[36] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations (ICLR), 2016. htttps://arxiv.. http://wwww..rraab///boa////// www//a//........... 03983
[37] Nya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations (ICI.R), 2019. hataga///peenerrrree.eerrrrrrmap
[38] Naveen Mellempudi, Sudarshan Srinivasan, Dipankar Das, and Bharat Kaul. Mixed precision training with 8-bit floating point. CoRR, abs// 905. 12334, 2019. http://a/// wwwporgggaaaaa.oooee 12334
[39] Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev. Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXiv:1700.03740, 2017.
[40] Paulius Micikevicius, Dusan Stosie, Neil Burgess, Marius Cornea, Pradeep Dubey, Richard Grisenthwaite, Sangwon Ha, Alexander Heinecke, Patrick Judd, John Kamalu, ct al. Fp8 formats for deep learning. arXiv preprint or Xiv:2209. 05433, 2022.
[41] Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, et al. A theory on adam instability in large-scale machine learning. arXiv preprint arXiv:2304.0987), 2023.
[42] Gunho Park, Baeseong Park, Jung Kwon, Byeongwook Kim, Youngjoo Lee, and Dongsoo Lee. nuqmm: Quantized matmul for efficient inference of large-scale generative language models. arXiv preprint arXiv:2206.09557 2022.
[43] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), 2019. https://arzii: http://axivv....aab///aasa
hts////s/11911 19t./1912.co 01703
[44] Hieu Pham, Zihang Dai, Golnaz Ghiasi, Hanxiao Liu, Adams Wei Yu, Minh-Thame Luong, Mingxing Tan. and Quoe V. Le. Combined scaling for zero-shot transfer learning, 2021. https://arxiv. org.ooggg/a.o....o... 2111.
[45] Haotong Qin, Ruihao Gong, Xianglong Liu, Xiao Bai, Jingkuan Song, and Nicu Sebe. Binary neural networks: A survey. CoRR, abs/20004003333 2020. / www.://ggw..r..ooo../// 03333
[46] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning (JCML), 2021. https: //ta://///ii//ee/a org/ abs/ 2abb///1/ /21.//111.000 00000
[47] Jack Rac, Schastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint ar arXiv:21 2021.
[48] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning. pages 8821-8831. PMLR, 2021.
12

---

[49] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.0600 06125, 2022.
[50] Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. arXiv preprint UmXiv:1904.000 092377, 2019.
[51] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[52] Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ili,, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b-parameter open-access multilingual language model. arXiv preprint arXiv:2211.05100 2022.
[53] Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al. Laion-5b: An open large-scale dataset for training next generation image-text models. arXiv preprint arXiv:2210.08400 2022.
[54] Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference Machine Learning, pages 4596 4604. PMLR, 2018
[55] Sheng Shen, Zhen Dong, Jiayu Ye, Linjian Ma, Zhewei Yao, Amir Gholami, Michael W Mahoney, and Kurt Keutzer. Q-bert: Hessian based ultra low precision quantization of bert. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 8815-8821, 2020.
[56] Xiao Sun, Jungwook Choi, Chia- Yu Chen, Naigang Wang, Swagath Venkataramani, Vijayalakshmi Srinivasan, Xiaodong Cui, Wei Zhang, and Kailash Gopalakrishnan. Hybrid 8-bit floating point (HFP8) training and inference for deep neural networks. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer. Florence d' Alché-Buus Emily B. Fox, and Roman Gamett, editors. Advances in Neural Information Processing Systems 32: Anmual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 14, 2019, Vancouver; BC, Canada, pages 4901 4910, 2019, URL htttps://proceeeeee neurips cc/ ptte:/aaeer/aae 2ww///aah///eeeeee 651
ca2d398f Abstract.hter
[57] Philippe Tillet, Hsiang- Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10-19, 2019.
[58] Hugo Touvron, Matthiuu Cord, Alexandre Sablayrolles. Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In Proceedings of the IEEE/CVE International Conference on Computer Vision, pages 32-42, 2021.
[59] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Maric-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.1397 2023.
[60] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017.
[61] Naigang Wang, Jungwook Choi, Daniel Brand, Chia-Yo Yu Chen, and Kailash Gopalakrishnan. Training deep neural networks with 8-bit floating point numbers. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.. 2018 URL http:ttttps///proreeiinnsse .eiggg...eriigg neurips.cc. cc/paper. Paper. pdf
[62]
Naigang Wang, Jungwook Choi, Daniel Brand, Chia-Yu Yu Chen, and Kailash Gopalakrishnan. Training deep neural networks with bi floating point numbers. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa- Dasa-Bianchi, and Roman Garnett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NewrIPS 2018, December 3-8, 2018, Montréal, Canada, pages 7686-7695, 2018, https:: httttpp:///rocceeingg html
[63] Shibo Wang and Pankaj Kanwar. Bfloat16: The secret to high performance on cloud tpus. 2019.
https://cloud. woogle..gooo....ooooooooo. httpg////cc///ooo////ottt//////s http://ootuctssta--aciiiieeee machine- machine-Lerrnnnggg
bfloat16- the- www.se-er-tttoeeeeeeeeeeeeeeee to-
high-
13

---

[64] Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon. Simon Kornblith, al. Model soups: aver- aging weights of multiple fine-tuned models improves accuracy without increasing inference time, In International Conference on Machine Learning, pages 23965 23998. PMLR, 2022.
[65] Guangxuan Xiao, Lin, Mickael Seznec, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211. 2022.
[66] Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang. Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-tratt post-training quantization for large-scale transformers. arXiv preprint arXiv:2206.08861. 2022.
[67] Aohan Zeng, Xiao Liu, Zhengxiao Du, Zihan Wang. Hanyu Lai, Ming Ding. Zhuoyi Yang. Yifan Xu, Wendi Zheng, Xiao Xia, et al. Glm-130b: An open bilingual pre- rre-trained model. arXiv preprint arXiv:2210.02414. 2022.
[68] Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin Dan Busbridge, Jason Ramapuram, Yizhe Zhang. Jiatao Gu, and Josh Susskind. Stabilizing transformer training by preventing attention entropy collapse. arXiv preprint arXiv:2.30..005 062966 2023.
[69] Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers, 2021.
https://arxiv. arg/abs//106..... 04560
[70] Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer. Sigmoid loss for language image pre-training. arXiv preprint vrXiv:2333.53344 2023.
[71] Hongyi Zhang. Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint arXivv1901.09321, 2019.
[72] Susan Zhang. Open pretrained transformers lecture, 2023. p9IxoSkv2-X
[73] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.0 01068, 2022.
[74] Wei Zhang, Hou, Yichun Yin, Lifeng Shang, Xiao Chen, Xin Jiang, and Qun Liu. Ternarybert: Distillation aware ultra-low bit bert. In EMNLP, 2020.
[75] Changsheng Zhao, Ting Hua, Yilin Shen. Qian Lou, and Hongxia Jin Automatic dixed-preccsion quantization scarch of bert. arXiv preprint 112.:2212..1438,, 2021.
[76] Kang Zhao. Sida Huang, Pan, Yinghan Li, Yingya Zhang, Zhenyu Gu, and Yinghui Xu. Distribution adaptive int8 quantization for training cnns. In Proceedings of the Thirty-Fifth AAAI Conference on Artificial Intelligence, 2021.
[77] Chenzhuo Zhu, Song Han, Huizi Mao, and William Dally. Trained temary quantization. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017 Conference Track Proceedings. OpenReview.net, 2017. URL httttpp://openrrveee.nett// net/:/oum//t/ooicc
S1_ S1_pAu9x1
[78] Feng Zhu, Ruihao Gong, Fengwei Yu, Xianglong Liu, Yanfei Wang. Zhelong Li, Xiuqi Yang, and Junjic Yan. Towards unified int8 training for convolutional neural network. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1969-1979, 2020.
14

---

Standard fp16 (sum of parts]
Unear layar, batch sequence length 32k
35
Unear layer summary, karying divensions
SwitchBack int8 (sum of parts)
50
Average fp15
30-
Average int8 matmul
Average quartize operation
time
20
28
15
sequerce length
10
10
betch sequerce length 32768 batch sequence length 055
batch
131072
5
4099
dim
dim
Figure 6: (Left) Individually profiling operations which constitute forward and backward pass in a lincar layer for i) SwitchBack using triton kernels and ii) fp16 bascline using torch.matmon Times averaged over linear layer from dim to 4 dim and linear layer from 4 dim dim--repreeentative of the linear layers in transformer MLP. (Right) The % speedup of SwitchBack over a standard fp16 linear layer when all operations in Figure 6 (left) are summed.
A Additional Related Work on Quantization
The literature on training neural networks in low-bit precision is vast. The main differentiating factor of our work is that we train relatively large models in fact, we train the largest 8-bit vision transformers to date.
The literature agrees that quantization of very large networks is more difficult than for smaller networks [88 17 655 24). As such, we divide our related work into three parts: (I) large-scale low-precision neural network (larger than BERT-large), and (2) low-precision training of smaller networks.
Large-scale Low-precision Neural Networks. Our work is currently the only work that does low-precision (8-bit and below) training of very large networks with more than 230M parameters. Other related work studies inference at scale. SmoothQuant (65), ZeroQuant [66], NuQmm 42 and LLM.int8() [77 study inference with Int8 matrix multiplication. Another line of work studies large models inference with more than 250M parameters by considering 16-bit inputs and k- k-bit weights [16 24 677.
Small Scale Low-precision Training Training of small-scale low-precision neural networks can take many shapes and forms, such as quantization for integer only devices, quantization for mobile device, or quantization accelerate training. One way to break up these directions is through the data type used and the neural network trained. One major direction is to quantize convolutional
neural networks often for fast and memory efficient usage edge devices [77 6 13
78
23
76 30).
Further work in this area is discussed in the survey by (45) Another line work is centered around 8-bit float data types which Lastly, a common application is to finetune (similar to training) BERT models to particular datasets. This not only decreases the model footprint and increases inference speed but adjusts the model to new data
be used accelerate training of neural networks [1] 56
62 38 3
32 74 55 75
B Achieving speed-ups with SwitchBack
We now test the speedups offered by SwitchBack by first examining individual operations and then end-to-end training.
We profile all of the operations which constitute forward and backward pass for single linear layer in Figure 66) (left) for both SwitchBack and the bascline. For SwitchBack we profile our custom triton kernels and for the baseline we profile Norch.matmin Overall, we observe that int8 multiplies occupy just over half the time as standard fp16 matmuls, and that quantize operations are roughly an order of magnitude less time than a matmul. Note that our int8 matmuls are fused with the dequantize operation.
Figure (right) displays the % speedup of SwitchBack over a standard fp16 layer when all operations in Figure 6 (left) are summed. Overall, the advantage of SwitchBack is greater for larger dim and batch_siize x sequence lenegtheeennghh. Overall, the speedup ranges from 5% to 35%. We see a bump at dim 1280 because standard PyTorc matmuls do not have optimized kernels for matrices of this
15

---

25
25
Quantization averhoad
syterch opt maxd Iear
engen
20
wnghn 133002
25
-5
10
1024
2048
4096 VST
dim
Figure (Left) Measuring the of time occupied by quantize operations for SwitchBack linear layer, which is usually less than 20% and decreases with dim. (Right) Benchmarking speedups for end-to-end CLIP training single node (with 4 A100 GPUs, per- Ber-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 projections as well the MI.P). speedups reported over i) custom linear layer implemented with SwitchBack that uses torch.autograd,
autch.a (Algorithm 5). which matches our implementation and ii) using the standard Py Porch Ln.Linear which includes additional hackground C++/CUDA optimizations which we do not replicate. LLM.int80 [17] does not provide speed-ups over the autorrautograd or nn. nn.eine baseline at this scale-we compare the speed of SwitchBack and Lnt in Figure 8
LLM nitil
$4.7
54.6
Specedu
54.4
54.3
54.2
$4.1
MrT
Large
MT Huge
Figure 8: Benchmarking speedups of SwitchBack compared to LLM.int80 [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 's autotune feature which provides fine-grained optimized kernels for matrices of any size. Our kernels are easy to modify as they are written Triton [57]. and the code run the benchmarks and produce Figure G is open sourced. In doing so, we invite the community to further improve the kernels and provide benchmark for measuring this progress. Due to computational constraints we have not tested dim > 4096 and it's possible the kemnlls 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 ( 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 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 77, 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. wwrch.autogra..liii linear (Algorithm 5) and ii) the PyTorch optimized linear layer Linear. In both cases the speedups increase when going from CLIP ViT-Base to CLIP ViT-Huge. However, there an additional 12..5 speedup when comparing SwitchBack to the baseline linear layer which uses torch. autograd We believe this comparison fair because SwitchBack also implemented using autograd,, while the standard Py PyTorch 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 12.5%% speedup. Finally, we note that the kernels from LLM.int8) [17 do not provide speedups over fp16 at the scale consider.
16

---

Additional code and figures
C.1
Additional Code
This Section provides additional pseudocode:
Algorithm the memory effecient variant of SwitchBack.
Algorithm the weights. For SwitchBackQ. the forward pass given by
A is the variant of SwitchBack which uses row- and column-wise quantization for
127 state (X) (tattee (W))
(W)T * Qrow (X) Crow (W)*
(4)
int8 matmul
where * is an elementwise product. Again, we append -tanspose to function in Algorithm| to mean that the operation is fused with a transpose.
Algorithm is a standard linear layer implemented with torch.autograd.
Algorithm 3 Memory efficient SwitchBackM
(llccd/ecc//sttmuuuttttrrrrrr Function): @staticnethod
dar fetvard(ctx, x, w) inputs
[n, weights
i_ins,, rox-vis=_quantize(()
W_ints, state_Ñ tensor-vise_quaniiee
save X_int,, state_x, N_ints, state_
output
return natami, natami_int,, ints_ and_ dequadeeee K_int8, V v_int8.t(,, state_X, state_l
Ostaticnethod
backrard(ctzz, G): [b, gredient to output
X_int,, state_X, V_int8, state_V caxe
dequantise_roo-xtta(....t state_X)
del inin
R_gradient natmul (p) X)
G_int8 rov-viso__uantzze((
del 8_int8 N_ints.tO w.ints.tO.cootiie..o.o. contiguous(
Use Sbit nataul only for K_gradient
X_gradient
intt, aadd__and_de_e dequamittt
intn, L_inta.t)), state_X, state_k
return ggradient, N_gradient
SvitelBackHLLneerr(nm Linear); def forvard X): return SottchBeccWMattullaalplll seighweight)
17

---

Algorithm 4 SwitchBack with row-wise and column-wise quantization for the weights SwitchBack
clasa SwitchBacckQWttmll (autograd.Fuun Function):
Ostaticnethon
foreard X,
X [b, n] inputa
[n, weights
save tensors in ctx have_for_backvard X,
ctx. save_for______e___eeeeeeeeeeeeee backward
x intn, state__ prx-otse__uantii20 N_int8, state_N pex-vise_quaatize
Return output
return natmul, into and and_anddddde dequanize
natmul_int@__am
K_int,, u_inta.t(), state_x, atate_k
@ataticnethod def backvard(ctz, G): G [b, m] gradient to output
Recover tensora from ctx X, M stw.saaeeforrbbalee backvard
Grrowiise Sovvise_quantzeeCC
W_intB, state_N colum--zise_quane
only K_gradient
Xggaadiett matmul int8 aartt__and_eeee deq_aattu G_int, M_int8.t(), state_X, state_k
i_gradient aattmu__f16(C..tt X)
X_gradient, N_gradient
class SwitcchBackkL.inaar Litnarr ferv X): return BatBBBcccKKammmaaaappyccc selg.weight)
Algorithm A standard linear layer implemented with autograu
clasa Standaadd (autograd. Function):
@staticmethod
daf (ctx, X, W) inputa
M [n. weights
save ctx. save_for wawe_for_baccoard X,
output
return torch.matmul (X, R.tO)
dstaticnethod
def backvard(ctz, Q):
G [b, gradient to output
Recover tensors from ctx x, V fox. backsard
Xgraddiett
(G, W)
W_gradient torch.matmul (G.+(,, X)
return graddinnt, N_gradient
class Standardlinesss (ww.inneeennnnnninnana Linear):
(srlf, X):
return Standard (x self.
18

---

C.2
Additional Figures
This section presents additional figures.
Figure
shows the loss curves throughout training for the models in Figure !
Figure
10
presents more fine-grained version of Figure 6
Figure 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 loss spikes. They find that loss spikes increase with high batch size and learning rate, but can be mitigated by reducing AdamW B2. However, reducing B2 by too much can slow learning.
Figure 14 shows that using schedule for 32 of the form 1 iteration - does not improve accuracy.
VIT-Base model
VIT-Huge model
Loss
Loss
Loss
10000 5500 20000 Iteration
5000 10000 15000 20000 iteration
5000 10000 15000 20000
5000 10000 15000 20000 Iteration
Iteration
SwitchBack
LLM baseline
1o8 tensor-wise baseine
Figure 9: Loss curves for the CLIP ViT-Base and CLIP VIT-Huge models evaluated in Figure | The left two plots display results for int8 training while the right two plots display results for float8 (fp8) training
19

---

Standard fp16 (sum parts)
Uneer sequence 32k
dimersonss
Matmul XW [standent)
30
Matmu (standard)
Sspeedup
2.5
Inc8 Natmul XW (switchback)
30
Int8 Matmul CW (swichback)
20
quantize (switchback)
20
Row wise quantize (swhchback)
Tersor quartize W (switchback)
betch soquercee length 16284
quartize and rspose W (switchback)
batch sequurcce length 32768
length
botch sequence length 131072
1024
2048
4096
1024
2048 dim
4096
dim
Figure 10: more fine-grained version of Figure 6
ViT-Huge,
VIT-Huge,
Tronsformer blck output
207
MLP weight gradiem
207
10-1
2500
5000
7500
10000
12500
15000
17500
20000
2500
5000
75.00
100000
12500
15000
17500
20000
Iteration
Iteraatio
VIT-Huge, layer caleinit
VIT-Huge, layer scale init 0
Transfomer block blockkoul output
300
MLP weight gradient
10-1
20-4
101
10
10a
10:1
10--
2500
3000
7300
10000
12500
13000
1/300
20000
2500
3000
75:00
10000
12500
13000
15000
20090
Iteration
Iteration
VIT-Large, layer
ViT-Large, layer
Trinnfomerr block output
MLP weight gradtent
10-2
107
20.5
301
20
10
VIT-Lage, layer scale Init
VIT-Large, layer scale Init
Trandfomer block output.
300
MLP weight gradient
307
10-4
2500
5000
7500
10000
12500
15000
17500
20000
2500
5000
75:00
10000
12500
15000
17500 20000
lteration
Iteration
Block
Black
Bock 20 moan
90
Figure 11: The mean and max for (left) the gradient to MLP weight and (right) the output of transformer block throughout training. Different rows correspond to different choice of model size and layer scale.
20

---

Base model batch 4096
Base model. batch size 16384
Base model batch size 65536
3000
12000
15000
20000
5000
10300
15200
20000
5200
10000
19000
20000
Beration
Iteration:
[Eer attanion
betad
0:90
beta2
beta2
Figure 12: Loss spikes increase with batch size for fixed learning rate and model size. Reducing AdamW 82 from its default in PyTorch of 0.999 mitigates loss spikes. Reducing 3n much slows training.
Base model. LR Se-4
Base model, LR 2e:3
Base model. LR 5e-3
LOss
6
5000
10000
15000
20000
3000
10000
15000
20000
5000
10000
15000
20000
Iteration
Iteration
beta2 0.99
botb2 0.98
beta2
0.95
beto2 0.9
beta2 0.5
Figure 13: I Loss spikes increase with learning rate for fixed batch size and model size. Reducing Adam W 88 from its default in PyTorch of 0.999 mitigates loss spikes. Reducing much slows training.
VIT-uge
58.0
57.5
ImageNet
57.0
56.5
56.0
55.5
updalle clipping
update clipping beta2 warmup
0.98
0.985
0.99
0.995
Beta2
Figure 14: We try a schedule for 3g which is used in AdaFactor 32 warmup. This means that 82 at iteration & is iteration and show the r-axis the final iteration. This warm-up does improve accuracy setting.
54 and PaLM 11 and refer to the experiment In this Figure try À 0.45, 0.5, 0.65
21

---

D Analysis
Consider a matrix multiplication UV for E Brnk and E Rkxm, This matmul consists of
computing inner products between vectors of length k.
This section shows that error due quantization increases with 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 high precision, & is batch size times sequence length, which often & 32000 in our experiments. For the other operations which comprise matmul, k is less than 4 embed dim which is v 8000 in our experiments. These dimensions are standard for CLIP training experiments [46. 99.
D.I Analyzing variance due to quantization for inner products
This section measures the variance due to quantization for the inner product between and v. Let v vectors of length & vectors with each element drawn i...d. from a distribution with mean 0. Let u; have variance of and v, have variance oy-
Next, let and be the quantized versions of and v, respectively. We model quantization error as úg u4 + and ig = + & where e,, & are i.dd mean centered random variables with variance 9.
The aim this section to show that variance due to quantization grows with k. Our analysis is conservative because do not assume the variance of G, ç, increase with k though in practice believe they would as the absmax of u and v increases with k.
We first examine the variance of ûzú; By using that all random variable are mean centered, this variance is given by,
Var(ū,;;) E ((;
(5)
E [(u + Es) (v; + (:))
(6)
e,ū; + E. + + E [(u,e] +
(7)
E [ufu]] + E + G + Var(u,e)) + o0 +
(8)
(9)
Next, we use linearity of variance for independent random variables to calculate Var ((ù, 0)). This given by,
Var ((ù, 81)
Var (e)(
(10)
Σ Var(u,ve) + Σ o
( of oz
(11)
Var ((u, e)) + k. 00 +
o + o
(12)
D.2 Takeaways
We have shown that for inner products with length & vectors, variance due to quantization increases with A: This means the variance of output units/features due to quantization increases with i which can thought of making the outputs more noisy. Noise compounds throughout the network and will eventually drown out useful signal-for large the network features or gradient will no longer lead to effective learning.
D.3
Why LLM.in88) fails: LLMs CLIP models
This Section details our hypothesis for why SwitchBack outperforms I.LMMint88 for CLIP training, which is conditioned on the analysis in Section (.T 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 GPT-3 5 or LLaMA [59], compared to CLIP models [46.
22

---

Pretrained large language models (LLMs) tend to have larger weight matrices relative to their batch sizes when compared CLIP models. CLIP models perform best when the batch size is As consequence, L.MSs and CLIP models have their most noisy operations for different matrix multiplications. L.MMs are most noisy in the forward pass XWT and during layer-to- layer back propagation Y& = x 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 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 good model for training.
large 46 44 Ω
For CLIP training with ViT-Huge, we have 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 x 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 Section D.] is a good model for training. Notice that the CLIP weight gradient is twice as noisy compared to the most noisy LLaMA 65B operations we assume that all quantization operations have the same error variance.
As such, low-precision LLM training and CLIP requires different parts of the training.
quantization routines for
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 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 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 that the reader can also evaluate if these heuristics appear reasonable.
We define RMS spikes events as (t: RMS, 2 2.3) while loss spike events are defined as the set of t where loss at time f 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 B2, 14 out 15 come 1-8 iterations after an RMS spike in the patch embedding layer www.uull.oovl..eeeee.... spike follows 1-8 iterations after an RMS spike by chance
With only 76 total RMS spike events, the probability that loss 1%.
Figure repeats this analysis for ViT-Large, wherein 13 out of 15 loss spikes follow an RMS spike by 1-8 iterations. The probability that loss spike follows an RMS spike by chance is 1.0%.
Figure 17 zooms in 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 Figure 15 but examines the RMS of random layer in the middle of the transformer-not the patch embedding layer. In this case, none of the loss spikes follow RMS spikes.
23

---

10.0
7.5
Hage 6dta
Loss
Loss spike 1-8 kerations
after RMS spite (? nut 7)
itarations
after RMS spliko (0 out of
0.0
2500
12500
17500
Belaz
Ivisual. canial..catt. weight RMS spike (total 251
RWS
2500
5000
7500
10000
12500
15000
17500
20000
10.0
(VFillage, Bebz 0.90
7.5
Loss
kerations
after RMS spite out B)
RMS splik
0.0
2500
7500
12500
15000
17500
conv2. weight)
RAN5
RMS spike (total 411
2500
5000
7500
10000
32500
15000
17500
20000
10.0
ViF4lage Beta? 0.35
kerations
spike
2.5
Loss spike not 1-8 iterations
spike
0.0
2500
5000
7500
10000
32500
15000
17500
20000
Hage
RMS Ivisual convi. weight) RMS spikee (totai 10)
RMS
10.0
Hige 5ata2
Loss
5.0
Loss kerations
spike out
Loss spiice not 1-8 iterations
2.5
after spike (0 of
2500
5000
7500
10000
12500
15000
17500
20000
MiTHugge 8eb2 0.3
RMS Ivisual.comv1. weighti spiie
RMS
0-
2500
7500
10000
12500
15000
17500
20000
Iteration
Figure 15: Observing a predictive relation between RMS spikes and loss spikes. For CLIP ViT-Huge and multiple 32 values, we use hcuristies (Appendix E to automatically identify loss spikes which we use for analysis. Out of 15 total loss spikes, 14 follow RMS spike in the patch embedding layer (RMS 2.3) by 1-8 iterations. We show loss spikes which identified by our heuristic. We use green if they follow RMS spike otherwise use red. RMS Spike indicates that the second moment estimator out of date (see Section and Shazeer and Stern
(54) ). The chance that a loss spike follows 1-8 iterations after an RMS spike by chance is
1%.
24

---

10.0
VIrLarge 0.39
7.5
Loss spite 1-8 iterations after RMS spike 19 out of 10)
Loss
iterations
after RMS spike (1 out of 10)
0.0
7500
12500
17500
Belaz 0.99
Ivisual.cam canial weight RMS spike (total 521
RM5
2500
5000
7500
10000
12500
15000
17500
20000
10.0
MiFllarge Beta? 0.90
7.5
Loss
Loss
after RMS spite out 4)
kerations
RMS splik
0.0
2500
7500
12500
15000
17500
weight)
RMN
RMS spikke (total 29
2500
5000
7500
10000
12500
15000
17500
20000
10.0
VTF- srge. Beta? 0.95
kerations
spike
2.5
Loss spike not 1-8 iterations
spike
0.0
2500
5000
7500
10000
12500
15000
17500
20000
Lirgar 0.35
RMS Ivisual convi. weight)
RMS
RMS spikee (totai 15)
10.0
Large Bataz 0.3
Loss
5.0
Loss kerations
spike out
Loss spiice not 1-8 iterations
2.5
after spike (0 of
2500
5000
7500
10000
12500
15000
17500
20000
Mar Beta 0.8
RMS Ivisual.comv] weight spike
RMS
0-
2500
7500
10000
12500
15000
17500
20000
Iteration
Figure 16: Observing a predictive relation between RMS spikes and loss spikes. For CLIP ViT-Large and multiple 32 values, we use hcuristies (Appendix E to automatically identify loss spikes which we use for analysis. Out of 15 total loss spikes, 13 follow RMS spike in the patch embedding layer (RMS 2.3) by 1-8 iterations. We show loss spikes which identified by our heuristie. We use green if they follow RMS spike otherwise use red. RMS Spike indicates that the second moment estimator out of date (see Section and Shazeer and Stern 1.0%.
(54) ). The chance that a loss spike follows 1-8 iterations after an RMS spike by chance is
25

---

10.0
y Melaaae Bela? 0.99
7.5
Loss
meerations
after RMS spike
spiko
0.0
2800
3000
3200
3400
3600
3800
4000
4200
(iT.Hage.
RMS (visual. convi weight)
RMS
RMS spike
0
2800
3000
3200
3400
3600
3800
4000
4200
Iteration
Figure 17: Zooming in on section of Figure 15
ViT-Hage,
7.5
Loss
Loss spike 1-8 Iterations after RMS spike
5.0
1-8 iterations
after RMS spike
0.0
4100
4200
[visual.conv weight)
RMS
RMS spike
4100
4200
Iteration
Figure 18: Examining the "failure" case in Figure 15 red loss spike does not really appear to be spike at all. However, adjusting our heuristic led to the issue of true spikes not being identified.
We believe this is not really failure as the non-predicted
10.0
VIT arge, Reta
7.5
Loss
Iterations
after RMS spiiee
spike iterationn
0.0
5600
5050
5700
5750
5000
5850
5900
5950
6000
4
Vil-Lagee,
(visual.comw1 weight)
RMS
RMS spike
5600
5700
5800
5900
6000
10.0
Large. 6cta2 0.85
Loss
5.0
spike iterations
RMS spike
2.5
Loss spike not 1-8 Iteratioon
IMS spike
45000
4550
4600
4650
4/00
4730
LVrTLagge, Betsž
0.95
RMS (visual. wwn.w...ww...ooommooooommm.. weight)
RMS
RMS
Iteration
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.
26

---

Loss
5.0
7.5
Loss spice 1-6 Rerations after RMS spike (9 out of 7)
Loss spike not 1-8 iterations
2.5
after RMS spike (7 out of 7)
2500
5000
7500
10000
12500
15000
17500
20000
ViF-Ilage Belaz - 0.99
RMS [blor 20 20....ttttnn..m.... attnnin ppan spike (total
RM5
0
2500
5000
7500
10000
12500
15000
17500
20000
10.0
ViE-Hage 66az
7.5
Loss
Loss spike 1-8 kerations after RMS spite (0 nut 8)
itarations
after RMS spliko (8 out of
0.0
7500
17500
Belaz
RMS (block prap
RWS
RMS spike (total 4)
2500
5000
7500
10000 Itaration
12500
15000
17500
20000
Figure 20: This figure repeats part of Figure 15 but examines the RMS of random layer in the middle of the transformer blocks. wtw....t.innn.r......... iaaiiinn_proooo not the patch embedding layer. RMS spikes no longer precede loss spikes.
StableAdam W continued
F1 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 [ 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 advocate using a hybrid. We tried porting other features from AdaFactor to AdamW such the 82 schedule but did not find them to help (Figure [4]. Morcover, while PaLM uses an AdaFactor-Adant hybrid, we believe they don't use update clipping.
Another question is, why not updates by any value, and is therefore immune to the stuck-in suck-in-the-pass scenario. We believe this may be promising path forward. However, while we observe that Lion outperforms Adam W at small scale, Lion still slightly under-perfoor mnder-performs Adam for CLIP ViT-Huge seale in experiments.
optimizer such as Lion 7 which does not divide
A final question is, why consider of in the numerator for computing RMS, and not vf? We also tried o² and found the performance worse.
F22 Implementation considerations
To prevent divide by O issues when computing RMS, we compute RMS, /E htf/mmaxxmmmmu,e e2]] where is the AdamW hyperparamer for which we use le-6 and maximum is elementwise maximum. This is instead of RMS, / [g//u..
G Loss spikes and the loss scalar
This Section ties the low precision training results 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 more limited representible range. Therefore, reducing loss spikes is important step for successful low precision training.
27

---

[
2.4
Fexture max
Love
2.2
25
Less
M010
1-/1
2900
2925
2950
2975
3000
2500
2925
2950
2075
3000
2900
2025
2950
2975
3000
Iteration
Iteration
Beration
7000 dliis beiiiig)
104
(austtrarsk restreckk aEE
Grad atamas
10
10
2900
2925
2950
2975
2000
2300
2925
2050
2975
3000
2900
2925
2350
2975
3000
Iteratiinn
retratiio
Bteation
RMS spike embedsing layer which preceeds loss spike
Inf gradient causing grad scaler decrease
Figure 21: Avoiding loss spikes is helpful for low precision training. As shown in this figure, loss spikes coincide with with activation spikes and gradient spikes. Large activations/rradients precision training a more limited representible range
issues during low
Supporting data is illustrated by Figure 21 in which an RMS spike precedes loss spikes which coincides with spikes in the activations (i.c., 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 an important step to enabling low- now-precision training. Second, spikes in gradient magnitude can be transient and therefore 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 the main layers where Inf/NaN encountered. Concretely, while we only track every tenth block, we never observe any Inf/NaN for any transformer block greater than O. However, with the PyTorcc default loss scalar an Inf/NaN in single layer will skip the update for the whole network.
few takeaways from this observation.Fiis First, reducing loss spikes is
This motivates the loss scalar that we use in our experiments when one is required (except for in Figure 2TP. We use 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 ViT-Huge scale where previously the scalar became too low and training diverged (9. We also believe adaptive block-wiss wise scalar as Ramesh al. would remedy this issue. One interesting remark 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 recovers the stability solution of Chen et [8] which is freeze the embedding layer. As final remark, 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.
28


Reviews:
[{"review_id": "fOXhqiBccG", "reviewer": "Reviewer_AbFj", "summary": "The authors introduce new methods for  accelerating and stabilizing training for large language-vision models. For acceleration, SwitchBack is proposed, which use high-precision for backwardd pass to compute the gradients for the weights. For stability, the introduce an\nAdamW-Adafactor hybrid (StableAdamW).", "questions": "See the weaknesses.", "limitations": "yes", "rating": 3, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "1. This work aims to accelerate and stabilize the training of language-vision models, which is particularly important for large-scale language-vision applications.\n2. The analysis of relationship between loss spike and AdamW second moment is interesting.", "weaknesses": "1. My main concern is about the novelty. The idea of using high-precision for backward weight gradients is studied by previous works such as the Gradients Bifurcation methd [1]. Float8 training are also widely studied. For the proposed StableAdamW, I didn't see the essential difference with AdaFactor, despite the explanation around line 255.\n2. It seems that the proposed methods are not designed particularly for language-vision models. It is better to compare with existing methods on traditional CV or NLP tasks.\n\nminor issues:\n\nline 331, divides the learning rate for iteration $t$ by $max(RMS_t,1)$.\n\nline 154, why constrains the float8 with in -1 and 1 ? float8cast(x/absmax(x))\n\n[1] Scalable Methods for 8-bit Training of Neural Networks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your review, we hope this rebuttal addresses your concerns. \n\n**Novelty**\n\nWe thank you for highlighting the gradient bifurcation method, and we will update the paper to include this important reference. As prior work has observed (e.g., [1]), quantization becomes more difficult at scale. Therefore, we believe that our experiments in this large scale setting, and for large vision-language models in particular, are valuable for the community. We emphasize that the largest vision models that have been studied before have been an order of magnitude smaller. In addition, we show that the technique can be simplified for float8 vs. int8.\n\nMoreover, we do not claim methodological novelty over AdaFactor, but instead aim to draw the community's attention to an important feature of AdaFactor that is often overlooked, and show that this feature can be effectively transferred to Adam. AdaFactor is decreasing in popularity as it has been shown to underperform Adam at scale [2]. However, our results indicate that this may be due to factored moments and not some of AdaFactor’s other innovations which can effectively be transferred to Adam. Finally, our work is novel as we establish a predictive relationship for loss spikes.\n\n**Why CLIP?**\n\nWe focus on the contrastive vision-language pre-training setting because we believe that it is of increasing importance as it underlies approaches from generative models to zero-shot classifiers. While the results may be more general, contrastive approaches require a very large batch size. As we discuss in Section D.3 of the supplementary material (L686), we believe that this distinction explains why LLM.int8() underperforms the bfloat16 baseline in the setting that we consider.\n\n**Why constrain the float8 to [-1, 1]?**\n\nWe apologize for our lack of clarity – this is a typo. After converting to [-1, 1] range we multiply by the max value in float8.\n\n[1] https://arxiv.org/abs/2208.07339\n\n[2] https://arxiv.org/abs/2112.11446"}, {"review_id": "5xRBNICdBC", "reviewer": "Reviewer_sdBF", "summary": "This paper proposed methods for accelerating and stabilizing CLIP training. To accelerate training, the authors proposed the SwitchBack method, which quantizes the precision to int8 for the first two matrix multiplies but switches back to higher precision for the weight gradient. The proposed method speeds up 13-25% CLIP ViT-Huge training. To stabilize training, the authors proposed to combine the AdamW and the Adafactor optimizers and call the proposed new optimizer StableAdamW. This is observed from the experiments that when the squared gradients become underestimated by AdamW's second-moment estimator, loss spikes will occur after several iterations. ", "questions": "$x_n$ should be $x_b$ in Eq. 1? \n\nWhy only examine RMSt for the visual transformer patch embedding layer, visual.conv1.weight? \n\nWhy not compare SwitchBack with the quantization method in [1]?\n\nReference:\n\n[1] Bai et al., Towards Efficient Post-training Quantization of Pre-trained Language Models, NeurIPS 2022. ", "limitations": "Mentioned in Weakness. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The motivation for the ideas is very clear. To train large-scale language-vision models, speed is very critical. Training stability is also very important because we need to train with large-scale data without performance degradation. These two factors are essential for training large-scale language-vision models. \n\nThe authors proposed the SwitchBack method that quantizes the precision to int8 for fast matrix multiplication and then transforms back to the original floating point precision for the weight gradient. The authors provide PyTorch torch illustrating this process in Algorithm 1. Experiments show that the proposed method could speed up 13-25% of CLIP ViT-Huge training. \n\nThe authors proposed the StableAdamW optimizer for training large-scale language-vision models. The authors observed that if the squared gradients are underestimated by AdamW's second-moment estimator, then the loss spike would occur in the subsequent few iterations. The authors proposed to combine the AdamW and the Adafactor to get their proposed StableAdamW optimizer. Experimental results show that StableAdamW stabilizes training and helps the model achieve higher zero-shot performance than the AdamW optimizer and the AdamW optimizer with gradient clipping. ", "weaknesses": "The experiments are only conducted for the CLIP training. It's unknown whether the proposed SwitchBack and StableAdamW will work for other language-vision pretraining models such as BEiT v2, BEiT v3, or BLIP. Also, since CLIP only uses contrastive loss, it's not clear whether the proposed methods will work for other losses. \n\nIt's unclear whether the authors used the global contrastive loss computed across all GPUs or the local contrastive loss computed on each GPU. The global contrastive loss also affects speed, accuracy, and stability.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your insightful review and careful attention to details. We hope the following answers your questions.\n\n- Question1: Regarding the typo in equation 1. Thanks very much for this catch, this is indeed a typo.\n- Question 2: We find that the patch embedding layer is the source of instability, which also aligns with the findings of [1]. When we create the same plot for other layers, we do not observe the predictive relationship which we observe with the patch embedding. More information on this point is provided in L741 of the supplementary material, in particular Figure 20 of the supplementary material examines RMS and loss for a non-embedding layer.\n- Question 3: Thank you for highlighting [2], we will revise to include this important reference. However, [2] focuses on post-training quantization, i.e., quantizing after training, while our paper focuses on training in low precision. Training in low precision is a much more difficult problem.\n- Weakness 1: We agree that focusing on CLIP is a limitation, but also believe that large language-vision models are increasingly important as they underlie generative models and zero-shot methods, and that these models present a unique set of challenges. We will revise the paper to make this shortcoming more apparent.\n- Weakness 2: We use the global contrastive loss, and will revise to include this information.\n\n[1] https://arxiv.org/abs/2104.02057"}, {"review_id": "pOKXECRpy6", "reviewer": "Reviewer_rd6a", "summary": "This paper introduces an efficient and stable INT8 training method for models similar to CLIP. The proposed method offers a training latency improvement of 10-20%, which could account for a significant portion of training costs for larger models. The authors leverage LLM.int8() kernels for training, taking into account the quantization challenge for weights. Additionally, the paper explores the potential issues associated with FP8 training for such models.", "questions": "included in weaknesses", "limitations": "included in weaknesses", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is beneficial to the community due to the rising interest in low-precision training, such as INT8 and FP8. In particular, understanding how FP8, which is supported by H100 GPUs, performs across different models is essential. This work provides useful insights into these topics.\n- The authors propose system/hardware-aware methods, including open-sourcing triton implementations and fused kernels, which could prove valuable for many researchers.\n- The zero-init layer-scale method, a simple yet effective approach to tensor-wise quantization, is proposed.\n- The systematic study of loss spike offers many points for consideration and deeper analysis.\n- The introduction of the novel predictive power of patch embedding RMS is noteworthy.", "weaknesses": "- The paper's SwitchBack approach seems to lack novelty, as it simply observes the correlation of inner dimensions with quantization noise. The authors could explore more ingenious solutions to tackle the large inner dimension of weight gradient matmul, rather than just resorting to higher precision computing.\n- Figure 1 raises a concern as LLM.int8() might not serve as an adequate baseline. LLM.int8 was designed for LLMs, not models like CLIP. Additional PTQ baselines, including MRE methods, should be considered for int8 training.\n- There's a fundamental curiosity about the applicability of low-precision training like INT8/FP8. While we can establish an FP16 baseline for well-known models, it becomes challenging when we train newly updated models or datasets with no baseline. If the training fails, it's unclear if the failure stems from precision issues or model issues. Although this paper's analysis and proposal could assist in these processes, the argument still stands: why use low-precision training if stable results can be obtained with slightly slower speeds? Because this open question applies to all low-precision training and not solely to this paper, I will respect feedbacks from other reviewers in this aspect.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your insightful comments. We hope the following addresses any concerns.\n\n- Weakness 1: We absolutely agree that resorting back to higher precision is a weakness and a more clever technique could alleviate this. We thank you for highlighting this as we believe it’s one of the most promising avenues for future research, and will revise the paper to indicate this shortcoming.\n- Weakness 2: Our focus is on faster training via quantization and not quantization after training. Appendix E of the LLM.Int8() paper establishes it as a practical method for large scale quantized training, though it is commonly regarded as a PTQ technique.\n- Weakness 3: Thank you very much for this very insightful question about why quantized training is even warranted. Indeed, we believe that this is very much an open question. However, if the current scaling trend persists for more years, it is likely that training runs in the future will cost upwards of 100 million US dollars. Therefore, any reduction in training time could be the difference between whether a run is feasible or not. Given the potential savings, we believe institutions will (and perhaps already are) taking the risk."}, {"review_id": "GfRu047UQ3", "reviewer": "Reviewer_Nme8", "summary": "The paper is a study of different ingredients necessary to train an int8-quantized CLIP model. They mainly address two aspects: \n  1. quantization: they built on top of techniques previously applied to LLM inference (LLM.uint8) and expand it to the CLIP setting to optimize both inference and training; the 13-25% speedup of training and negligible (~0.1%) quality drop is achieved by carefully selecting which parts of the model are quantized to 8 bits (forward pass and part of backward pass) and which to 16 bits (part of backward pass), \n  2. stabilization: they observed spikes in the loss and correlate it with times where the state of Adam optimizer gets outdated; this is mitigated by applying the same trick to Adam which was already established in the literature to stabilize AdaFactor\n", "questions": "1.  Is the training instability present only when quantization is used, or even without it?\n2.  How to explain that loss curves for grad clipping and adam clipping are the same, but there is a difference in imagenet zero-shot accuracy in Table 5? Is this difference repeatable over several runs, or could it be measuring noise?\n", "limitations": "The authors have adequately addressed the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.  The authors manage to train the biggest quantized CLIP up-to-date, with negligible quality drop. Learning that this is possible to do is a valuable contribution.\n2.  The source of training instability at scale is identified and characterized, and a way to mitigate it is proposed.\n", "weaknesses": "The proposed quantization method targets linear layers, which are “usually >90% of compute in standard transformer models”. However, it seems that there is significant complexity introduced for a relatively small speedup of 13-25%. Was that expected? What are the remaining bottlenecks which remain to be solved? It would be interesting to get more details around this topic.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your thoughtful review. We hope the following addresses your interesting questions and comments.\n\n- Weakness 1: We believe there are a number of reasons that the speed-up is only 13-25%. For one, 16-bit baselines have had years of hardware support and optimization, while the same is not true for eight bit formats. We believe that more speedups can be expected as hardware support improves. For example, initial data from H100 GPUs indicates that 8-bit speedups are much higher for these GPUs, but we did not include these benchmarks since we have no access to H100 GPUs. In addition, as we illustrate in Figure 6 and 10 of the supplementary material, the quantization operations (i.e., changing from 16 bit to 8 bit precision) incur overhead which reduces the overall speed-up. However, as we show in Figure 7 of the supplementary material, this overhead decreases in proportion at scale.\n- Question 1: The instabilities we discuss in Section 3 are present even without quantization.\n- Question 2: We believe that this is not noise because we consistently observe a difference across all four values of beta2 that we try as illustrated in Figure 5 (right).\n"}, {"review_id": "53xu3uRVhj", "reviewer": "Reviewer_35W5", "summary": "The paper proposes methods to improve the training of large, vision-language models. On the one had, the authors propose a new linear layer for int8 quantized trainings, dubbed SwitchBack. On the other hand, a new optimizer is presented, StableAdamW, which results from combining AdamW with the update clipping technique proposed by AdaFactor [52]. These two simple but effective measures allow for faster and more stable training of large vision-language models.", "questions": "- Minor typo in line 282? “for in AdamW.”\n- I assume that code will be released. Could the authors confirmed? If the results presented here are indeed generalizable to vision-language models in general, it could have major implications for the advancement of the community.", "limitations": "Limitations have been addressed.", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Simply having to use 8 bit precision for the first two matrix multiplies of a linear layer (forward pass to computer output + backward pass to compute gradients of the input) is a powerful way to achieve notable speed-ups while keeping the same performance as when using 16-bit training. I find this a very insightful contribution that can accelerate training for many projects that leverage transformer architectures.\n- The contributed StableAdamW (AdamW + update clipping from AdaFactor) addresses instabilities that can arise during training of large vision-language models. This is a well-known issues in the community, which could have major impact for future research. The fact that StableAdamW outperforms other stability control measures such as gradient clipping or setting a lower β2 already opens up the door for very easy stability fixes across the board for vision-language models, at least for those CLIP-based.\n- Detailed comparisons and results are presented.", "weaknesses": "- Could be interesting if the authors briefly discussed whether the results found in this paper generalize to downstream tasks, e.g., object detection or semantic segmentation. In other words, can other researchers directly apply SwitchBack layers in other transformers architectures, such us OneFormer? Can we also directly use StableAdamW? This is mentioned in the limitations section, which is appreciated. But any further points on this could be valuable to the community, as many of the breakthroughs in classification tasks have been found not to always easily translate into gains on downstream tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your thoughtful comments and thorough review. \n\n- Weakness 1: We agree that the paper would benefit from a discussion on whether we can expect the results to generalize to downstream applications. Our analysis in Section D of the supplementary material suggests that the approach is best applied when the inner dimension of the matrix multiplication for the weight gradient computation is large, which is typically the case for large batch training. Moreover, we believe that one of the benefits of CLIP style models is that they can be used in downstream applications without modifying the weights. That is, we believe that our models may be used for a drop-in replacement in stable diffusion to provide speed-ups. We will revise the paper to elaborate on this discussion.\n- Question 1: Thank you very much for catching this, it is indeed a typo.\n- Question 2: We confirm that code will be released. Indeed, we have focused on making our triton kernels clear and hackable. While we have optimized our kernels to the best of our ability, we believe that by open-sourcing them and making them hackable, the community will be able to iterate and further the already observed speed-ups.\n"}]

Author rebuttal:

Meta-review: This paper develops a simple and effective strategy to enhance the training efficiency and stability of large vision-language models. Overall, all reviewers appreciate its clear presentation, strong motivation, and solid empirical results. A few concerns are raised, including 1) novelty seems not enough; 2) experiments are only on CLIP models; and 3) additional ablations/clarifications are needed. 

The rebuttal addresses most of them, but reviewers are still concerned that the technique novelty of this paper is not very strong. While the ACs acknowledge that this concern is legitimate, they recognize the value of this work and believe the research community will be interested in this type of work on accelerating the training of large-scale vision-language models. Weighing the pros and cons, the ACs conclude that the paper's advantages surpass its limitations and thus recommend its acceptance.
Decision: Accept (poster)