[論文レビュー] 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

ひとこと要約

本論文では、重み、活性化、勾配の各テンソルに対して動的スケーリング係数を調整することで、最大700億パラメータの言語モデルを8ビット浮動小数点（FP8）形式でスケーラブルにトレーニングおよび推論する手法を提示する。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トレーニングおよび推論を可能にする。
• FP16やINT8の代替手段として実用的かつハードウェア効率的な選択肢を提供するため、FP8の計算的利点を活用する。
• 大規模モデルにおいて静的スケーリング戦略を上回る、信頼性の高いFP8量子化手法を確立する。

提案手法

• 各テンソル（重み、活性化、勾配）の最大絶対値に基づき、動的スケーリングバイアスを計算するFP8-AMAXという動的スケーリング戦略を導入する。
• compute_bias関数による動的指数バイアス調整を用い、フォーマットの最大値とテンソルの振幅比のlog2に基づきスケーリングを計算する。この際、設定可能なマージンを適用する。
• FP8操作の前後にスケーリングを適用する：入力をFP8にキャストする前にスケーリングし、行列積の後にアンスケーリングすることで精度を保持する。
• 重みと活性化にはFP8_E4、勾配にはFP8_E5という別々のFP8フォーマットを採用し、範囲と精度のバランスを図る。
• トレーニング（勾配スケーリングおよび重み更新を含む）および推論（トレーニング後量子化）の両パイプラインにこの手法を統合する。
• 収束の堅牢性を検証するため、損失スケーリングスイープを実施し、固定スケーリング（FP8-CSCALE）と比較して、適応的スケーリング（FP8-AMAX）の性能を評価する。

(a) Forward pass for FP16 inference.

実験結果

リサーチクエスチョン

• RQ1 大規模言語モデルのトレーニングおよび推論にFP8を効果的に使用できるか、精度の著しい低下が生じないか？
• RQ2 固定スケーリング（FP8-CSCALE）と比較して、動的テンソルスケーリング（FP8-AMAX）は、モデルサイズにかかわらず収束性および安定性に優れているか？
• RQ3 重み、活性化、勾配においてアンダーフローおよびオーバーフローを防ぐために、どのスケーリング戦略が数値的安定性を確保するか？
• RQ4 トレーニング中にスケーリングバイアスはどのように変化するか？これには、更新頻度および計算コストにどのような影響があるか？
• RQ5 FP8量子化は、多様なLLMアーキテクチャおよびタスクにおいて、FP16と同等の性能を達成できるか？

主な発見

• FP8-AMAXは、MNLI、QQP、SST-2タスクにおいて、1億1100万～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.