PaperlayerView exact revision on arXiv

Source-linked AI summary

Stable and low-precision training for large-scale vision-language models

Mitchell Wortsman, Tim Dettmers, Luke Zettlemoyer, Ari Morcos, Ali Farhadi, Ludwig Schmidt

arXiv:2304.13013v2cs.LGcs.CV

TL;DR

The paper addresses speed and stability in large-scale CLIP training with SwitchBack and an AdamW–AdaFactor hybrid. SwitchBack matches bfloat16 within 0.1 percentage points while accelerating CLIP ViT-Huge training 13–25%.

  • Problem

    Large-scale CLIP training must be fast enough to process substantial data and stable against loss spikes that degrade performance.

    Source: §1 · p. 1

  • Method

    The paper introduces SwitchBack for low-precision linear layers, zero-initialized layer-scale to discourage large feature magnitudes, and StableAdamW with AdaFactor update clipping.

    Source: §1 · p. 2, §2.3 · p. 8

  • Results

    SwitchBack matches bfloat16 within 0.1 percentage points and speeds CLIP ViT-Huge training 13–25%; StableAdamW removes instabilities and outperforms gradient clipping at tested scales.

    Source: §1 · pp. 1–2

  • Takeaways & Limitations

    Reducing loss spikes is an important step for successful low-precision training, and the authors share Triton kernels for further work.

    Source: §3.6 · p. 12, §4 · p. 13

    Float8 training is simulated, and checkpoints have low accuracy because compute was limited and the experiments prioritized loss-spike trials over competitive performance.

    Source: §4 · p. 13

Abstract

from arXiv · show

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.

1 Introduction

The paper studies speed and stability as critical elements of large-scale CLIP training. It proposes low-precision linear layers and an optimizer hybrid to address them.

1 Introduction

  • SwitchBack matches bfloat16 zero-shot ImageNet accuracy within 0.1 percentage point on CLIP ViT-Huge and speeds training 13% versus optimized PyTorch, 25% versus autograd.The figure also reports that SwitchBack outperforms LLM.int8().Source: Figure 1 · p. 1, §1 · pp. 1–2
  • In simulated fp8 training, tensor-wise baselines diverge above 420M scale, while zero-initialized layer-scale keeps feature magnitudes small enough for high accuracy.The authors report using layer-scale initialized with zeros to discourage large feature magnitudes.Source: §1 · p. 2
  • Loss spikes coincide with AdamW’s patch-embedding second-moment estimator underestimating squared gradients; StableAdamW avoids instabilities and outperforms gradient clipping at tested scales.StableAdamW applies AdaFactor update clipping, lowering the learning rate when the gradient-square-to-second-moment ratio is large.Source: §1 · p. 2

2 8-bit training

This section examines low-precision training for vision-language models, covering int8 CLIP training with SwitchBack and simulated fp8 training with zero-initialized layer-scale. It focuses on accuracy and training stability across these approaches.

2.1 Preliminaries and related work

  • At the time described, Ampere GPUs supported int8 while float8 support required Hopper GPUs, and the 1B-parameter CLIP ViT-Huge models made quantization more challenging than smaller models.The paper notes that Hopper GPUs were difficult to obtain at the time.Source: §2.1 · p. 3

2.2.1 Method

  • SwitchBack uses 8-bit precision for forward and input-gradient matrix multiplies but higher precision for weight gradients, whose long batch-by-sequence dot products increase quantization variance.The method uses row-wise quantization for inputs and gradients and tensor-wise quantization for weights.Source: §2.2.1 · p. 3

2.2.3 Results

  • For CLIP ViT-Huge, SwitchBack drops 0.1 percentage points vs bfloat16 in int8 and simulated float8; LLM.int8() drops 5.9 percentage points, while tensor-wise float8 training diverges.Figure 2 plots the CLIP ViT-Base and ViT-Huge loss curves for int8 and float8 training.Source: Figure 2 · p. 4, §2.2.3 · p. 6

2.2.1 Method

  • SwitchBackM saves only 8-bit tensors to reduce memory, with an extra backward dequantization and negligible ViT-Huge accuracy differences; SwitchBackQ showed no accuracy improvement at tested scales.SwitchBackQ uses row-wise and column-wise weight quantization instead of tensor-wise quantization.Source: §2.2.1 · p. 4

2.2.2 Experimental setup

  • The SwitchBack evaluation trained CLIP on LAION-2B for 380 million images with patch-dropout 0.5, then evaluated zero-shot on ImageNet.The study focused on comparing low-precision methods rather than achieving high final accuracy.Source: §2.2.2 · p. 5

2.2.3 Results

  • SwitchBack delivered 5%–35% speedups over fp16 per-layer and ~12.5% additional end-to-end speedup over torch.autograd, with end-to-end gains rising from ViT-Base to ViT-Huge.Quantization occupies at most 25% of linear-layer time, falling to around 10% or below for large dim.Source: Figure 3 · p. 5, Figure 4 · p. 6, §2.2.3 · pp. 6–7

2.3 Float8 training by reducing feature magnitude

  • Zero-initialized layer-scale enabled high-accuracy simulated float8 with tensor-wise quantization by discouraging large feature magnitudes; without it, tested methods diverged on ViT-Large.The simulation uses exact float8-representable values but performs arithmetic in float16.Source: Figure 5 · p. 7, §2.3 · p. 7
  • Residual connections can increase norms in deeper networks, while zero-initialized layer-scale makes the transformer an identity function at initialization.Layer-scale multiplies each self-attention and MLP block output by a learnable vector; although typically initialized to 10^-4 or 10^-6, it is set to 0 here.Source: §2.3 · pp. 7–8
  • Zero-initialized layer-scale controlled average output feature magnitude, which otherwise became high in later blocks.Figure 5 (right) demonstrates this intervention; large feature magnitudes are identified as a low-precision training concern [17].Source: §2.3 · p. 8
  • Only simulated fp8 ViT-Large training with zero-initialized layer-scale avoided divergence; tensor-wise fp8, gradient clipping, and KQ layernorm variants diverged.The comparison is shown in Figure 5 (left). The authors say the remaining difference from bfloat16 is primarily due to layer-scale and believe hyperparameter tuning could match standard-training accuracy.Source: §2.3 · p. 8

3 Stability

In CLIP training, loss spikes become more common as model size, batch size, or learning rate increases, and are linked to an outdated AdamW second-moment estimate in the patch embedding layer. The section develops StableAdamW to address these instabilities and connects loss spikes with low-precision training issues.

3.1 Preliminaries and related work

  • The study distinguishes fast CLIP loss spikes from slow divergence and finds they arise when the second-moment estimator is out of date in early layers.Unlike the instabilities in [14] [69], the authors report spikes despite a long warmup and identify the patch embedding layer as the primary site.Source: §3.1 · p. 8

3.2 Experimental setup

  • Experiments use short CLIP runs on LAION with 20k iterations, patch-dropout 0.5, and 5k warmup iterations.Unless otherwise noted, batch size is 16384, learning rate 2e-3, and weight decay 0.2.Source: §3.2 · pp. 8, 10

3.3 Loss spikes increase with model size, batch size, and learning rate

  • Loss spikes increase with model size, batch size, and learning rate; lowering AdamW β2 can avoid them, but lowering it too far slows learning and worsens performance.Figures 6–8 show the trends; the authors note that β2 reduction has a performance trade-off.Source: Figure 6 · p. 9, Figure 7 · p. 9, Figure 8 · p. 9

3.4 On β2 and an out-of-date second moment estimator

  • In the patch embedding layer, an outdated second-moment estimate underestimates squared gradients, and 28 of 30 detected loss spikes follow an RMS spike by 1–8 iterations.The predictive relationship was not found for RMS in other transformer layers; the chance probability reported is 1%.Source: §3.4 · pp. 10–11

3.5 StableAdamW: AdamW with update clipping from AdaFactor

  • StableAdamW ports AdaFactor update clipping to AdamW with d = 1, dividing the learning rate by max(RMSt, 1) independently for each tensor; with clipping, β2 = 0.99 tends to perform better.The intervention slows updates when the second-moment estimate is no longer reliable.Source: §3.5 · pp. 11–12
  • StableAdamW (AdamW with update clipping) outperforms gradient clipping in zero-shot ImageNet accuracy for CLIP ViT-Huge; both interventions remove instability.The comparison is shown in Figure 10.Source: §3.5 · p. 12, Figure 10 · p. 12

3.6 Loss spikes and the loss scalar

  • A fixed, per-tensor Inf/NaN loss scalar enabled fp16 CLIP ViT-Huge training, though loss spikes do not always lower it.Loss spikes can coincide with activation and gradient spikes, including Inf/NaNs; a scalar drop can also occur for other reasons.Source: §3.6 · pp. 12–13

4 Conclusion

The conclusion presents experiments on accelerating and stabilizing large multimodal training and shares Triton kernels. The authors identify non-exhaustiveness and limited-compute checkpoints as limitations.

4 Conclusion

  • The work shares experiments on accelerating and stabilizing large multimodal model training, along with Triton kernels for further development.The authors describe these resources as useful to the community.Source: §4 · p. 13
  • The study is non-exhaustive: it simulates float8 training and does not examine stability effects of width scaling or training hyperparameters such as those in [19].The authors also report that their checkpoints have low accuracy because compute was limited and their aim was to study loss spikes across trials, not achieve competitive performance.Source: §4 · p. 13

A Additional Related Work on Quantization

The related-work discussion distinguishes large-scale low-precision neural networks from smaller-scale training directions. It positions this work as training very large models at 8-bit precision and below.

A Additional Related Work on Quantization

  • The paper identifies its scale as a differentiator, reporting training the largest 8-bit vision transformers to date.The broader low-bit neural-network literature is described as vast.Source: A Additional Related Work on… · p. 19
  • The authors note that quantizing very large networks is more difficult than quantizing smaller ones [18] [17] [66] [24].They use this distinction to organize the related work.Source: A Additional Related Work on… · p. 19
  • The authors describe their work as the only low-precision (8-bit and below) training study of networks above 230M parameters at the time of writing.They distinguish it from related work on large-model inference.Source: A Additional Related Work on… · p. 19
  • The discussion separates small-scale low-precision work into CNN quantization for edge devices, float8 training, and BERT fine-tuning.The survey by [45] covers further work on CNN quantization.Source: A Additional Related Work on… · p. 19

B Additional code and figures

The appendix points to additional SwitchBack and standard-linear-layer pseudocode and describes how to select implementations for OpenCLIP training.

B.1 Additional Code

  • The appendix section supplies additional pseudocode for the implementations.Source: §B.1 · p. 19
  • Algorithm 3 is the memory-efficient SwitchBack variant, Algorithm 4 uses row- and column-wise weight quantization, and Algorithm 5 is a standard linear layer using torch.autograd.The appendix also identifies StandardLinear for Algorithm 5.Source: §B.1 · p. 19
  • The listed implementations are available in the bitsandbytes Triton-based modules, with OpenCLIP flags selecting global, memory-efficient, or vectorwise SwitchBack implementations.The mappings include SwitchBackLinearGlobal for Algorithm 1, SwitchBackLinearGlobalMemEfficient for Algorithm 3, and SwitchBackLinearVectorwise for Algorithm 4.Source: §B.1 · p. 19

B.2 Additional Figures

The additional figures provide finer-grained speed and feature-statistic views, and report that a particular β2 schedule does not improve accuracy.

B.2 Additional Figures

  • Figure 12 provides a more fine-grained version of Figure 3.Source: §B.2 · p. 19
  • Figure 13 compares SwitchBack's speed-up with LLM.int8().Source: §B.2 · p. 19
  • Figure 14 shows mean and maximum gradient and activation (feature) values throughout training.Source: §B.2 · p. 19
  • A β2 schedule of 1 − iteration^−λ does not improve accuracy.Source: §B.2 · p. 19

C Analysis

This analysis relates quantization error in inner products to their dimension and explains how SwitchBack assigns precision across linear-layer computations. It concludes that growing quantization noise can impair learning as it compounds through the network.

C Analysis

  • Quantization error increases with inner-product dimension k, and SwitchBack leaves the weight-gradient computation in high precision while using int8 for other matrix multiplications.The implementation computes the weight gradient with matmul_fp16 and the input gradient with int8.Source: C Analysis · pp. 19, 22

C.1 Analyzing variance due to quantization for inner products

  • The variance analysis models quantized vectors using additive, mean-centered errors and examines how quantization affects inner products.The vectors are modeled as i.i.d. and mean-zero, with i.i.d. mean-centered quantization errors.Source: §C.1 · p. 22

C.2 Takeaways

  • Quantization-induced variance grows with k, making output features noisier; as this noise compounds, features or gradients may no longer support effective learning.The analysis links the inner-product variance increase to noisier output units and potential loss of useful learning signal at large k.Source: §C.2 · p. 23

C.3 Why LLM.int8() fails: LLMs vs CLIP models

The section explains why CLIP and language-model training can have different high-noise operations under quantization. In CLIP, large batch sizes make weight-gradient computations especially noisy, motivating precision choices that differ by model type.

C.3 Why LLM.int8() fails: LLMs vs CLIP models

  • LLMs and CLIP have their most noisy operations in different matrix multiplications because LLMs have larger weight matrices relative to batch size, while CLIP performs best with large batches.Under the stated analysis, LLM weight gradients are 4x to 16x less noisy than LLM forward and layer-to-layer backpropagation operations.Source: §C.3 · p. 23
  • CLIP's weight-gradient computation is estimated to be 51.2x to 12.8x more noisy than its forward and layer-to-layer backpropagation operations.For ViT-Huge, the passage gives a batch size of 65536 per GPU and weight matrices of 1280×5120.Source: §C.3 · pp. 23–24
  • The section concludes that low-precision LLM and CLIP training require high-precision quantization routines for different parts of training.The authors suggest that LLM.int8() may fail for CLIP because its noisy weight-gradient computation may not provide enough signal for SGD to converge to a local minimum.Source: §C.3 · p. 24

D RMS Spikes precede Loss Spikes

The section tests whether RMS spikes in the patch embedding layer precede loss spikes using heuristics for identifying both events. The reported relationship is common in ViT-Huge and ViT-Large, but is not found for a random middle transformer layer.

D RMS Spikes precede Loss Spikes

  • The spike counts use heuristics: RMS spikes are defined by RMSt ≥ 2.3, and loss spikes by loss exceeding the running mean by 3.2 running standard deviations.The analysis ignores the first 1000 iterations and deduplicates events occurring within 10 iterations.Source: D RMS Spikes precede Loss Spikes · p. 24
  • 14 of 15 ViT-Huge loss spikes occurred 1–8 iterations after an RMS spike in the patch embedding layer.The chance probability for a loss spike following an RMS spike was reported as less than 1%.Source: D RMS Spikes precede Loss Spikes · p. 24
  • 13 of 15 ViT-Large loss spikes followed an RMS spike by 1–8 iterations, with a reported chance probability of 1.0%.Source: D RMS Spikes precede Loss Spikes · p. 24
  • Cases of apparent missed loss spikes mainly reflected false positives from the loss-spike detection heuristic, while none followed an RMS spike in a random middle transformer layer.Source: D RMS Spikes precede Loss Spikes · p. 24

E.1 Q&A

The Q&A argues for an AdamW–AdaFactor hybrid rather than AdaFactor alone, and reports mixed results for alternative optimizer choices and schedules. In the tested CLIP ViT-Huge setting, Lion slightly under-performed AdamW.

E.1 Q&A

  • The authors advocate an AdamW–AdaFactor hybrid because AdaFactor under-performs AdamW at scale, which they attribute likely to factored moments rather than update-clipping.They report that porting AdaFactor's β2 schedule to AdamW did not help.Source: §E.1 · p. 24
  • Lion outperforms AdamW at small scale but slightly under-performs it for CLIP ViT-Huge in the reported experiments.Source: §E.1 · p. 24
  • Using v^2_t rather than g^2_t in the numerator for computing RMSt produced worse performance.Source: §E.1 · p. 25

E.2 Implementation considerations

The implementation notes add an epsilon safeguard when computing RMS. Across CLIP ViT-Huge and ViT-Large, patch-embedding RMS spikes generally precede loss spikes, unlike RMS spikes in a sampled middle transformer layer.

E.2 Implementation considerations

  • An elementwise maximum with epsilon is used when computing RMS to prevent division-by-zero issues; the AdamW epsilon is 1e-6.Source: §E.2 · p. 25
  • 14 of 15 CLIP ViT-Huge loss spikes and 13 of 15 CLIP ViT-Large loss spikes followed patch-embedding RMS spikes by 1–8 iterations.The passages report chance probabilities of <1% and 1.0%, respectively.Source: Figure 16 · p. 26, Figure 17 · p. 27
  • The authors attribute apparent unmatched loss spikes to heuristic issues, while noting that changing the heuristic can miss true spikes.Source: Figure 19 · p. 28, Figure 20 · p. 29
  • RMS spikes in the sampled middle transformer layer did not precede loss spikes, unlike patch-embedding RMS spikes.Source: Figure 16 · p. 26, Figure 21 · p. 29

Related papers

Loading 2304.13013v2…