Metaは、パラメータを積み重ねたり、トレーニング時間に依存したりすることなくViTトレーニングプロセスを加速し、スループットを4倍に増加させます。

現在、ビジュアルトランスフォーマー（ViT）モデルは、画像分類、オブジェクト検出、セグメンテーションなどのさまざまなコンピュータービジョンタスクで広く使用されており、視覚表現と認識においてSOTAの結果を達成できます。コンピューター ビジョン モデルのパフォーマンスは、パラメーターの数やトレーニング時間と正の相関関係にあることが多いため、AI コミュニティではますます大規模な ViT モデルの実験が行われています。

ただし、モデルのサイズがテラフロップスを超え始めると、この分野ではいくつかの大きなボトルネックが発生していることに注意する必要があります。 1 つのモデルをトレーニングするには数か月かかり、数千の GPU が必要になるため、アクセラレータの要件が増加し、大規模な ViT モデルでは多くの専門家が「排除」されることになります。

ViT モデルの使用範囲を拡大するために、Meta AI の研究者はより効率的なトレーニング方法を開発しました。アクセラレータを最大限に活用するには、トレーニングを最適化することが非常に重要です。ただし、このプロセスには時間がかかり、労力がかかり、かなりの専門知識が必要です。秩序だった実験を設定するには、研究者は無数の最適化スキームから選択する必要があります。トレーニング プロセス中に実行される何百万もの計算は、非効率性の影響を受け、妨げられる可能性があります。

Meta AI は、画像分類コード ライブラリ PyCls の ViT 実装に一連の最適化を適用することで、計算効率とストレージ効率を向上できることを発見しました。 PyCI を使用してトレーニングされた ViT モデルの場合、Meta AI のアプローチにより、トレーニング速度とアクセラレータあたりのスループット (TFLOPS) を向上させることができます。

下の図は、最適化されたコード ベース PyCI を使用した V100 ベースラインと比較した、チップあたりのアクセラレータ スループットの相対的な増加を示しています。A100 最適化アクセラレータ スループットは、V100 ベースラインの 4.05 倍です。

仕組み

Meta AI はまず PyCIs コードベースを分析して、トレーニング効率が低い潜在的な原因を特定し、最終的には数値形式の選択に焦点を当てました。デフォルトでは、ほとんどのアプリケーションは、ニューラル ネットワークの値を表すために 32 ビットの単精度浮動小数点形式を使用します。 16 ビット半精度形式 (FP16) に変換すると、モデルのメモリ使用量と実行時間を削減できますが、精度も低下することがよくあります。

研究者たちは妥協案として混合精度を採用した。これにより、システムは単精度形式で計算を実行してトレーニングを高速化し、メモリ使用量を削減すると同時に、結果を単精度で保存して精度を維持します。ネットワークの一部を手動で半精度に変換する代わりに、数値形式を自動的に切り替える自動混合精度トレーニングのさまざまなモードを試しました。自動混合精度のより高度なモードは、主に半精度演算とモデルの重みに依存します。研究者が採用したバランスの取れた設定により、精度を犠牲にすることなくトレーニングを大幅に高速化できます。

プロセスをより効率的にするために、研究者らは、GPU 上でパラメータ、勾配、およびオプティマイザーの状態をシャーディングする FairScale ライブラリの Fully Sharder Data Parallel (FSDP) トレーニング アルゴリズムを最大限に活用しました。 FSDP アルゴリズムを使用すると、研究者はより少ない GPU を使用してより大きなモデルを構築できます。さらに、研究者らは、MTA オプティマイザー、プールされた ViT 分類器、およびバッチ秒入力テンソル レイアウトを使用して、冗長な転置操作をスキップしました。

下の図の X 軸は可能な最適化を示し、Y 軸は分散データ並列 (DDP) ベンチマークと比較した ViT-H/16 でトレーニングした場合のアクセラレータ スループットの相対的な増加を示しています。

研究者らは、合計パッチサイズ 560 で、アクセラレータ チップあたりの 1 秒あたりの浮動小数点演算で測定されたアクセラレータ スループットの 1.51 倍の向上を達成しました。画像サイズを 224 ピクセルから 256 ピクセルに増やすことで、スループットを 1.86 倍に向上させることができました。ただし、画像サイズを変更するとハイパーパラメータも変更されるため、モデルの精度に影響します。フル FP16 モードでトレーニングすると、相対スループットは 2.18 倍に増加します。精度が低下することもありますが、実験では精度の低下は 10% 未満です。

下の図の Y 軸はエポック時間で、ImageNet-1K データセット全体に対する 1 回のトレーニングの期間です。ここでは、通常 224 ピクセルの画像サイズを使用する既存の構成の実際のトレーニング時間に焦点を当てます。

Meta AI の研究者は最適化スキームを使用して、エポック時間 (ImageNet-1K データセット全体のトレーニング実行時間) を 0.65 時間から 0.43 時間に短縮しました。

下のグラフの X 軸は、特定の構成における A100 GPU アクセラレータ チップの数を表し、Y 軸はチップあたりの TFLOPS 単位の絶対スループットを表します。

この調査では、さまざまな GPU 構成の影響についても説明しています。いずれの場合も、システムは分散データ並列 (DDP) ベースラインよりも高いスループットを実現します。チップの数が増えると、デバイス間通信のオーバーヘッドによりスループットがわずかに低下することが分かります。ただし、64 個の GPU を搭載した場合でも、Meta のシステムは DDP ベースラインよりも 1.83 倍高速です。

新しい研究の重要性

ViT トレーニングで達成可能なスループットを 2 倍にすると、トレーニング クラスターのサイズを実質的に 2 倍にすることができ、アクセラレータの使用率を向上させることで AI モデルの炭素排出量を直接削減できます。最近のメガモデルの開発により、モデルの大型化とトレーニング時間の延長の傾向が見られるようになったため、この最適化により、研究コミュニティは最先端の技術をさらに進歩させ、ターンアラウンドタイムを短縮し、生産性を向上させることができると期待されています。