[论文解读] Training and inference of large language models using 8-bit floating point

Sergio P. Perez, Yan Zhang+7 more|arXiv (Cornell University)|Sep 29, 2023

Parallel Computing and Optimization Techniques被引用 4

一句话总结

本文提出了一种可扩展的方法，通过动态调整权重、激活值和梯度的每张量缩放因子，实现使用8位浮点数（FP8）格式训练和推理大规模语言模型（最大达700亿参数）。FP8-AMAX方法在GPT和Llama 2模型上实现了与全精度FP16基线几乎相同的准确性，即使在大规模场景下也表现出强大鲁棒性。

ABSTRACT

FP8 formats are gaining popularity to boost the computational efficiency for training and inference of large deep learning models. Their main challenge is that a careful choice of scaling is needed to prevent degradation due to the reduced dynamic range compared to higher-precision formats. Although there exists ample literature about selecting such scalings for INT formats, this critical aspect has yet to be addressed for FP8. This paper presents a methodology to select the scalings for FP8 linear layers, based on dynamically updating per-tensor scales for the weights, gradients and activations. We apply this methodology to train and validate large language models of the type of GPT and Llama 2 using FP8, for model sizes ranging from 111M to 70B. To facilitate the understanding of the FP8 dynamics, our results are accompanied by plots of the per-tensor scale distribution for weights, activations and gradients during both training and inference.

研究动机与目标

• 解决在大规模语言模型训练和推理过程中FP8格式动态范围有限的挑战。
• 通过动态调整缩放因子，克服FP8中因下溢和上溢导致的数值不稳定性。
• 实现在最大700亿参数模型上端到端的FP8训练与推理，且精度损失极小。
• 利用FP8的计算优势，提供一种实用且硬件高效的FP16和INT8替代方案。
• 建立一种可靠的FP8量化方法，其性能优于大型模型中的静态缩放策略。

提出的方法

• 提出FP8-AMAX，一种动态缩放策略，基于每张量的最大绝对值，为权重、激活值和梯度计算每张量的缩放偏置。
• 通过compute_bias函数实现动态指数偏置调整，该函数基于格式最大值与张量振幅之比的log2值计算缩放因子，并支持可配置的余量。
• 在FP8操作前后应用缩放：在转换为FP8之前缩放输入，在矩阵乘法后解缩放以保持精度。
• 使用不同的FP8格式：权重和激活值使用FP8_E4，梯度使用FP8_E5，以在范围和精度之间取得平衡。
• 将该方法集成到训练（含梯度缩放和权重更新）和推理（训练后量化）流程中。
• 通过损失缩放搜索验证收敛鲁棒性，将FP8-CSCALE（固定缩放）与FP8-AMAX（自适应缩放）进行对比。

(a) Forward pass for FP16 inference.

实验结果

研究问题

• RQ1 FP8能否在不造成显著精度损失的情况下，有效用于大规模语言模型的训练和推理？
• RQ2 在不同模型规模下，动态每张量缩放（FP8-AMAX）与固定缩放（FP8-CSCALE）在收敛性和稳定性方面有何对比？
• RQ3 何种缩放策略可确保FP8中数值稳定性，防止权重、激活值和梯度的下溢和上溢？
• RQ4 缩放偏置在训练过程中如何演变？这对更新频率和计算成本有何影响？
• RQ5 FP8量化能否在多种大型语言模型架构和任务上实现与FP16相当的性能？

主要发现

• FP8-AMAX在MNLI、QQP和SST-2任务上，成功实现了从1.11亿到130亿参数的GPT和Llama 2模型微调，验证准确率均保持在FP16基线的99.5%以内。
• 对于130亿参数模型，FP8-CSCALE仅在极窄的缩放范围[-1, 0]内收敛，且不同随机种子下收敛性无法保证。
• FP8-AMAX在所有测试的模型规模（包括67亿和130亿参数）中均实现一致收敛，微调过程中未观察到发散。
• FP8-AMAX与FP16的损失函数轨迹几乎无法区分，表明训练动力学稳定。
• 每张量缩放分布显示，FP8-AMAX在整个训练过程中保持稳定且自适应的缩放，而FP8-CSCALE则呈现刚性缩放。
• 该方法成功扩展至700亿参数的Llama 2模型，证明了在大规模场景下实现完整FP8训练与推理的可行性。

(b) Forward pass for FP8 inference.