GPTのようなモデルのトレーニング速度が26.5％向上、清華大学の朱俊らはINT4アルゴリズムを使用してニューラルネットワークのトレーニングを加速

アクティベーション、重み、勾配を 4 ビットに量子化することは、ニューラル ネットワークのトレーニングを高速化するために非常に価値があることがわかっています。しかし、既存の 4 ビット トレーニング方法では、最新のハードウェアではサポートされていないカスタム デジタル形式が必要です。この論文では、清華大学の Zhu Jun らが、INT4 アルゴリズムを使用してすべての行列乗算を実装する Transformer トレーニング方法を提案しました。

モデルのトレーニング速度は、活性化値、重み、勾配などの要素の要件と密接に関係しています。

ニューラル ネットワークのトレーニングには一定の計算能力が必要であり、低精度アルゴリズム (完全量子化トレーニングまたは FQT トレーニング) を使用すると、計算とメモリの効率が向上することが期待されます。 FQT は、元の完全精度計算グラフに量子化器と逆量子化器を追加し、高価な浮動小数点演算を安価な低精度浮動小数点演算に置き換えます。

FQT の研究は、収束速度と精度を犠牲にしながら、トレーニングの数値精度を下げることを目的としています。必要な数値精度は FP16 から FP8、INT32+INT8、INT8+INT5 に削減されます。 FP8 トレーニングは、Transformer エンジンを搭載した Nvidia H100 GPU で実行され、大規模な Transformer トレーニングの驚異的な高速化を実現します。

最近、トレーニング数値の精度は 4 ビットにまで下げられました。 Sun らは、INT4 アクティベーション/重みと FP4 勾配を使用していくつかの最新のネットワークをトレーニングすることに成功しました。Chmiel らは、精度をさらに向上させるカスタム 4 ビット対数形式を提案しました。ただし、これらの 4 ビット トレーニング メソッドは、最新のハードウェアではサポートされていないカスタム数値形式を必要とするため、高速化に直接使用することはできません。

4 ビットのような低レベルでのトレーニングは、最適化に大きな課題をもたらします。まず、順方向伝播の微分不可能な量子化器によって損失関数のグラフが不均一になり、勾配ベースのオプティマイザーが簡単に局所最適値に陥ってしまう可能性があります。第二に、勾配は精度が低いと近似値でしか計算できません。この不正確な勾配により、トレーニング プロセスが遅くなり、不安定なトレーニングや発散したトレーニングの原因にもなります。

この論文では、一般的なニューラル ネットワーク Transformer 用の新しい INT4 トレーニング アルゴリズムを提案します。 Transformer のトレーニングに使用される高価な線形演算はすべて、行列乗算 (MM) の形式で記述できます。 MM 形式により、研究者はより柔軟な量子化器を設計できるようになります。この量子化器は、Transformer 内の特定のアクティベーション、重み、および勾配構造を通じて、FP32 行列乗算をより適切に近似します。この論文の量子化器は、ランダム値線形代数の分野における最近の進歩も活用しています。

写真

論文アドレス: https://arxiv.org/pdf/2306.11987.pdf

研究によると、フォワードパスの場合、精度の低下の主な原因はアクティベーションの外れ値にあることがわかっています。この外れ値を抑制するために、この研究では、変換された活性化行列を量子化するために使用されるアダマール量子化器を提案しました。変換はブロック対角アダマール行列であり、外れ値によって運ばれる情報を外れ値の近くの行列項目に拡散し、外れ値の数値範囲を縮小します。

バックプロパゲーションでは、活性化勾配の構造化されたスパース性を活用します。研究によると、一部のトークンの勾配は非常に大きいが、同時に、他のほとんどのトークンの勾配は非常に小さく、より大きな勾配の量子化残差よりもさらに小さいことが示されています。したがって、これらの小さな勾配を計算する代わりに、計算リソースを使用して、より大きな勾配の残差を計算することができます。

この論文では、順方向および逆方向の量子化技術を組み合わせて、Transformer のすべての線形演算に INT4 MM を使用するアルゴリズムを提案します。この研究では、自然言語理解、質問応答、機械翻訳、画像分類など、さまざまなタスクでトランスフォーマーをトレーニングするためのアルゴリズムを評価しました。既存の 4 ビット トレーニング作業と比較して、提案されたアルゴリズムは同等かそれ以上の精度を実現します。さらに、このアルゴリズムはカスタム数値形式 (FP4 や対数形式など) を必要としないため、最新のハードウェア (GPU など) と互換性があります。本研究で提案されたプロトタイプの量子化 + INT4 MM 演算子は、FP16 MM ベースラインよりも 2.2 倍高速で、トレーニング速度が 35.1% 向上します。

前方伝播

トレーニングプロセス中、研究者は INT4 アルゴリズムを使用してすべての線形演算子を高速化し、計算強度が低いすべての非線形演算子を FP16 形式に設定しました。 Transformer のすべての線形演算子は、行列乗算形式で記述できます。デモンストレーションの目的で、彼らは次のような単純な行列乗算の高速化を検討しました。

写真

この行列乗算の主な使用例は、完全接続層です。

学習したステップサイズの量子化

トレーニングを高速化するには、整数演算を使用してフォワードパスを計算する必要があります。そのため、研究者は学習済みステップサイズ量子化器 (LSQ) を利用しました。静的量子化方式である LSQ の量子化スケールは入力に依存しないため、動的量子化方式よりもコストがかかりません。対照的に、動的量子化方法では、各反復で量子化スケールを動的に計算する必要があります。

FP行列Xが与えられると、LSQは次の式(2)を使用してXを整数に量子化する。

写真

外れ値をアクティブにする

4 ビットのアクティベーション/重みを持つ FQT (完全量子化トレーニング) に LSQ を単純に適用すると、アクティベーションの外れ値により精度が低下します。下の図 1 (a) に示すように、他の項目よりも大きさがはるかに大きい外れ​​値項目がいくつかアクティブ化されています。

この場合、ステップ サイズ s_X は、量子化の粒度と表現可能な値の範囲との間のトレードオフになります。 s_X が大きい場合、他のほとんどの項を粗く表現する代わりに、外れ値を適切に表現できます。 s_Xが小さい場合、[−Q_Ns_X、Q_Ps_X]の範囲外の項は切り捨てられる必要があります。

アダマール量子化

研究者らは、外れ値問題を解決するためにアダマール量子化器 (HQ) の使用を提案しました。その主なアイデアは、外れ値の少ない別の線形空間で行列を量子化することです。

活性化マトリックス内の外れ値は、特徴レベルの構造を形成する可能性があります。これらの外れ値は通常、いくつかの次元に集中しています。つまり、X 内のいくつかの列だけが他の列よりも大幅に大きくなります。線形変換として、アダマール変換は外れ値を他の項に分散させることができます。具体的には、アダマール変換 H_k は 2^k × 2^k 行列です。

外れ値を抑制するために、研究者らは X と W の変換バージョンを量子化しました。

量子化された行列を組み合わせることで、研究者は次の結果を得ました。

逆変換は互いに打ち消し合い、MM は次のように実装できます。

写真

バックプロパゲーション

研究者らは、INT4 操作を使用して線形層のバックプロパゲーションを高速化しました。式（３）で定義される線形演算子HQ-MMには、活性化X、重みW、ステップサイズs_Xとs_Wの4つの入力があります。損失関数 L に対する出力 ∇_YL の勾配が与えられた場合、これら 4 つの入力の勾配を計算する必要があります。

勾配の構造的スパース性

研究者たちは、勾配行列∇_Y がトレーニング中に非常にスパースになることが多いことに気づきました。スパース構造は、∇_Y のいくつかの行 (つまり、トークン) に大きなエントリがある一方で、他のほとんどの行はすべてゼロのベクトルに近いというものです。彼らは、すべての行について行ごとのノルム ∇(∇_Y)_i:∇ のヒストグラムを以下の図 2 にプロットしました。

写真

ビット分割と平均分数サンプリング

研究者らは、構造化されたスパース性を活用してバックプロパゲーション中に MM を正確に計算するための勾配量子化器の設計方法について議論しました。大まかな考え方としては、多くの行の勾配は非常に小さいため、パラメータの勾配への影響も小さいですが、多くの計算が無駄になります。さらに、INT4 では大きな行を正確に表現することはできません。

このスパース性を活用するために、研究者は各トークンの勾配を上位 4 ビットと下位 4 ビットに分割するビット分割を提案しました。次に、スコアを平均化することで、最も多くの情報を持つ勾配が選択されます。これは、RandNLA の重要度サンプリング手法です。

実験結果

この研究では、言語モデルの微調整、機械翻訳、画像分類など、さまざまなタスクで INT4 トレーニング アルゴリズムを評価しました。この研究では、CUDA と cutlass2 を使用して、提案された HQ-MM および LSS-MM アルゴリズムを実装しました。 LSQ を埋め込み層として単純に使用するだけでなく、すべての浮動小数点線形演算子を INT4 に置き換え、最後の分類器層の完全な精度を維持します。そして、その過程で、研究者たちは、すべての評価モデルにデフォルトのアーキテクチャ、オプティマイザー、スケジューラー、ハイパーパラメータを使用しました。

モデルの精度が収束します。以下の表 1 は、各タスクにおける収束モデルの精度を示しています。

写真

言語モデルの微調整。 LSQ+LUQ と比較して、提案されたアルゴリズムは、BERT ベース モデルで平均精度を 5.5% 向上させ、BERT ラージ モデルで平均精度を 25% 向上させます。

研究チームはアルゴリズムのデモンストレーションも行い、さらに SQUAD、SQUAD 2.0、Adversarial QA、CoNLL-2003、SWAG データセットで結果をデモンストレーションしました。すべてのタスクにおいて、この方法は LSQ+LUQ と比較して優れたパフォーマンスを実現します。 LSQ+LUQ と比較すると、この方法は SQUAD と SQUAD 2.0 でそれぞれ 1.8% と 3.6% 改善されます。より困難な敵対的 QA では、この方法により F1 スコアが 6.8% 向上します。 SWAG と CoNLL-2003 では、この方法により精度がそれぞれ 6.7% と 4.2% 向上します。

機械翻訳。この研究では、事前トレーニングにも提案された方法を使用しました。この方法は、WMT 14 En-Deデータセット上で機械翻訳用のTransformerベース[51]モデルをトレーニングします。

HQ+LSS の BLEU 分解率は約 1.0% で、これは Ultra-low の 2.1% 未満であり、LUQ 論文で報告されている 0.3% よりも高い値です。それにもかかわらず、HQ + LSS はこの事前トレーニング タスクで既存の方法と同等のパフォーマンスを発揮し、最新のハードウェアをサポートします。

画像分類。事前トレーニング済みの ViT チェックポイントを ImageNet21k にロードし、CIFAR-10、CIFAR-100、および ImageNet1k で微調整します。

LSQ+LUQと比較して、この研究方法により、ViT-B/32とViT-L/32の精度がそれぞれ1.1%と0.2%向上しました。 ImageNet1kでは、この方法により、LSQ+LUQと比較して、ViT-B/32の精度が2%、ViT-L/32の精度が2.6%、ViT-L/32の精度が0.2%向上しました。

研究チームは、ImageNet1K 上の事前トレーニング済みの DeiT-Small モデルでアルゴリズムの有効性をさらにテストしました。その結果、HQ+LSS は、LSQ+LUQ と比較して同等の精度レベルに収束しながらも、よりハードウェアに優しいことがわかりました。

アブレーション研究

私たちは、困難な CoLA データセットに対する前方法と後方法の有効性を独立して実証するためにアブレーション研究を実施しています。順方向伝播におけるさまざまな量子化器の有効性を調べるために、バックプロパゲーションを FP16 に設定しました。結果は下の図3(a)に示されています。

バックプロパゲーションについては、研究者らは単純なミニマックス量子化器、LUQ、および独自の LSS を比較し、フォワードパスを FP16 に設定しました。結果は、下の図 3 (b) に示されています。ビット幅は 2 より大きいですが、LSS によって達成された結果は LUQ の結果と同等か、わずかに上回っています。

写真

計算とメモリの効率

研究者らは、提案した HQ-MM (HQ)、重み勾配を計算するための LSS (LSSWeight)、活性化勾配を計算するための LSS (LSSAct)、平均スループット (INT4)、および cutlass によって提供されたベースライン Tensor Core FP16 GEMM 実装 (FP16) のスループットを、以下の図 4 の NVIDIA RTX 3090 GPU 上で比較しました。ピーク スループットは 142 FP16 TFLOPs および 568 INT4 TFLOPs です。

写真

研究者らはまた、8 つの NVIDIA A100 GPU 上で BERT のような言語モデルと GPT のような言語モデルをトレーニングするための FP16 PyTorch AMP と独自の INT4 トレーニング アルゴリズムのトレーニング スループットを比較しました。彼らは、隠れ層のサイズ、中間完全接続層のサイズ、バッチ サイズを変更し、INT4 トレーニングの高速化を図 5 にプロットしました。

結果は、INT4 トレーニング アルゴリズムが BERT のようなモデルで最大 35.1% の加速を達成し、GPT のようなモデルで最大 26.5% の加速を達成することを示しています。

写真