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

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.

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. - 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. - 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.

Questions

- Minor typo in line 282? “for in AdamW.” - 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.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

Limitations have been addressed.

Summary

The paper is a study of different ingredients necessary to train an int8-quantized CLIP model. They mainly address two aspects: 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), 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

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. 2. The source of training instability at scale is identified and characterized, and a way to mitigate it is proposed.

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.

Questions

1. Is the training instability present only when quantization is used, or even without it? 2. 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors have adequately addressed the limitations.

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.

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. - The authors propose system/hardware-aware methods, including open-sourcing triton implementations and fused kernels, which could prove valuable for many researchers. - The zero-init layer-scale method, a simple yet effective approach to tensor-wise quantization, is proposed. - The systematic study of loss spike offers many points for consideration and deeper analysis. - 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. - 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. - 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.

Questions

included in weaknesses

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

included in weaknesses

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.

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. The 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. The 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. It'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.

Questions

$x_n$ should be $x_b$ in Eq. 1? Why only examine RMSt for the visual transformer patch embedding layer, visual.conv1.weight? Why not compare SwitchBack with the quantization method in [1]? Reference: [1] Bai et al., Towards Efficient Post-training Quantization of Pre-trained Language Models, NeurIPS 2022.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Mentioned in Weakness.

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 AdamW-Adafactor hybrid (StableAdamW).

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. 2. 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. 2. 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. minor issues: line 331, divides the learning rate for iteration $t$ by $max(RMS_t,1)$. line 154, why constrains the float8 with in -1 and 1 ? float8cast(x/absmax(x)) [1] Scalable Methods for 8-bit Training of Neural Networks.

Questions

See the weaknesses.

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

2 fair

Contribution

1 poor

Limitations

yes

Response to rebuttal

Thank the authors for the reply! Most of my questions are addressed.

Response to rebuttal

Thank the authors for the response! My main concern is still the lack of novelty. The proposed SwitchBack is previously studied in Gradients Bifurcation, the proposed StableAdamW has no methodological novelty over AdaFactor as responsed by the authors. I understand the authors' statement about drawing the community's attention to the useful methods studied in this paper, however, I think the contribution is not enough to get accepted. I think it it critical for a training method to generalize well. Otherwise, as raised by Reviewer rd6a, 'if the training fails, it's unclear if the failure stems from precision issues or model issues'. Thus, I think it is critical to test the generalization ability of the proposed method to a wide range of tasks and models. I have also read other reviewers' comments, but I still tend to reject this paper based on the above two concerns, which are not well addressed by the authors' response.