[论文解读] Stable and low-precision training for large-scale vision-language models

Mitchell Wortsman, Tim Dettmers+4 more|arXiv (Cornell University)|Apr 25, 2023

Multimodal Machine Learning Applications被引用 5

一句话总结

本文提出 SwitchBack，一种用于大模型视觉-语言模型稳定且高效 8 位（int8 和 float8）训练的新型线性层，在 CLIP ViT-Huge 上实现了与 bfloat16 训练几乎相同的精度，同时实现了 13-25% 的加速。此外，本文还提出了 StableAdamW，一种结合 AdamW 与 AdaFactor 的优化器，通过应用更新裁剪来消除损失尖峰，显著提升了大规模训练的稳定性。

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.

研究动机与目标

• 通过低精度（int8 和 float8）计算加速大规模视觉-语言模型训练，同时不牺牲精度。
• 通过识别并缓解导致模型性能下降的损失尖峰，提升训练稳定性。
• 实现首个成功的大规模 1B 参数 CLIP ViT-Huge 模型 int8 训练，其性能与 bfloat16 相比仅相差 0.1%。
• 开发实用且开源的工具，如 Triton 内核，以支持未来低精度训练的研究。
• 证明通过适当的初始化和自适应优化，即使在无原生硬件支持的情况下，也能实现稳定的 float8 训练。

提出的方法

• SwitchBack 使用 16 位矩阵乘法计算权重梯度，同时在前向传播和输入梯度计算中使用 int8，从而减少高维注意力层中的量化噪声。
• 该方法采用混合精度策略：推理和输入梯度使用 8 位运算，梯度累积使用 16 位，相比标准量化方法提升了精度。
• 对于 float8 训练，SwitchBack 通过精确的 fp8 模拟实现，结果表明其性能优于在大规模训练中发散的张量级量化基线方法。
• StableAdamW 将 AdamW 与 AdaFactor 的更新裁剪技术相结合，通过监测平方梯度与其二阶矩估计器的比值，动态调整学习率。
• 损失标量被修改为在张量层面检查 Inf/NaN 并保持固定，防止因单个层的不稳定性导致全局训练失败。
• 采用全零的层尺度初始化方法，以抑制过大的特征幅值，从而在使用标准量化技术时实现稳定的 float8 训练。

实验结果

研究问题

• RQ1 像 SwitchBack 这样的新型线性层设计能否在实现大视觉-语言模型 int8 训练的同时，提供显著的加速并保持高精度？
• RQ2 大规模 CLIP 训练中为何会出现损失尖峰？是否可以预测或预防？
• RQ3 在无原生硬件支持的情况下，能否使 float8 训练稳定化？哪些初始化或优化技术可实现这一点？
• RQ4 StableAdamW（即 AdamW 与 AdaFactor 更新裁剪的结合）与梯度裁剪相比，在稳定性与最终性能方面表现如何？
• RQ5 在低精度训练中，大激活值和梯度在导致不稳定性方面起什么作用？

主要发现

• 与 bfloat16 相比，SwitchBack 在 CLIP ViT-Huge 训练中实现了 13-25% 的端到端加速，零样本 ImageNet 准确率相差不超过 0.1 个百分点。
• SwitchBack 实现了迄今最大的 int8 训练规模，应用于 1B 参数的 CLIP ViT-Huge 模型，其精度优于 LLM.int8()。
• 使用标准张量级量化方法的 float8 训练在参数量超过 420M 时会发散，但结合使用全零的层尺度初始化后，SwitchBack 实现了稳定训练。
• StableAdamW（一种结合 AdamW 与 AdaFactor 更新裁剪的 AdamW-AdaFactor 混合优化器）可消除损失尖峰，在大规模训练中表现优于梯度裁剪，零样本 ImageNet 准确率更高。
• 损失尖峰在发生前，通常伴随 patch 嵌入层中平方梯度与其移动平均值的比值下降，表明二阶矩估计器存在低估现象。
• 使用固定、按张量的损失标量，在张量层面检查 Inf/NaN，可防止全局训练失败，并实现 ViT-Huge 规模下的稳定 fp16 混合精度训练。