ビデオメモリを節約する新しい方法: 2 ビット活性化圧縮を使用して PyTorch でニューラル ネットワークをトレーニングする

[[410937]]

この記事では、カリフォルニア大学バークレー校の PyTorch をベースにしたアクティベーション圧縮トレーニング フレームワークである ActNN を紹介します。同じメモリ制約の下で、ActNN は 2 ビットのアクティベーション圧縮を使用することで、バッチ サイズを 6 ～ 14 倍、モデル サイズまたは入力画像を 6 ～ 10 倍に増やすことができます。 ActNN関連の論文はICML 2021のロングトークとして採択されており、コードはgithubでオープンソース化されています。

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

コード https://github.com/ucbrise/actnn

AIトレーニングが「記憶の壁」にぶつかる

AlexNet、ResNet、GPT-3 に至るまで、ディープラーニングのパフォーマンスの飛躍的進歩は、モデル サイズの驚異的な増加と切り離せません。業界では、大型モデルの方がパフォーマンスが優れているというコンセンサスが生まれています。過去数年間で、最先端のモデルをトレーニングするために必要な計算能力が飛躍的に増加しただけでなく、最先端のモデルをトレーニングするために必要なメモリも飛躍的に増加しました。下の図に示すように、大規模な Transformer モデルのパラメータ数は 2 年ごとに 240 倍の割合で指数関​​数的に増加しています。しかし、単一の GPU のメモリはゆっくりとしか増加しておらず、2 年ごとに 2 倍になっています。さらに、モデルをトレーニングする際には、モデルパラメータだけでなく、中間結果のアクティベーション値やオプティマイザーのステータスも保存する必要があり、より多くのメモリが必要になります。限られた GPU メモリでこれらの大規模モデルをどのようにトレーニングするかが課題となります。

出典: Gholami A、Yao Z、Kim S、Mahoney MW、Keutzer K. AI とメモリ ウォール。RiseLab Medium ブログ投稿、カリフォルニア大学バークレー校

トレーニングメモリを節約する方法

現在、トレーニング メモリを節約する方法は主に 3 つあります。1. 再計算 (勾配チェックポイント/再マテリアライゼーション)、2. スワップに CPU メモリを使用する、3. 分散トレーニングを使用して複数の GPU にテンソルを保存する、です。これら 3 つの方法は互いに競合せず、組み合わせて使用​​できます。ほとんどの機械学習フレームワークはこれらの手法を何らかの形でサポートしており、関連する論文も数多くあります。しかし、これらの戦略を効率的かつ自動的に実装することは容易ではありません。既存の方法とは異なり、新しい圧縮ベースのメモリ節約フレームワークである ActNN を提案します。理論的な証明を提供すると同時に、PyTorch に基づく効率的で使いやすい実装も提供します。表 1 は、ActNN といくつかの既存のメモリ節約システムを比較したものです。 ActNN は PyTorch の動的グラフ実行モードをサポートしており、事前に複雑な戦略検索を必要としません。 ActNN は、PyTorch を変更または再コンパイルせずにインポートできる独立した Python ライブラリです。既存の作品と比較すると、ActNN は柔軟性があり、使いやすいです。同時に、ActNN は理論的には既存のテクノロジーに重ね合わせることもできます。

ActNN: 2ビット活性化圧縮トレーニング

多層ニューラル ネットワークをトレーニングする場合、順方向伝播では、逆方向伝播の勾配を計算するために各層の中間結果を保存する必要があります。これらの中間結果は「アクティベーション」とも呼ばれ、特にバッチ サイズが大きい場合や入力画像が大きい場合に、メモリ消費の大部分を占めます。 ActNN の原理は、これらのアクティベーション値を圧縮してメモリを節約することです。下の図に示すように、左の図は通常の順方向伝播と逆方向伝播を示しています。順方向伝播中は、すべての層の fp32 活性化値が逆方向伝播用に保存されます。損失を計算するときにメモリ使用量がピークに達します。右の図は、ActNN のトレーニング方法を示しています。順方向伝播中は、圧縮操作 Q によってアクティベーション値が圧縮され、保存されます。逆方向伝播中は、解凍操作 Q^-1 によってアクティベーション値が解凍され、勾配が計算されます。

目的がメモリの節約だけであれば、さまざまな圧縮アルゴリズムを使用できますが、既存の圧縮アルゴリズムのほとんどは GPU 上で効率的に実行できず、大きなオーバーヘッドが発生します。 ActNN はここで圧縮アルゴリズムとして 2 ビット量子化を使用することを選択しました。量子化は安価な操作であり、非可逆圧縮を使用して大きな圧縮率を実現できる優れた数学的特性を備えています。

fp32 浮動小数点数を 2 ビット整数に量子化すると非可逆圧縮となり、いくつかのエラーが発生します。この論文では、量子化によって導入された誤差がトレーニングの収束にどのように影響するかを理論的に分析します。

まず、ランダム化された量子化戦略があり、非可逆量子化圧縮後に、推定された非可逆勾配は元の勾配の不偏推定値になります。

この条件下では、既存の確率的勾配降下法の収束定理を適用し、最終収束時の誤差は勾配の分散によって制限されると結論付けます。

次に、量子化圧縮を使用した後に確率的勾配降下法によって計算された勾配の分散を導出します。

等号の右側の最初の項は、ミニバッチをサンプリングするときに確率的勾配降下法によって生成される分散であり、等号の右側の 2 番目の項は、非可逆圧縮によって導入される追加の分散です。この式は非可逆圧縮の影響を明確に示しています。非可逆量子化圧縮によって導入される分散が元の確率的勾配降下法の分散よりもはるかに小さい場合、ActNN によって導入される非可逆圧縮はトレーニングの収束に影響を与えないことに注意してください。式の導出と視覚化の詳細については、記事の最後にある論文リンクを参照してください。この論文では、さまざまな演算子 (conv2d、バッチ ノルム、線形など) の詳細な分析を提供します。

上記の式に着想を得て、非可逆圧縮によって生じる追加の分散を減らすためのいくつかの新しい量子化手法を提案します。異なるサンプル、異なる緯度、異なるレイヤー間の勾配の異質な特性を活用するために、新しい量子化手法 (グループごとの量子化、細粒度混合精度、実行時適応) を導入します。最終的な圧縮アルゴリズムでは、より重要なアクティベーション値にさらに多くのビットが割り当てられます。平均して、各浮動小数点数には 2 ビットが割り当てられます。

圧縮アルゴリズムを実装する際に調整できるパラメータは多数あります。ここでは、メモリの節約とトレーニング速度の間にトレードオフがあります。一般的に、より複雑な圧縮アルゴリズムを使用すると、より多くのメモリを節約できますが、追加のオーバーヘッドも発生し、トレーニングが遅くなります。ユーザーにさらなる柔軟性を提供するために、ActNN ではユーザーが選択できる 5 つの最適化レベル L1 ～ L5 を提供しています。最適化レベルが低いほど、節約できるメモリは少なくなりますが、実行速度は速くなります。最適化レベルを高くすると、より多くのメモリが節約されますが、実行速度も遅くなります。最高の最適化レベル L5 では、ActNN は単純なメモリ スワップ戦略を組み合わせて、圧縮されたアクティベーション値を CPU メモリに移動し、さらにメモリを節約します。

成し遂げる

PyTorch で ActNN アルゴリズムを実装するのは非常に簡単です。 PyTorch nn モジュールの場合、順方向関数に量子化圧縮を追加し、逆方向関数に解凍操作を追加するだけで済みます。すべての計算は以前と同様に fp32 で実行され、疑似コードを下の図に示します。

ActNN は、最も一般的に使用される PyTorch nn.Module の量子化および圧縮バージョンを実装します。ユーザーは、モデル内のすべての PyTorch nn.Module を ActNN の対応するモジュールに置き換えるだけで (nn.Conv2d を actnn.Conv2d に置き換えるなど)、他のコードを変更せずにメモリを節約できます。 ActNN は、コード行を自動的に置き換えるラッパーも提供します。

実験結果

ActNN は非可逆圧縮を実行するため、最も重要なことは、まず ActNN がモデルの精度に影響を与えるかどうかを確認することです。下の図は、ActNN を使用して ImageNet 上で ResNet-50 をトレーニングした結果を示しています。 FP はプレーン fp32 トレーニングの略で、BLPA は NeurIPS 2019 の関連研究です。 ActNN の 2 ビット圧縮モードでは、モデルの精度がほとんど失われないことがわかります。 1.25 ビットのより極端なケースでは、ActNN も収束できますが、精度が多少低下します。しかし、BLPA に関するこれまでの研究では、サイズが 4 ビット未満の場合、収束に失敗します。

また、画像セグメンテーション、オブジェクト検出、自己教師学習など、複数のタスクに関する実験も実施しました。 ActNN は、2 ビット圧縮モードで通常の fp32 とほぼ同じ結果を達成できます。一部のタスクでは、ActNN はより大きなバッチ サイズを使用できるため、より良いテスト結果を達成できます。詳細な実験結果とトレーニング記録については、記事末尾の論文と github リンクを参照してください。

その後、ActNNと通常のfp32トレーニングの実際のメモリ使用量を比較しました。下の表に示すように、ActNN はアクティベーション値で使用されるメモリを 12 倍、トレーニングに使用される合計メモリを 4 ～ 7 倍圧縮できます。この実際のメモリ圧縮効果は理論的な導出と一致しています。アクティベーション値の圧縮率が 32 ビット / 2 ビット = 16 ではなく 12 なのはなぜですか?これは主に、ActNN がインプレース ReLU を使用できず、解凍のために少量の余分な min と scale を保存する必要があるためです。

最後に、ActNNのトレーニング速度をテストしました。 ActNN はトレーニング中に圧縮を実行するため、これらの圧縮によってメモリが節約されますが、追加の計算オーバーヘッドも発生します。一般的に言えば、節約するメモリが増えるほど、入力するオーバーヘッドが増え、トレーニングの速度が遅くなります。 NVIDIA T4 (16 GB メモリ) 上の既存のメモリ効率の高いシステムと ActNN のトレーニング速度を比較しました。下の図に示すように、DTR (ICLR 2020)、BLPA (NeurIPS 2019)、swap は、それぞれ再計算、圧縮、メモリスワッピングに基づく 3 つの方法です。赤い十字はメモリ不足を表します。 Y 軸はトレーニング スループット (1 秒あたりの画像数) であり、値が高いほど優れています。緑の曲線は、さまざまな最適化レベルでの ActNN の最適な結果です。 ActNN は、最大のバッチ サイズに到達できる (つまり、最も多くのメモリを節約できる) だけでなく、すべてのバッチ サイズでベースラインよりもトレーニング速度が速いことがわかります。

さらに多くのネットワークもテストしました。同じメモリ制約の下で、ActNN はバッチ サイズを 6 ～ 14 倍、モデル サイズまたは入力画像を 6 ～ 10 倍に増やすことができます。詳しい実験設定と結果については、記事末尾の論文リンクをご覧ください。

PyTorchで使用する2行のコード

インポート行為nn
モデル = actnn.QModule(モデル)

ActNN は自動モデル変換パッケージを提供します。トレーニング スクリプトに 2 行のコードを挿入するだけで、通常の PyTorch モデルを ActNN を使用するモデルに変換できます。同時に、ActNN はカスタマイズされた使用シナリオをサポートするために、より高度な API も提供します。

その他の例については、github リンクを参照してください。画像認識、画像セグメンテーション、オブジェクト検出、自己教師学習などの複数のタスクで actnn を使用する完全な例とトレーニング レコードを提供します。ぜひお試しください。