LLaMA の微調整によりビデオメモリの要件が半分に削減され、清華大学は 4 ビットの最適化を提案

大規模モデルのトレーニングと微調整にはビデオ メモリに対する要件が高く、オプティマイザーの状態は主要なビデオ メモリ オーバーヘッドの 1 つです。最近、清華大学の Zhu Jun 氏と Chen Jianfei 氏のチームは、ニューラル ネットワーク トレーニング用の 4 ビット オプティマイザーを提案しました。このオプティマイザーは、モデル トレーニングのメモリ オーバーヘッドを節約しながら、フル精度オプティマイザーに匹敵する精度を実現します。

4 ビット オプティマイザーは、さまざまな事前トレーニングおよび微調整タスクで実験されており、精度を維持しながら LLaMA-7B の微調整のメモリ オーバーヘッドを最大 57% 削減できます。

論文: https://arxiv.org/abs/2309.01507

コード: https://github.com/thu-ml/low-bit-optimizers

モデルトレーニングのメモリボトルネック

GPT-3、Gopher、LLaMA に至るまで、業界では大規模なモデルの方がパフォーマンスが優れているというコンセンサスが形成されています。しかし、対照的に、単一の GPU のメモリ サイズはゆっくりと増加しており、大規模なモデルのトレーニングではメモリが主なボトルネックになっています。限られた GPU メモリで大規模なモデルをどのようにトレーニングするかが重要な問題になっています。

これを行うには、まずビデオ メモリ消費の原因を特定する必要があります。実際には、ソースには 3 つの種類があります。

1. 「データメモリ」には、ニューラルネットワークの各層から出力される入力データと活性化値が含まれます。そのサイズは、バッチサイズと画像解像度/コンテキストの長さに直接影響されます。

2. 「モデル メモリ」には、モデル パラメータ、勾配、オプティマイザの状態が含まれます。サイズはモデル パラメータの数に比例します。

3. 「一時ビデオ メモリ」には、GPU カーネルの計算で使用される一時メモリやその他のキャッシュが含まれます。モデルのサイズが大きくなるにつれて、モデルメモリの割合が徐々に増加し、主なボトルネックになります。

オプティマイザー状態のサイズは、使用されるオプティマイザーによって異なります。現在、Transformer のトレーニングでは AdamW オプティマイザーがよく使用されますが、トレーニング中に 2 つのオプティマイザー状態、つまり第 1 モーメントと第 2 モーメントを保存および更新する必要があります。モデルパラメータの数が N の場合、AdamW のオプティマイザー状態の数は 2N となり、明らかにメモリのオーバーヘッドが大きくなります。

LLaMA-7B を例にとると、モデルには約 7B 個のパラメータが含まれています。フル精度 (32 ビット) AdamW オプティマイザーを使用して微調整すると、オプティマイザー状態が占有するビデオ メモリのサイズは約 52.2 GB になります。さらに、単純な SGD オプティマイザーは追加の状態を必要とせず、オプティマイザー状態によって占有されるメモリを節約しますが、モデルのパフォーマンスを保証することは困難です。したがって、この論文では、オプティマイザのパフォーマンスが損なわれないようにしながら、モデル メモリ内のオプティマイザの状態を減らす方法に焦点を当てています。

オプティマイザのメモリを節約する方法

現在、トレーニング アルゴリズムでオプティマイザーのメモリ オーバーヘッドを節約するための主な方法は 3 つあります。

1. 低ランク因数分解の考え方に基づいて、オプティマイザーの状態に対して低ランク近似を実行します。

2. LoRA などの少数のパラメータのみをトレーニングすることで、オプティマイザーの状態の大部分を保存しないようにします。

3. 圧縮ベースの方法では、低精度の数値形式を使用してオプティマイザーの状態を表します。

特に、Dettmers ら (ICLR 2022) は、モメンタムと AdamW を使用した SGD に対応する 8 ビット オプティマイザーを提案しました。ブロック単位の量子化と動的指数数値形式技術を使用することで、言語モデリング、画像分類、自己教師学習、機械翻訳などのタスクで、元の完全精度オプティマイザーに匹敵する結果を達成しました。

これに基づいて、本論文では、最適化状態の数値精度をさらに 4 ビットに削減し、異なる最適化状態に対する量子化方法を提案し、最終的に 4 ビットの AdamW 最適化装置を提案します。同時に、この論文では、圧縮と低ランク分解法を組み合わせる可能性を探り、4 ビットの Factor オプティマイザーを提案しています。このハイブリッド オプティマイザーは、優れたパフォーマンスと優れたメモリ効率の両方を備えています。この論文では、自然言語理解、画像分類、機械翻訳、大規模モデルの命令の微調整など、いくつかの古典的なタスクで 4 ビット オプティマイザーを評価します。

すべてのタスクにおいて、4 ビット オプティマイザーは、より少ないメモリを使用しながら、フル精度オプティマイザーと同等の結果を達成します。

問題設定

圧縮に基づくメモリ効率の高い最適化のためのフレームワーク

まず、アルゴリズム 1 に示されている、一般的に使用されるオプティマイザーに圧縮操作を導入する方法を理解する必要があります。ここで、A は勾配ベースのオプティマイザー (SGD や AdamW など) です。オプティマイザーは、既存のパラメータ w、勾配 g、およびオプティマイザー状態 s を入力し、新しいパラメータとオプティマイザー状態を出力します。アルゴリズム 1 では、完全精度の s_t は一時的に存在しますが、低精度 (s_t) ̅ は GPU メモリに永続的に保存されます。この方法がビデオ メモリを節約できる重要な理由は、ニューラル ネットワークのパラメーターが、各レイヤーのパラメーター ベクトルの連結で構成されていることが多いためです。したがって、オプティマイザの更新もレイヤーごと/テンソルごとに実行されます。したがって、アルゴリズム 1 では、最大で 1 つのパラメータのオプティマイザ状態が完全精度形式でメモリに残され、他のレイヤーに対応するオプティマイザ状態は圧縮された状態になります。

主な圧縮方法：量子化

量子化は、低精度の値を使用して高精度のデータを表す手法です。この論文では、量子化操作を正規化とマッピングの 2 つの部分に分離し、より軽量な方法で新しい量子化方法を設計および実験できるようにします。正規化とマッピングの 2 つの操作は、要素ごとに完全精度データに順番に適用されます。正規化は、テンソル内の各要素を単位区間に投影する役割を担います。テンソルごとの正規化とブロックごとの正規化は次のように定義されます。

正規化方法によって、粒度、外れ値を処理する能力、追加のメモリ オーバーヘッドが異なります。マッピング操作は、正規化された値を低精度で表現できる整数にマッピングする役割を担います。正式に言えば、ビット幅b（つまり、量子化後に各値はbビットで表される）と定義済み関数Tが与えられると、

マップ操作は次のように定義されます。

したがって、適切な T をどのように設計するかが量子化誤差を減らす上で重要な役割を果たします。この論文では、主に線形マッピングと動的指数マッピングについて検討します。最後に、逆量子化のプロセスでは、マッピングと正規化の逆演算子を順番に適用します。

一次モーメント圧縮法

以下では、主に AdamW のオプティマイザー状態 (一次モーメントと二次モーメント) のさまざまな量子化方法を提案します。一次モーメントの場合、私たちの量子化方法は主に Dettmers らの方法 (ICLR 2022) に基づいており、ブロック正規化 (ブロック サイズは 2048) と動的指数マッピングを使用します。

予備実験では、ビット幅を 8 ビットから 4 ビットに直接削減し、一次モーメントは量子化に対して非常に堅牢であり、多くのタスクで一致する結果が得られた一方で、一部のタスクではパフォーマンスの低下も見られることを発見しました。パフォーマンスをさらに向上させるために、一次モーメントのパターンを注意深く調査したところ、単一のテンソルに多くの外れ値があることが分かりました。

これまでの研究では、パラメータと活性化値の外れ値のパターンが研究されてきました。パラメータの分布は比較的滑らかですが、活性化値はチャネル分布の特徴を持っています。オプティマイザー状態における外れ値の分布は複雑であり、一部のテンソルでは固定行に外れ値があり、他のテンソルでは固定列に外れ値があることがわかります。

列に沿って外れ値が分布しているテンソルの場合、行優先のブロック正規化が困難になる可能性があります。したがって、本論文では、ブロック サイズが 128 のより小さなブロックを使用することを提案します。これにより、追加のメモリ オーバーヘッドを制御可能な範囲内に維持しながら、量子化誤差を削減できます。下の図は、さまざまなブロック サイズの量子化誤差を示しています。

2次モーメント圧縮法

一次モーメントと比較すると、二次モーメントの量子化はより困難であり、トレーニングの不安定性につながります。この論文では、ゼロ点問題が二次モーメントを量子化する上での主なボトルネックであると特定し、異常な外れ値分布の正規化方法の改善であるランク 1 正規化を提案しています。この論文では、2次モーメントの因数分解も試みています。

ゼロ点問題

パラメータ、活性化値、勾配の量子化では、ゼロ点が不可欠であることが多く、量子化後に最も頻度の高い点でもあります。しかし、アダムの反復式では、更新のサイズは2次モーメントの-1/2乗に比例するため、ゼロに近い範囲での変化は更新のサイズに大きな影響を与え、不安定性を引き起こします。

下の図は、量子化前後のアダム2次モーメントの-1/2乗の分布をヒストグラムの形式で示しています。つまり、h (v) = 1/(√v + 10^(-6) )です。ゼロが含まれる場合（図b）、ほとんどの値は10^6に押し上げられ、近似誤差が大きくなります。簡単な解決策は、動的指数マップ内のゼロを削除することです。これを実行すると (図 c)、2 次モーメントの近似がより正確になります。実際の場面では、低精度の数値の表現力を効果的に活用するために、ゼロ点を除去した線形マッピングを使用することを提案しており、実験では良好な結果が得られています。

ランク1正規化

本論文では、1 次モーメントと 2 次モーメントの複雑な外れ値分布に基づき、SM3 オプティマイザーにヒントを得て、ランク 1 正規化と呼ばれる新しい正規化手法を提案します。非負行列テンソルx∈R^(n×m)の場合、その1次元統計は次のように定義されます。

ランク 1 正規化は次のように定義できます。

ランク 1 正規化は、テンソルの 1 次元情報をよりきめ細かく活用し、行または列に分散された外れ値をよりインテリジェントかつ効率的に処理できます。さらに、ランク 1 正規化は高次元テンソルに簡単に拡張でき、テンソルのサイズが大きくなるにつれて、発生する追加のメモリ オーバーヘッドはブロック正規化よりも小さくなります。

さらに、本論文では、Adafactor オプティマイザーの 2 次モーメントの低ランク分解法がゼロ点問題を効果的に回避できることがわかったため、低ランク分解法と量子化法の組み合わせも検討されました。下の図は、2 次モーメントに関する一連のアブレーション実験を示しており、ゼロ点問題が 2 次モーメントの量子化のボトルネックであることを確認し、ランク 1 正規化と低ランク分解法の有効性も検証しています。

実験結果

観察された現象と使用方法に基づいて、この研究では最終的に 4 ビット AdamW と 4 ビット Factor という 2 つの低精度最適化装置を提案し、それらを 8 ビット AdamW、Adafactor、SM3 などの他の最適化装置と比較しました。この研究では、自然言語理解、画像分類、機械翻訳、大規模モデルの指示の微調整など、幅広いタスクを評価することを選択しました。次の表は、さまざまなタスクにおける各オプティマイザーのパフォーマンスを示しています。

NLU、QA、NLG を含むすべての微調整タスクにおいて、4 ビット オプティマイザーは 32 ビット AdamW に匹敵するか、それを上回ることがわかります。同時に、すべての事前トレーニング タスクにおいて、CLS、MT、および 4 ビット オプティマイザーは、完全な精度に匹敵するレベルに達します。指示を微調整するタスクから、4 ビット AdamW は事前トレーニング済みモデルの機能を損なうことなく、指示に従う能力をより適切に獲得できるようにしていることがわかります。

その後、4 ビット オプティマイザーのメモリと計算効率をテストしました。結果を次の表に示します。 8 ビット オプティマイザと比較して、本論文で提案する 4 ビット オプティマイザはより多くのメモリを節約でき、LLaMA-7B の微調整の実験では 57.7% という最高の節約率を達成しました。さらに、計算効率に影響を与えずにメモリを節約できる AdamW の 4 ビット融合演算子バージョンも提供しています。 LLaMA-7B の命令微調整タスクでは、4 ビット AdamW によりキャッシュ圧力が軽減されるため、トレーニングの高速化効果も得られます。詳細な実験設定と結果については、論文のリンクをご覧ください。

PyTorchで使用するために1行のコードを置き換えます

import lpmm optimizer = lpmm.optim.AdamW (model.parameters (), lr=1e-3, betas=(0.9, 0.999))

すぐに使用できる 4 ビット オプティマイザーを提供しています。元のオプティマイザーを 4 ビット オプティマイザーに置き換えるだけです。現在、Adam と SGD の低精度バージョンをサポートしています。同時に、カスタマイズされた使用シナリオをサポートするために量子化パラメータを変更するためのインターフェースも提供します。