[翻訳] DeepSpeed: 誰もが使用できる非常に大規模なモデルトレーニングツール

今年 2 月にDeepSpeed をリリースしました。これは、新しいメモリ最適化テクノロジである ZeRO (Zero Redundancy Optimizer) を含むオープンソースのディープラーニングトレーニング最適化ライブラリであり、規模の拡大、速度の向上、コストの制御、可用性の向上によって大規模モデルのトレーニングを大幅に促進します。DeepSpeed は、研究者によるチューリング自然言語生成モデル ( Turing-NLG ) の開発を支援してきました。このモデルは、出版時点では世界最大の言語モデル (170 億のパラメーターを含む) であり、最高の精度を備えていました。5 月に、2,000 億のパラメーターを使用したモデルトレーニングをサポートし、最新テクノロジーよりも最大 10 倍高速なZeRO-2 と、最速の BERT トレーニングを強化する一連のコンピューティング、IO、およびコンバージェンス最適化機能をリリースしました。それ以来、私たちは速いペースで革新を続け、深層学習モデルのトレーニングにおける速度と規模の限界を押し上げてきました。

本日は、ディープラーニングトレーニングを極限まで推し進めるだけでなく、データサイエンティストによるスーパーコンピューターでのトレーニングから、ローエンドクラスターやクラスターでのトレーニングに至るまで、このテクノロジーをより広く使用できるようにする新しい開発を皆さんと共有できることを非常にうれしく思います。 GPU が 1 つだけでも。具体的には、DeepSpeed は、 AI at Scaleイニシアチブをさらに拡大するために 4 つの体系的な新しいテクノロジーを追加しました。また、Microsoft の AI 製品とプラットフォームの革新も促進します。これらのテクノロジーは、コンピューティング、メモリ、通信を非常に効率的に利用できるようになり、数十億から数兆のパラメーターを使用してモデルをトレーニングするのに役立ちます。これらのテクノロジーは非常に長い入力シーケンスもサポートしており、シングルカード GPU、キロカード GPU を備えたハイエンドクラスター、または低速イーサネットを備えたローエンドクラスターで使用できます。

3D 並列化で兆パラメータのモデルトレーニングを実現: DeepSpeed は、ZeRO がサポートするデータ並列処理、パイプライン並列処理、テンソルスライスモデル並列処理の 3 つの並列手法を柔軟に組み合わせて実装します。3D 並列処理は、さまざまなワークロードのニーズに適応して、ほぼ完璧なメモリスケーラビリティとスループットスケーリング効率を達成しながら、何兆ものパラメータを持つ非常に大規模なモデルをサポートします。さらに、通信効率が向上したため、ネットワーク帯域幅が限られている従来のクラスター上で、数十億のパラメーターを使用してモデルを 2 ～ 7 倍の速さでトレーニングできるようになります。
ZeRO-Offload により、1 つの GPU カードで 10 倍大きなモデルをトレーニングできるようになります。CPU メモリ と GPU メモリの両方を利用して大規模なモデルをトレーニングするために、ZeRO-2 を拡張しました。当社のユーザーは、単一の NVIDIA V100 GPUを搭載したマシンを使用して、ビデオメモリを使い果たすことなく最大 130 億のパラメータを持つモデルを実行し、モデルを既存の方法の最大 10 倍のサイズに拡張して、競争力のあるスループットを維持できます。この機能により、数十億のパラメーターを使用したモデルトレーニングが民主化され、多くのディープラーニング実践者がより大規模で優れたモデルを探索できる窓が開かれます。
DeepSpeed スパースアテンションで 10 倍長いシーケンスを 6 倍高速に実行: DeepSpeed は、テキスト入力、画像入力、音声入力などのモデル入力の長いシーケンスをサポートするツールテクノロジであるスパースアテンションカーネルを提供します。従来の高密度 Transformer と比較して、桁違いに長い入力シーケンスをサポートし、同等の精度を維持しながら最大 6 倍高速な実行を実現します。また、最新のスパース実装よりも 1.5 ～ 3 倍高速です。さらに、当社のスパースカーネルはスパース形式を柔軟にサポートし、ユーザーがカスタムスパース構造を革新できるようにします。
1 ビット Adam は通信を 5 分の 1 に削減します。Adam は、大規模な深層学習モデルのトレーニングシナリオで効果的な (おそらく最も広く使用されている) オプティマイザーです。ただし、通信効率最適化アルゴリズムと互換性がないことがよくあります。したがって、デバイス間で分散して拡張する場合、通信オーバーヘッドがボトルネックになる可能性があります。新しい 1 ビット Adam アルゴリズムとその効率的な実装を紹介します。このアルゴリズムは、Adam と同様の収束率を達成しながら、通信量を最大 5 分の 1 に削減します。通信が制限されたシナリオでは、分散トレーニング速度が 3.5 倍向上したことが観察され、これによりアルゴリズムをさまざまなタイプの GPU クラスターやネットワーク環境に拡張できます。

携帯電話のスクリーンショット

このブログ投稿では、これら 4 つのテクノロジーについて詳しく説明します。私たちは、これらのエキサイティングな最適化テクニックをオープンソースプロジェクト DeepSpeedで発表しました。

3D パラレル: 兆パラメータモデルへのスケーリング

最新の GPU クラスターでのコンピューティングの急速な成長に伴い、驚くべき機能を備えた強力な 10 兆個のパラメーターモデルをトレーニングすることはもはや手の届かないことではなく、近い将来に実現可能になる可能性があります。DeepSpeed は、データ並列トレーニング、モデル並列トレーニング、パイプライン並列トレーニングという 3 つの強力な手法を組み合わせて、兆規模のモデルをトレーニングし、数千の GPU に拡張します。これら 3 つの共生により、ディープラーニングトレーニングを、それぞれの戦略を単独で使用した場合の可能性をはるかに超えて拡張することができます。3D 並列処理は、兆パラメータモデルのトレーニングにおける 2 つの基本的な課題、つまりメモリ効率と計算効率を同時に解決します。その結果、DeepSpeed は速度を犠牲にすることなく、ビデオメモリ内の最大のモデルに合わせて拡張できます。

巨大なモデルをトレーニングする際のメモリと計算効率の課題を理解する

ビデオメモリの効率: 1 兆個のパラメータモデルをトレーニングするために必要なビデオメモリは、単一の GPU のビデオメモリサイズをはるかに超えています。混合精度トレーニングに Adam オプティマイザーを使用する場合、モデルの状態 (パラメーター、勾配、オプティマイザーの状態) を保存するために約 16 TB のビデオメモリが必要です。比較のために、最も先進的な Nvidia A100 GPU のビデオメモリは 40 GB しかありません。モデルの状態を保存するだけでも、そのような GPU が 400 個必要になります。

アクティベーション関数によって消費される追加メモリは、バッチサイズに応じて増加します。バッチが 1 に設定されている場合、兆パラメータモデルをトレーニングすると、アクティベーション関数用に 1 TB を超えるメモリ (以下、アクティベーションメモリと呼びます) が生成されます。チェックポイント処理でビデオメモリをアクティブ化し、計算でビデオメモリを交換すると、このビデオメモリを約 20 GB に削減できますが、トレーニングにはまだ多すぎます。

このような大規模なモデルをビデオメモリを使い果たすことなくトレーニングできるようにするには、モデルの状態とアクティブなビデオメモリを複数の GPU デバイスに効率的に分割する必要があります。

計算効率: 1 兆個のパラメータモデルのエンドツーエンドのトレーニングには約 5,000 Zflops (つまり、5 の後に 24 のゼロが続きます。この推定は OpenAI の研究によるスケーリングの法則に基づいています ) が必要であると推定されています。これは、このようなモデルをトレーニングするには、4,000 台の A100 を 50% の計算効率で約 100 日間実行する必要があることを意味します。

大規模なスーパーコンピューティング GPU クラスターには 4,000 個を超える GPU を搭載できますが、バッチサイズの制限により、この規模で高いコンピューティング効率を達成することは依然として課題です。通信時間に対する計算時間の比率が増加すると、計算効率が向上します。この比率はバッチサイズに比例します。ただし、モデルをトレーニングするためのバッチサイズには上限があり、それを超えると収束が大幅に低下する可能性があります。

実際、最大のモデルの 1 つであるGPT-3 のトレーニングバッチサイズは約 1500 です。約 4000 個の GPU が使用されている場合、バッチサイズを 4000 に自由に設定できますが、各カードのバッチサイズは 1 のみとなり、スケーラビリティに影響します。

データ並列処理、モデル並列処理、パイプライン並列処理の間のトレードオフを理解する

データ並列処理は、深層学習で一般的に使用される手法です。この手法では、入力トレーニングデータの各バッチがデータ並列ワーカー間で均等に分割されます。バックプロパゲーションの後、オプティマイザが各ワーカーに対して同じ更新を確実に行うために、勾配を伝達して削減する必要があります。データ並列処理には、高い計算効率や実装の労力の軽減など、いくつかの明確な利点があります。ただし、データ並列処理のバッチサイズはワーカーの数に応じて増加するため、収束に影響を与えずに常にバッチサイズを増やすことはできないことがよくあります。

メモリ効率: データ並列処理では、すべてのワーカー間でモデルとオプティマイザがコピーされるため、メモリ効率は高くありません。DeepSpeed は、データ並列処理のメモリ効率を向上させるように設計された一連のオプティマイザーであるZeROを開発しました。この作業は、冗長性を削減するためにオプティマイザーの状態をワーカー間で分割する ZeRO のフェーズ 1 に依存しています。
計算効率: 並列度を上げても、各ワーカーが実行する計算量は一定になります。データの並列処理により、小規模でのほぼ線形のスケーリングが可能になります。ただし、ワーカー間の勾配を減らすための通信オーバーヘッドはモデルのサイズに正の相関があるため、モデルが大きい場合や通信帯域幅が低い場合、計算効率は制限されます。。勾配累積は、通信コストを償却するための一般的な戦略です。さらにバッチサイズを増やし、ローカルで マイクロバッチを使用して順方向および逆方向の伝播を複数回実行して勾配を蓄積し、その後勾配の低減とオプティマイザーの更新を実行します。

モデルの並列処理は、広範な種類の技術です。モデルのさまざまなレイヤーを複数のワーカーに分割します。その性質上、モデルの並列処理の計算と通信はモデルの構造によって異なるため、実装には多大な労力がかかります。DeepSpeed は、NVIDIA の Megatron-LMを借用して、Transformer ベースの言語モデルに大規模なモデルの並列処理を提供します。モデルの並列処理では、ワーカーの数に比例してメモリ使用量が削減され、3 つの並列処理の中で最もメモリ効率が高くなります。ただし、その代償として計算効率が最も低くなります。

ビデオメモリの効率: モデルの並列処理により、ワーカーの数に比例してビデオメモリの使用量が削減されます。重要なのは、これが単一ネットワーク層のアクティブメモリを削減する唯一の方法であるということです。DeepSpeed は、アクティブなメモリをモデルの並列ワーカー間で分割することにより、メモリ効率をさらに向上させます。
計算効率: 各順方向伝播および逆方向伝播で追加の通信アクティブ化値が必要となるため、モデル並列処理の計算効率は非常に低くなります。モデルの並列処理には高い通信帯域幅が必要であり、通信帯域幅が限られているノードにはうまく拡張できません。さらに、各モデルのパラレルワーカーは、各通信ステージ間で実行される計算量を削減するため、計算効率に影響を与えます。モデルの並列処理は、メモリと計算効率の間のトレードオフを生み出すために、データの並列処理と組み合わせて使用されることがよくあります。

パイプライン並列トレーニングエンジンも、DeepSpeed のこのリリースに含まれています。パイプライン並列処理では、モデルの各層を並列処理できるステージに分割します。ステージがマイクロバッチの順方向パスを完了すると、アクティブ化されたメモリがパイプラインの次のステージに伝達されます。同様に、次のステージがバックプロパゲーションを完了すると、勾配はパイプラインを通じて逆方向に伝達されます。パイプラインの各ステージを並行して計算できるようにするには、複数のマイクロバッチを同時に計算する必要があります。PipeDreamなど、メモリと計算効率、および収束動作をトレードオフするいくつかのメソッドが開発されています。DeepSpeed が採用する手法は、勾配累積によって並列処理を実現し、同じ総バッチサイズの下で従来のデータ並列およびモデル並列トレーニングと同じ収束条件を維持します。

ビデオメモリの効率: パイプラインの並列処理によって削減されるビデオメモリはパイプラインのステージ数に比例するため、モデルのサイズはワーカーの数に比例して拡張できます。ただし、パイプラインの並列処理によって各層のアクティベーション関数のメモリフットプリントは削減されません。さらに、各ワーカーは、同時に実行される各マイクロバッチのアクティベーション値を保存する必要があります。この結果、パイプラインの最初のステージのアクティベーションメモリは、単一の mirco バッチの合計アクティベーションメモリとほぼ同じになります。兆パラメータモデルでは、マイクロバッチ用に約 19 GB のビデオメモリのアクティベーションメモリが必要になります。これは、新しい Nvidia A100 GPU の総ビデオメモリのほぼ半分です。
計算効率：パイプライン並列処理は、通信量が各ステージの境界における各層の活性化値にのみ比例するため、最も通信量が低くなります。ただし、無限に拡張することはできません。モデルの並列処理と同様に、パイプラインサイズを増やすと、各パイプラインステージで必要な計算量が減り、計算対通信の比率が低下します。良好な計算効率を達成するには、パイプライン並列処理で各ステージの計算負荷が完全にバランスが取れていることも必要です。

さらに、パイプラインの並列処理により、各バッチの開始時と終了時にパイプラインを補充または排出する必要があるため、バブルオーバーヘッドが発生します。1 つのパイプラインステージと比較して、4 倍または 8 倍の勾配蓄積ステップ数 (およびバッチサイズ) でトレーニングすると、それぞれ 81% と 90% のスケーラビリティが達成されました。

3D 並列処理により、高いメモリ効率と高い計算効率を同時に実現

データ、モデル、パイプラインの並列処理はすべて、メモリと計算効率の向上において特定の役割を果たします。図 1 は、当社の 3D 戦略を示しています。

ビデオメモリ効率: まず、モデルの各層を異なるパイプラインステージに分割し、モデル全体で各ステージの層を並列にさらに分割します。この 2D の組み合わせにより、モデル、オプティマイザー、およびアクティベーション関数によって消費されるメモリが同時に削減されます。ただし、通信オーバーヘッドを導入せずにモデルを無限に分割することはできないため、計算効率が制限されます。

計算効率: 計算効率を犠牲にすることなく、モデルとパイプラインの並列処理がサポートできる以上にワーカーの数をスケールするために、ZeRO でサポートされるデータ並列処理 (ZeRO-DP) を使用します。ZeRO-DPは、オプティマイザ状態量を分割することでメモリ利用効率をさらに向上させるだけでなく、通信トポロジに基づくマッピング関係を利用することで、最小限の通信オーバーヘッドで任意の数のGPUに拡張することができます。

通信トポロジに基づく 3D マッピング (図 2) : 2 つの主要なアーキテクチャプロパティを活用することで、3D 並列処理の各次元をワーカーに注意深くマッピングし、最大の計算効率を達成します。

ノード内およびノード間の通信帯域幅を最適化する: モデルの並列処理は 3 つの戦略の中で通信オーバーヘッドが最も大きいため、より大きなノード内帯域幅を利用するためにノード内にモデルの並列ワーカーグループを配置することを優先します。ここでは、NVIDIA Megatron-LM に基づいてテンソルセグメンテーションモデルの並列処理を実行します。モデル並列処理グループがノード内のすべてのワーカーを占有していない場合は、データ並列処理グループをノード内に配置することを選択します。それ以外の場合は、ノード間でデータの並列処理を実行します。パイプライン並列処理では通信が最小限に抑えられるため、通信帯域幅に制限されることなく、ノード間でパイプラインステージをスケジュールできます。
並列通信による帯域幅の増加: 各データ並列グループが通信する必要がある勾配の量は、パイプラインとモデルの並列処理の規模に応じて直線的に減少するため、総通信量はデータ並列処理のみを使用する場合よりも少なくなります。さらに、各データ並列グループは一部のワーカー内で独立して通信し、グループ間の通信は相互に並列することができます。その結果、トラフィックが削減され、局所性と並列性が向上することにより、データ並列通信の有効帯域幅が増加します。

この図は、32 個のワーカーによる 3D 並列処理の例を示しています。ニューラルネットワークの層は 4 つのパイプラインステージに分割されます。各パイプラインステージのレイヤーは、4 つのモデル並列ワーカーにさらに分割されます。最後に、パイプラインステージごとに 2 つのデータ並列インスタンスがあり、ZeRO はオプティマイザーの状態を 2 つのレプリカ間で分割します。

図 1: 32 個のワーカーによる 3D 並列処理の例。ニューラルネットワークの各層は 4 つのパイプラインステージに分割されます。各パイプラインステージのレイヤーは、4 つのモデル並列ワーカーにさらに分割されます。最後に、パイプラインステージごとに 2 つのデータ並列インスタンスがあり、ZeRO はオプティマイザーの状態を 2 つのレプリカ間で分割します。

色付きのブロックは、図 1 のワーカーと、ノードあたり 4 つの GPU を備えた 8 ノードのシステム上の GPU へのマッピングを示しています。同じ色の GPU は同じノード上にあります。

図 2: ノードごとに 4 つの GPU を備えた 8 ノードのシステム上の GPU への図 1 のワーカーのマッピング。同じ色の GPU は同じノード上にあります。

兆パラメータモデルの 3D 並列トレーニングの詳細をご覧ください

8 方向のモデル並列処理、64 方向のパイプライン並列処理、および 8 方向のデータ並列処理を使用すると、1 兆パラメータのモデルを 4,096 個の NVIDIA A100 GPU でスケーラブルにトレーニングできます。

モデル並列処理とパイプライン並列処理を組み合わせることで、3D 並列処理により、複数のノード全体で優れたメモリ効率と高い計算効率が実現します。モデル並列処理により、ノード内のアクティベーションメモリとモデル状態量のストレージ効率が向上します。一方、パイプライン並列処理は、モデル並列処理のみを使用する場合と比較して、計算効率を犠牲にすることなくノード間でモデル状態を効率的に格納できます。マイクロバッチサイズ 1 の 1 兆個のパラメーターの例では、アクティベーション値チェックポイントと上記の 3D 並列処理を使用した後、モデル状態量は 30 GB のビデオメモリを消費し、分割されたアクティベーション値は 2.5 GB のメモリを消費します。。これにより、合計メモリフットプリントは 32.5 GB となり、このようなモデルに対応してトレーニングするために 40 GB のメモリを備えた NVIDIA A100 GPU を使用できるようになります。

モデル並列処理とパイプライン並列処理を組み合わせると、パイプライン並列処理により、非常に小さなバッチで最小限のバブルオーバーヘッドで高い計算効率を実現できます。8 ウェイモデルの並列処理では、モデルごとに 1 つのマイクロバッチを使用すると、GPU ごとに有効なマイクロバッチサイズが 1/8 になります。したがって、パイプライン並列処理の勾配累積ステップを 8 倍使用しても、各 GPU で合計累積バッチサイズが 1 になるだけであり、パイプライン並列処理では 90% の計算効率を達成できます。データ並列処理と組み合わせると、4096 個の GPU で合計有効バッチサイズが 4096 になり、それでも 90% のパイプライン効率を達成できます。

しかし、データの並列処理は計算効率にどのような影響を与えるのでしょうか? データの並列処理では、効率を維持するために各 GPU に大きなバッチが必要ではないでしょうか?

モデルの並列処理により、各 GPU の有効バッチサイズが 1 未満に減少する可能性があります。これにより、パイプラインの並列処理により、バッチが小さい場合でもパイプラインのバブルオーバーヘッドを隠すことができます。ノード間でパイプライン並列処理を使用すると、パイプラインの各ステージのデータ並列ノードが相互に独立して通信し、他のパイプラインステージと並行して通信できることに注意してください。実際、ハイエンド GPU クラスターで一般的な完全接続ネットワークトポロジでは、これはデータ並列トレーニングに利用できる有効な通信帯域幅に重要な意味を持ちます。パイプラインステージの各ノードは、対応するデータ並列ノードと並行して通信できるため、有効な通信帯域幅はパイプラインステージの数に比例します。64 個の並列パイプラインステージを設定すると、有効帯域幅は単一ノードとの間の帯域幅の 64 倍になります。パイプライン並列処理により、データ並列処理に非常に大きな有効帯域幅がもたらされ、計算と通信の比率が非常に低い小さなバッチ状況でもデータ並列処理を効率的に拡張できます。

線形スケーラビリティを備えた兆パラメータモデルのトレーニング

DeepSpeed は、わずか 800 個の NVIDIA V100 GPU を使用して、 1兆のパラメーターを使用して言語モデルをトレーニングできます (図 3)。モデルサイズとトレーニングスループットを実証し、メモリと計算効率の両方がモデルサイズに比例して増加することを観察しました。さまざまな構成にわたって、GPU あたり約 14 億個のパラメーターをトレーニングできます。これは、単一の GPU がメモリ不足なくサポートできる最大モデルサイズであり、完全なメモリスケーラビリティを示しています。また、V100 GPU あたり 47 Tflops のスループットで、ほぼ完璧な線形の計算効率スケーリングも達成しました。前述のハードウェアの場合、これは優れたスケーラビリティとスループットです。

図 3: GPU の数の関数としてのモデルサイズ (数十億のパラメーター) とトレーニングスループット (Pflops 単位) のグラフ。DeepSpeed は、32 GB のメモリを備えた 800 個の NVIDIA V100 Tensor コア GPU を使用して、1 兆のパラメーターでモデルをトレーニングできます。各構成では、 NVIDIA Megatron-LM によって提供される 16 ウェイモデル並列処理が使用され、残りの GPU がパイプライン並列処理を担当します。兆パラメータモデルには 298 層の Transformer があり、その隠れ層のサイズは 17408、トレーニングシーケンスの長さは 2048、バッチサイズは 2048 です。小規模なモデルの場合、GPU の数に比例して、Transformer レイヤーの数とバッチサイズを削減しました。

3D 並列処理によって GPT-3 スケールのトレーニングモデルがどのように高速化されるかについて詳しく説明します。

図 4: 800 個の GPU を使用して 2D と 3D を並行して使用し、1,800 億のパラメーターを持つ GPT-3 スケールモデルをトレーニングしたシステムパフォーマンス。このモデルには 100 の Transformer レイヤー、12288 の隠れレイヤーサイズ、および 96 個のアテンションヘッドがあります。トレーニングに使用されるバッチサイズは 2048、シーケンスの長さは 2048 です。ZeRO-1 はデータ並列処理と組み合わせて使用することもできます。P、M、D はそれぞれパイプライン、モデル、データの並列ディメンションを表します。

図 4 では、 3D 並列処理のベースラインとして 1,750 億を超えるパラメータを持つ最新のGPT-3モデルアーキテクチャを使用しています。

最初に2D 構成(C1 ～ C3)を評価しました。構成 C1 と C2 はパイプラインとモデルの並列処理のみを使用します。モデルをトレーニングできますが、モデルの過剰分解によりスループットと GPU 使用率が低下します。C3 はパイプライン処理とデータ並列処理のみを使用しようとしていますが、Megatron のモデル並列処理によってアクティベーションの量を減らすことなしにビデオメモリ不足の問題を解決することはできません。
3D 構成(C4 ～ C10) では、パイプラインの並列処理が順番に増加します。並列処理のバランスをとった中間構成では、最高のパフォーマンスを実現し、グラフィックスメモリ、コンピューティング、および通信の 3 つの最も高い効率を達成できます。
最適な 3D アプローチは、GPU あたり 49 Tflops を達成し、ハードウェアの理論上のピークを 40% 超えます。

ハイブリッド並列処理により、低帯域幅クラスター上で GPT-2 のトレーニングが 7 倍高速化される様子をご覧ください。

私たちは 15 億パラメータの GPT-2 モデルをトレーニングし、図 5 でハイブリッド並列処理の通信上の利点を実証しました。トレーニングの通信フェーズを強調するために、トレーニングはノード間帯域幅が低い 4 ノードクラスターで実行されます。

この場合、モデルが小さくなり、ノード内の帯域幅が低くなるため、モデルの並列処理には利点がありません。
パイプライン並列処理の通信量は、コンフィグレーションデータやモデル並列処理に比べて一桁少ないです。バッチサイズが小さい場合、トレーニング速度は 7 倍速くなります。
データ並列処理では、通信オーバーヘッドを均等に償却するために勾配累積によるバッチサイズの増加が使用されますが、バッチサイズが大きくなっても、構成されたパイプライン並列処理のパフォーマンスはデータ並列処理の 2 倍になります。
ハイブリッドパイプラインとデータ並列構成は、データ並列グループをノード内の GPU に制限することで勾配通信のボトルネックを回避するため、より高速なノード内帯域幅による勾配通信のメリットが得られます。

図 5: シーケンス長 1024 で GPT-2 (1.5B パラメーター) をトレーニングする場合のスループットとバッチサイズ。トレーニングには、それぞれ 4 つの V100 GPU と 16 GB のメモリを搭載した 4 つのノードが使用されました。GPU は、50 Gbps のノード内帯域幅と 4 Gbps のノード間帯域幅で相互に接続されます。DP は、ZeRO-1 対応データ並列処理を表します。すべてのメソッドは、勾配累積のステップ数を増やすことによってバッチサイズを調整します。

ZeRO-Offload: 10 倍大きいモデルを単一の GPU でトレーニング

ZeRO-Offload は、GPU とホスト CPU のコンピューティングリソースとストレージリソースを同時に利用することで、より少ない GPU リソースで効率的にトレーニングできるモデルの最大サイズを増やします。これにより、GPU あたり 30Tflop の高いトレーニングスループットを維持しながら、単一の V100 で現在の最高レベルの 10 倍である最大 1,300 億のパラメーターを使用してモデルをトレーニングすることができます。

ZeRO-Offload は、単一の GPU で数十億のパラメーターを使用してモデルをトレーニングできるようにすることで、大規模なモデルのトレーニングにアクセスできるようにし、ハードウェアリソースが限られている深層学習の実践者が参加できるようにします。

単一 GPU でデフォルトの PyTorch と ZeRO-Offload を使用してトレーニングできる最大モデルサイズのヒストグラム。

図 6: デフォルトの PyTorch と ZeRO-Offload を使用して単一の GPU でトレーニングできる最大モデルサイズ。

ZeRO-Offload の背後にあるコアテクノロジは、オプティマイザの状態と勾配をZeRO-2 に基づいて CPU メモリにオフロードすることです。この方法により、ZeRO-Offload は CPU へのコピーによる計算効率の損失を最小限に抑えながら、ZeRO-2 と同じ、または場合によってはそれを超える効率を達成できます。次の図は、Zero-OffLoad のアーキテクチャを示しています。

図 7: ZeRO-Offload の概要。

ZeRO-Offload が単一の GPU で数十億のパラメータモデルをトレーニングする方法を学びます

GPT や T5 などの数十億のパラメーターを使用してモデルをトレーニングするには、モデルと状態を保存するために複数の GPU が必要です。大規模なモデルのトレーニングでは、GPU 間でのモデルの並列処理によってメモリ制限の問題がほとんど解決されます。最近、モデルの状態 (オプティマイザーの状態、勾配、モデルパラメーター) を複数の並列 GPU に分散するメモリ効率の高いオプティマイザーである ZeRO をリリースしました。これにより、モデルの並列トレーニングを使用せずに数十億のパラメーターを持つモデルを処理できるようになります。ただし、ZeRO では分割されたモデルの状態を保存するために依然として多数のデータ並列 GPU が必要であるため、このモデルをトレーニングできる条件を備えている人は限られています。

ZeRO-Offload を使用すると、大規模なモデルのトレーニングを単一の GPU で実行できるため、そのようなトレーニングを民間で行うことができます。複数の GPU を使用せずに数十億のパラメーターを使用してモデルをトレーニングするために、ZeRO-Offload は、オプティマイザーの状態量と勾配を分割する ZeRO-2 の方法を継承しています。ZeRO-2 との違いは、ZeRO-Offload はオプティマイザーの状態と勾配の一部を各 GPU に保存せず、両方をローカルメモリに移動することです。オプティマイザーの状態は、トレーニングプロセス全体を通じてメモリに保持されます。勾配は、逆計算プロセス中に GPU 上で計算され、reduce-scatter によって平均化されます。その後、各データ並列プロセスは、平均化された勾配の独自の部分を CPU にオフロードし (図 7 の g オフロード)、その部分を破棄します。あなたの責任ではありません。

勾配がCPUに到達すると、分割された最適化状態量がCPU上で並行して更新されます( 図7のp update )。更新が完了すると、分割されたパラメータは GPU に戻され、all Gather 操作 ( 図 7 のg swap ) を使用して更新されます。Zero-Offload は、さまざまな CUDA ストリームを使用して通信 ( g offload や g swapなど) と計算 (バックプロパゲーションや p updateなど) をオーバーラップさせることにより、トレーニング効率も向上します。

モデルサイズ、トレーニング速度、スケーラビリティに関する ZeRO-Offload の利点をご覧ください。

10 倍のモデル拡張: 単一の 32 GB V100 GPU 上で、図 6 は、PyTorch が最大 13 億のパラメーターを使用してモデルをトレーニングできるのに対し、ZeRO-Offload は PyTorch の 10 倍である 130 億のパラメーターを使用してモデルをトレーニングできることを示しています。これは、ZeRO-Offload がトレーニングプロセス全体を通じて、GPU メモリの大部分を消費するオプティマイザの状態をネイティブメモリに保持し、逆伝播中に計算された勾配を CPU に移動するためです。したがって、節約された GPU メモリを使用して、より大きなモデルをトレーニングできます。

効率的なトレーニングスループット: 図 8 に示すように、100 億パラメータモデルをトレーニングする場合、トレーニングに 1 つの GPU のみが使用されている場合でも、ZeRO-Offload を使用すると、各 GPU で 30 Tflops を超えるスループットを実現できます。スループット GPU の数に応じてほぼ完全に直線的に増加します。

ZeRO-Offload は ZeRO-2 を完全に補完するもので、少数の GPU で大規模なモデルを効率的にトレーニングできます。ZeRO-Offload を使用すると、CPU メモリを利用してモデルに必要な GPU メモリを削減することで、1 ～ 16 GPU で大規模なモデルをトレーニングすることが可能になります。32 GPU では、ZeRO-Offload のパフォーマンスは ZeRO-2 のパフォーマンスよりわずかに高くなります。パフォーマンスの向上は、ZeRO-Offload によって節約された GPU メモリによってもたらされ、より大きなバッチでモデルをトレーニングできるため、 CPU オーバーヘッドへのコピーの影響により、GPU のコンピューティング効率をさらに向上させることができます。より多くの GPU (例: 64 と 128) を使用すると、ZeRO-2 は ZeRO-Offload よりも優れたパフォーマンスを発揮します。これは、どちらも同様のサイズのバッチを実行できるようになり、ZeRO-2 にはデータを CPU に移動するオーバーヘッドがなく、オプティマイザーの更新が GPU 上ではるかに高速になるためです。 CPUよりも。全体として、ZeRO-Offload は ZeRO-2 を補完するものであり、大規模なモデルトレーニング用の最適化ソリューションにより、ZeRO ファミリの最適化範囲を単一のデバイスから数千のデバイスまで拡張します。

ZeRO-Offload と ZeRO-2 を使用して 128 GPU で 100 億パラメーターの GPT-2 モデルをトレーニングした場合のスループットのヒストグラム。

図 8: 128 個の GPU を使用して 100 億パラメーターの GPT-2 モデルをトレーニングした ZeRO-Offload と ZeRO-2 のトレーニングスループットの比較。

DeepSpeed スパースアテンション: 10 倍長いシーケンスを 6 倍高速に実行

アテンションベースの深層学習モデル (トランスフォーマーなど) は、トークン間の距離が遠い場合でも、入力シーケンス内のトークン間の関係を捕捉するのに非常に効果的です。したがって、テキスト、画像、音声関連の入力でよく使用されます。これらの入力のシーケンス長は、数千のトークンになる場合があります。ただし、アテンションモジュールは長いシーケンス内の依存関係を効果的にキャプチャしますが、実際のアプリケーションでは、長いシーケンス入力のサポートは計算量とビデオメモリによって制限されます。計算量とメモリ要件は、シーケンスの長さ \(n\) に対して二次関数的に増加します。

この制限に対処するために、DeepSpeed はスパースアテンションカーネルを提供します。これは、ブロックスパース計算を通じてアテンション計算の計算要件とメモリ要件を桁違いに削減する手段テクノロジーです。このツールセットは、アテンション計算のメモリボトルネックを軽減するだけでなく、スパース計算も非常に効率的です。その API は、Transformer ベースのモデルに簡単に統合できます。さまざまなスパース構造を提供するだけでなく、ユーザー定義のブロックスパース構造にも柔軟に対応できます。

より具体的には、スパースアテンション (SA) は、近いトークン間のローカルアテンションを計算したり、ローカルアテンションを使用してグローバルアテンションを取得することでサマリートークンを計算したりするように設計できます。さらに、SA は、図 10 の青、オレンジ、緑のブロックに示すように、ランダムアテンションと、ローカルアテンション、グローバルアテンション、およびランダムアテンションの任意の組み合わせの両方をサポートします。これにより、SA はメモリフットプリントを \(O(wn)\) に減らすことができます。ここで 1\(<w≤n \) は、その値がアテンション構造に依存するパラメータです。

小さな色付きの四角形は、可変のスパース構造を示しています

図 10: 可変スパース構造

GPU での効率的な実装: スパースアテンションの基本的な実装はビデオメモリを節約しますが、計算的には密な計算よりも劣る可能性があります。これは主に、データがまばらであるためメモリアクセスが分散していることが原因です。効率的なスパースカーネルの開発は、特に GPU 上で困難なことがよくあります。DeepSpeed は、 Triton で開発された効率的なスパースアテンションカーネルを提供します。これらのカーネルは、ブロックのようなスパースパラダイムで構造化されており、アライメントされたメモリアクセスを可能にし、GPU スレッドの分岐を減らし、プロセッサ上のワークロードのバランスをとります。

システムパフォーマンス: 図 11 に示すように、SA は 10 倍長いシーケンスと最大 6.3 倍の計算速度向上をサポートします。左の図は、BERT-Base と BERT-Large で実行できる最長シーケンス長を示しています。私たちの実験には次の 3 つの設定があります。デンスモード、チェックポイントがアクティブ化されたデンスモード、およびチェックポイントがアクティブ化されたスパース (SA) モードです。BERT-Base および BERT-Large の高密度モードと比較して、SA のシーケンスはそれぞれ 10 倍および 16 倍長くなります。さらに、高密度モードと比較して、SA は総計算量を削減し、トレーニング速度を向上させます: 効率の向上はシーケンスの長さに応じて増加し、BERT-Base では最大 6.3 倍、BERT-Base では最大 6.3 倍高速になります。 5.3倍に。

図 11: BERT モデルでサポートされている最大シーケンス長 (左)、単一の NVIDIA V100 GPU で異なるシーケンス長を使用して BERT-Base (中央) と BERT-Large (右) をトレーニングするのにかかる時間。

SA が完全な密な注意と同等またはそれ以上の精度を達成する方法を学びましょう

まばらな注意を伴う関連研究 ( Sparse Transformer、Longformer、BigBird ) はすべて、私たちの経験と一致して、完全な注意よりも高い精度を示しています。メモリのオーバーヘッドが低くなり、計算が高速になることに加えて、実稼働モデルでは SA の精度が高く、収束が高速になることも観察されています。以下の図は、長いテキストを理解するためのBERT ベースの実動モデルのトレーニングの精度を示しています (シーケンス長 2048)。実験は次の 3 つの設定で実行されます: スクラッチからの高密度トレーニング、スクラッチからの SA トレーニング、およびシーケンス長 512 を使用した高密度チェックポイントから継続する SA トレーニング。スクラッチからの事前トレーニングの場合、SA は密な設定よりも速く、より高い精度で収束することが観察されています。さらに、時間と精度の点で、事前にトレーニングされたチェックポイントを SA でトレーニングし続けることの効果はさらに優れています。

図 12: 長文理解アプリケーションの精度

SA と最新の LongFormer を比較する

SA と、最近のスパース構造である Longformer およびその実装を比較しました。私たちの実験では、SA は「固定」スパース性を使用します。両方の実装の精度は同等です。システムパフォーマンスの点では、SA はトレーニングと推論の両方で Longformer よりも優れています。

Wikitext103 での事前トレーニング済み MLM の実行が 1.5 倍高速になりました
BERT-Baseの推論速度が3倍に向上（バッチサイズ1、シーケンス長2,048）

あらゆるブロックのスパース構造を処理する柔軟性: DeepSpeed Sparse Attendance スイートは、特定のスパース構造をターゲットにしていないため、モデル研究者があらゆるブロックのスパース構造を探索できるように効果的にサポートします。現在、 Fixed (OpenAI Sparse Transformer より)、[BigBird]( https://arxiv.org/pdf/2007.14062.pdf ) (Google より)、BSLongformer ( AI2 Longformerの一般的なスパース構造を追加しています。。また、図 10 に示すように、「変数」構造を持つテンプレートも定義します。これを使用すると、ランダム、ローカル、またはグローバルのアテンションパターンのブロックスパース構造を簡単にカスタマイズできます。

1 bit Adam: 通信量を5倍に削減し、学習速度を3.4倍に向上

BERT や GPT-3 などの大規模モデルのスケーリングトレーニングには、モデルの設計、アーキテクチャ、システム機能に基づいた慎重な最適化が必要です。システムの観点から見ると、特に標準の TCP を使用し、ネットワーク帯域幅が限られている商用システムでは、通信効率が大きなボトルネックになっています。

通信圧縮は、このようなシステムでのトレーニング時間を短縮するための重要な技術です。通信を圧縮する最も効果的な方法の 1 つはエラー補償圧縮であり、1 ビット圧縮でも安定した収束速度が得られます。ただし、最新の誤差補正技術は、確率的勾配降下法 (SGD) や Momentum SGD など、勾配に線形に関連する一部の単純なオプティマイザーにのみ適しています。これらの手法は、多くのタスク (BERT のようなモデルのトレーニングを含む) で最高の収束率と精度をもたらす Adam などの非線形オプティマイザーと統合できません。

Adam のような強力なオプティマイザにとって、勾配の非線形特性 (分散項) に依存しているため、誤差補正ベースの圧縮技術を開発することは非常に困難であり、高度な通信圧縮技術の実用的な価値が制限されています。

古典的な圧縮技術の背景を理解する

通信圧縮の 1 つの方法は 1 ビット圧縮であり、次のように表すことができます。

この圧縮では、各数値を 1 ビットで表すため、メモリ要件が 32 分の 1 に削減されます。問題は、この直接法では収束速度が大幅に低下し、実用的価値がほとんどないことです。最近の研究では、エラー補償圧縮を使用することで、通信圧縮下でもほぼ同じ収束率を保証できることが示されています。

誤差補正の考え方は次のように要約できます。1) 圧縮を実行し、2) 圧縮誤差を記憶し、3) 次の反復で圧縮誤差を追加し直します。SGD の場合、エラー圧縮は次と同等です。

このうち \(C(⋅)\) は 1 ビットの圧縮演算子です。この種の誤差圧縮の利点は、圧縮誤差 \(e_t\) と \(e_t-1\) の履歴値が最終的に互いに打ち消し合うことです。これにより、次のようになります。

この戦略は、SGD や Momentum SGD など、勾配に線形に依存するすべての最適化アルゴリズムで機能することが証明されています。

Adam に誤差補正を適用する際の課題を理解する

アダムのアルゴリズムの概要を以下に示します。更新ルールは次のとおりです。

上の式に示されているように、分散項 \(v_t\) と勾配 \(g_t\) には非スレッド関係があります。Adam に対して通常の誤差補正を実行すると、Adam が収束しないことがわかります (図 13 を参照)。

図 13: 勾配に対する非線形依存性のため、誤差補正圧縮は Adam では機能しません

1ビットAdamによる圧縮通信

Adam オプティマイザーの使用時に通信を圧縮するために、前処理を通じて勾配の非線形依存性の問題を解決する1 ビット Adamを開発しました。非線形項の分散 (\(v_t\)) の分散は、いくつかのトレーニングエポック後に大幅に減少し、その後 \(v_t\) を定数に設定しても収束率は変化しないことが観察されます。したがって、提案された 1 ビット Adam オプティマイザは、(図 14 に示すように) 2 つの部分で構成されます。ウォームアップフェーズは、本質的にオリジナルの Adam アルゴリズムです。圧縮ステージでは、分散項を一定に保ち、残りの線形項 (つまり、運動量) を 1 ビット表現に圧縮します。

アルゴリズムの圧縮段階は、しきい値パラメーターによって制御されます (図 14 を参照)。「分散」の変化が特定のしきい値を下回ったことが検出されると、圧縮段階に切り替わります。私たちの調査によると、ウォームアップ段階にはトレーニングステップ全体の 15 ～ 20% のみが必要です。

1 ビット Adam の基礎となるメカニズムについて詳しく学ぶ

1-bit Adam の重みは次の式に従って更新されます。i番目のワーカーの場合、圧縮フェーズ中に次のようになります。

テキストのスクリーンショット

携帯電話のスクリーンショット

図 14: 従来の Adam アルゴリズムと 1 ビット圧縮 Adam アルゴリズムを使用した分散トレーニングプロセスの比較

1 ビット Adam のシステム課題への対処

アルゴリズム上の課題に加えて、トレーニングシステムに 1 ビット Adam を適用するには 2 つのシステム上の課題があります。まず、運動量を 1 ビット表現に変換できる効率的なカーネルが必要です。次に、異なる GPU 間で圧縮された運動量を転送するための効率的な通信スキームが必要です。圧縮の目的は、全体のトレーニング時間を短縮して、帯域幅に制約のある汎用システムを大規模なモデルのトレーニングに使用できるようにすることです。DeepSpeed ではこれらの困難な問題に対処し、通信効率が限られたシステムでのトレーニングのシナリオ向けに 1 ビット Adam 実装を完全に最適化します。

通信に制約のあるシステムにおける 1 ビット Adam の利点

1 ビット Adam は、Adam と同等のコンバージェンス機能を提供し、 通信量を最大 5 倍削減できます。BERT-Large の事前トレーニングタスクに使用すると、SQuAD ファインに使用した場合の最大 3.5 倍のスループットを達成できます。チューニングタスクにより、最大 2.7 倍のスループットが向上します。エンドツーエンドのスループットの向上は、圧縮段階で観察された 6.6 倍 (図 15 左) および 6.2 倍 (図 15 右) の速度向上によるものです。当社の 1 ビット Adam オプティマイザーは 40 Gb イーサネットシステム上で非常によく拡張でき、そのパフォーマンスは 40 Gb InfiniBand QDR システム上の Adam のスケーラビリティに匹敵することは注目に値します。iPerf ベンチマークに基づくと、40 Gb イーサネットの実効帯域幅は 4.1 Gbps ですが、InfiniBand perftest マイクロベンチマークに基づくと、InfiniBand は 32 Gbps のピークに近い帯域幅を提供します。

図 15: NVIDIA V100 GPU での BERT-Large 事前トレーニング (左) と SQuAD 微調整 (右) の 1 ビット Adam スケーラビリティ。BERT 事前トレーニングのバッチサイズは 16/GPU、SQuAD 微調整のバッチサイズは 3/GPU です。

1 Bit Adam のレビュー結果をさらに詳しく見る

Adam と同じ収束: 1 ビット Adam を使用する場合の主な問題は、収束の速度です。同じ数のトレーニングサンプルを使用すると、1 ビット Adam が同じ収束速度と同等のパフォーマンスを達成できることがわかりました (図 16 を参照)。

図 16: 同じ数のトレーニングサンプルを使用すると、1 ビット Adam も Adam と同様に収束できます。

表 1 に、BERT-Base と BERT-Large の詳細な結果を示します。非圧縮の場合と圧縮の場合の両方で、1 ビット Adam のパフォーマンスは元のモデルと同等であり、一部は元のモデルより優れていることがわかります。

表 1: さまざまなテストタスクでの 1 ビット Adam の正確性の検証

通信量を最大 5 分の 1 に削減: 1 ビット Adam は、Adam と同じ収束機能を提供し、圧縮フェーズ中に通信を 16 分の 1 に削減します (16 ビット (FP16) トレーニングの場合)。BERT 事前トレーニングモデルの場合、ウォームアップフェーズはエンドツーエンドトレーニング時間の 15% にすぎないことが観察されているため、全体の通信量は 5 分の 1 に削減されます。

オリジナルのAdamと1bit Adamの通信量の比率の計算式は以下の通りです。

1 / (ウォームアップ + (1 – ウォームアップ)/16)

1 ビット Adam により、BERT-Large のトレーニングが 3.5 倍高速になります。 帯域幅の制約が制限された 2 つのシステムで BERT-Large をトレーニングした結果が提供されます。1) 40 Gbps イーサネット (図 17 左) および 2) 40 Gbps InfiniBand QDR (図 17 右) ）。圧縮フェーズでは、イーサネットを使用するとシステムスループットが 6.6 倍、InfiniBand を使用するとシステムスループットが 2 倍向上し、エンドツーエンドの速度 (ウォームアップおよび圧縮フェーズを含む) はそれぞれ 3.5 倍と 2.7 倍向上しました。。1 ビット Adam は主に、通信量の削減 (運動量通信の圧縮実装による) と、効率的な 1 ビットのノンブロッキングギャザーと オールギャザーオペレーションを通じて実装されたカスタムallreduceオペレーションから恩恵を受けます。

BERT の事前トレーニングに Adam オプティマイザーの代わりに LAMB を使用して、合計バッチサイズを増やして通信量を削減することも可能であることに注目してください。ただし、1 ビット Adam は、この要求の厳しいハイパーパラメータ調整を回避します。私たちの経験によれば、大規模なバッチでパラメータを調整することは通常より困難です。さらに、1 ビット Adam は、重要なバッチサイズが小さい作業にも非常に適しています (多くの微調整タスクなど、大きなバッチではうまく収束できません)。

図 17: 圧縮フェーズ中の 40 Gbps イーサネット (左) および InfiniBand (右) 上の 1 ビット Adam を使用した BERT-Large トレーニングのパフォーマンス

1 ビット Adam は、SQuAD 微調整タスクを 2.7 倍高速化します。1 ビット Adam は、大規模なトレーニングタスクでスケーラビリティを提供するだけでなく、SQuAD 微調整などのタスクでも適切に機能します。図 18 に示すように、1 ビット Adam はイーサネットベースと InfiniBand ベースの両方のシステムで適切に拡張され、イーサネットベースのシステムで最大 6.2 倍の高いスループット (圧縮段階で) を提供し、結果として 2.7 倍のエンドツー- 終了速度の向上 (ウォームアップ段階で 25%、圧縮段階で 75%)。SQuAD の微調整では、合計バッチサイズが 96 のときに F1 スコアが最も高くなることがわかります。バッチサイズがこれより大きい場合、収束率が低下し、追加のハイパーパラメータ調整が必要になります。したがって、32 GPU に拡張するには、各 GPU で 3 ～ 4 の小さなバッチを実行します。そのため、微調整タスクの通信が集中し、拡張が困難になります。1 ビット Adam はスケーラビリティの問題をうまく解決し、バッチを増やすことなく通信トラフィックを 3.4 倍削減し、2.7 倍のエンドツーエンドの高速化を実現します。

図 18: 40 Gbps イーサネット (左) および InfiniBand (右) 上の SQuAD 微調整タスクで 1 ビット Adam を使用した場合の圧縮ステージのパフォーマンス。

これらの新しいテクノロジーのコード、チュートリアル、ドキュメントについては、DeepSpeed Web サイトと Github リポジトリをチェックしてください。また、いくつかのテクノロジーをONNX ランタイムに統合しました。

私たちの素晴らしい協力者について:

学術協力者であるハーバード大学のフィリップ・ティレット氏に感謝いたします。彼は、 Tritonコンパイラを通じて、スパースアテンションアルゴリズムカーネルを私たちと一緒に開発しました。
ZeRO-Offload は、カリフォルニア大学マーセド校のインターンである Jie Ren 氏とともに開発されました。また、このトピックについて議論してくださった、 UC Merced の Dong Li 氏、 L2L 作業に携わった Microsoft の Bharadwaj Pudipeddi 氏と Maral Mesmakhuroshahi 氏にも感謝します。
1-bit Adam は、ロチェスター大学のインターンである Hanlin Tang によって共同開発されました。
また、NVIDIA、特に Megatron-LM チームの強力な協力にも感謝します。

DeepSpeed チームについて:

私たちは、大規模なシステムパフォーマンスの最適化に情熱を注ぐ研究者とエンジニアのグループです - Samyam Rajbhandari、Jeff Rasley、Olatunji Ruwase、Reza Yazdani Aminabadi、Elton Zheng、Arash Ashari、Jing Zhao、Minjia Zhang、Niranjan Uma Naresh、Shaden Smith、Ammar Ahmad Awan、Conglong Li、Yuxiong He (チームリーダー)。最近では、深層学習システムに焦点を当てており、深層学習システムのトレーニング速度、収束速度、開発速度の最適化に取り組んでいます。