DeepSpeed ZeRO++: ネットワーク通信を4倍削減し、大規模モデルやChatGPTのようなモデルのトレーニング効率を大幅に向上

大規模な AI モデルがデジタルの世界を変えています。大規模言語モデル (LLM) に基づく Turing-NLG、ChatGPT、GPT-4 などの生成言語モデルは広く使用されており、要約、コード生成、翻訳などのタスクを実行できます。同様に、DALL・E、Microsoft Designer、Bing Image Creator などの大規模なマルチモーダル生成モデルは、アート、建築、ビデオ、その他のデジタル資産を生成できるため、コンテンツ作成者、建築家、エンジニアはまったく新しい創造的な生産性を探求できます。

ただし、これらの大規模なモデルをトレーニングするには、数百または数千の GPU デバイス上で大量のメモリと計算リソースが必要になります。たとえば、Megatron-Turing NLG 530B モデルをトレーニングするには、4,000 個を超える NVidia A100 GPU を使用する必要があります。これらのリソースを効果的に活用するには、モデルを各デバイスのメモリに適切に分散し、これらのデバイス間で計算を効果的に並列化するための複雑な最適化システムが必要です。同時に、ディープラーニング コミュニティが大規模なモデルを簡単にトレーニングできるようにするには、これらの最適化が使いやすくなければなりません。

DeepSpeed の ZeRO 最適化ファミリーは、これらの課題に対する強力なソリューションを提供し、TNLG-17B、Bloom-176B、MPT-7B、Jurrasic-1 などの大規模なディープラーニング モデルのトレーニングに広く使用されています。変革能力にもかかわらず、いくつかの重要なシナリオでは、ZeRO は GPU 間で大量のデータ転送オーバーヘッドを発生させ、トレーニングの効率を低下させます。これは特に、次のシナリオで発生します: a) グローバル バッチ サイズが小さく、GPU の数が多いため、各 GPU のバッチ サイズが小さくなり、頻繁な通信が必要になる場合。または、b) ノード間のネットワーク帯域幅が制限されているローエンド クラスターでトレーニングする場合、通信のレイテンシが長くなります。このような場合、ZeRO のトレーニング効率は制限されます。

これらの制限に対処するために、ZeRO++ をリリースしました。 ZeRO++ は、モデルの品質を損なうことなく、ZeRO と比較して総通信量を 4 分の 1 に削減します。これには 2 つの重要な意味があります。

1. ZeRO++は大規模モデルの事前トレーニングと微調整を加速します

• GPU あたりのバッチ サイズが小さい: 数千の GPU で大規模なモデルを事前トレーニングする場合でも、数百または数十の GPU で微調整する場合でも、GPU あたりのバッチ サイズが小さい場合、ZeRO++ は ZeRO よりも 2.2 倍高いスループットを提供し、トレーニング時間とコストを直接削減します。
• 低帯域幅コンピューティング クラスター: ZeRO++ を使用すると、低帯域幅クラスターでも、4 倍の帯域幅を持つハイエンド クラスターと同様のスループットを実現できます。したがって、ZeRO++ は、より広範囲のクラスターにわたって大規模なモデルを効率的にトレーニングできます。

2. ZeRO++ が ChatGPT の RLHF トレーニングを加速

• ZeRO++ は主にトレーニング用に設計されていますが、通信オーバーヘッドは ZeRO のトレーニングと推論の両方に同様に適用されるため、その最適化は自動的に ZeRO-Inference にも適用されます。したがって、ZeRO++ は、RLHF がトレーニングと推論を組み合わせるため、人間のフィードバックによる強化学習 (RLHF) などのアルゴリズムの効率を向上させることができます。
• ZeRO++ は、DeepSpeed-Chat との統合により、オリジナルの ZeRO と比較して、RLHF トレーニングの生成フェーズの効率を最大 2 倍、強化学習トレーニング フェーズの効率を最大 1.3 倍向上させることができます。

次に、ZeRO とその通信オーバーヘッドについてさらに詳しく説明し、これらの問題に対処するために ZeRO++ で行われた主要な最適化について説明します。次に、さまざまなモデル サイズ、バッチ サイズ、帯域幅の制約に対して、ZeRO++ がトレーニング スループットに与える影響を示します。また、RLHF を使用した会話モデルのトレーニングを加速するために、ZeRO++ を DeepSpeed-Chat に適用する方法についても説明します。

ZeRO++ 詳細説明

図 2: ZeRO オプティマイザーのワークフロー図 (これは部分的な表示です。完全なプロセスについては、元の Zhihu の記事を参照してください)

ZeRO は、データ並列処理のメモリ効率の高いバージョンであり、モデルの状態は、ギャザー/ブロードキャストベースの通信を使用してトレーニング中にコピーおよび再構築する必要なく、すべての GPU に分割されて保存されます。これにより、ZeRO は、すべてのデバイスの集約された GPU メモリと計算能力を効率的に活用しながら、使いやすいデータ並列トレーニングを提供できるようになります。

モデルサイズがMであると仮定します。フォワード パス中、ZeRO は必要に応じて全収集/ブロードキャスト操作を実行し、各モデル レイヤーのパラメーターを収集します (合計サイズは M)。後方パスでは、ZeRO は各レイヤーのパラメーターに対して同様の通信パターンを採用し、ローカル勾配 (合計サイズ M) を計算します。さらに、ZeRO は、各ローカル勾配を計算した直後に、reduce または Reduce-Scatter 通信を使用して、ストレージ (合計サイズは M) を平均化および分割します。したがって、ZeRO には合計 3M のトラフィックがあり、2 つの all-gather/broadcast 操作と 1 つの Reduce-Scatter/Reduce 操作に均等に分散されます。

これらの通信オーバーヘッドを削減するために、ZeRO++ は上記の 3 つの通信セットを対象に、3 セットの通信最適化を実行します。

図3: qwZのパーティション量子化の図

ZeRO通信における重み量子化（qwZ）

まず、オールギャザー時のパラメータ通信量を削減するために、重み量子化を使用して、通信前に各モデルパラメータを FP16 (2 バイト) から INT8 (1 バイト) データ型に動的に削減し、通信後に重みを逆量子化します。ただし、重みを単純に量子化すると、モデルトレーニングの精度が低下します。良好なモデルトレーニング精度を維持するために、モデルパラメータの各サブセットを個別に量子化するパーティション量子化を使用します。現在、パーティション量子化のための高性能な実装は存在しません。そのため、私たちは高度に最適化された量子化 CUDA カーネルのセットをゼロから実装し、基本的な量子化よりも 3 倍の精度と 5 倍の速度を実現しました。

図4: 重みの階層的分割 (hpZ)

ZeRO モデルの重みの階層的分割 (hpZ)

次に、バックワードパス中の全収集重みの通信オーバーヘッドを削減するために、通信に GPU メモリを使用します。具体的には、ZeRO のようにモデルの重み全体をすべてのマシンに分散するのではなく、各マシンにモデルの完全なコピーを保持します。これにより、メモリのオーバーヘッドが高くなりますが、コストのかかるマシン間の全収集/ブロードキャストを、モデルの重みのマシン内の全収集/ブロードキャストに置き換えることができます。これにより、マシン内の通信帯域幅が高くなるため、通信速度が大幅に向上します。

図5: qgZエンドツーエンドワークフロー

ZeRO通信中の勾配量子化（qgZ）

3 番目に、勾配の Reduce-Scatter 通信コストを削減することはより困難です。量子化を直接適用して通信量を削減することは現実的ではないためです。量子化誤差を減らすために分割量子化を使用したとしても、勾配削減によって量子化誤差が蓄積され、増幅されます。この問題に対処するために、通信の前にのみ勾配を量子化し、削減操作の前に元の精度に逆量子化します。これを効率的に行うために、我々は qgZ と呼ばれる、圧縮された Reduce-Scatter 操作と機能的に同等な、新しい全対全ベースの量子化勾配通信パラダイムを発明しました。

qgZ は、次の 2 つの課題に対処することを目的としています。i) INT4/INT8 で単純に Reduce-Scatter を実装すると、精度が大幅に低下します。ii) 従来のツリーまたはリングベースの Reduce-Scatter で量子化を使用するには、量子化と逆量子化の手順を長く繰り返す必要があり、これにより、Reduce を完全な精度で実行した場合でも、エラーが蓄積され、遅延が大きくなります。これら 2 つの課題に対処するために、qgZ はツリーまたはリングベースの分散削減アルゴリズムを使用せず、新しい階層型の all-to-all アプローチに基づいています。

qgZ には主に 3 つのステップがあります。

• グラデーションスライスの並べ替え。
• ノード内通信と削減;
• ノード間の通信と削減。

まず、通信が発生する前に、勾配をスライスしてテンソル スライスを並べ替え、通信終了時の各 GPU の最終的な勾配の位置 (図 5 の緑色のブロック) が正しいことを確認します。次に、並べ替えられた勾配スライスを量子化し、各ノード内で全対全通信を実行し、全対全から受信した勾配スライスを逆量子化し、ローカル削減を実行します。 3 番目に、ローカル削減後の勾配を再量子化し、ノード間で全対全通信を実行し、受信した勾配を再度逆量子化し、最終的な高精度勾配削減を計算して、図 5 の緑色のブロックの結果を取得します。

この階層化アプローチの理由は、ノード間の通信の量を減らすためです。より正確には、ノードあたり N 個の GPU、M モデル サイズ、Z 量子化比が与えられた場合、シングル ホップの all-to-all では M*N/Z のクロスノード トラフィックが生成されます。対照的に、この階層的なアプローチでは、GPU あたりのノード間トラフィックを M/Z から M/(Z*N) に削減します。したがって、総通信量は M*N/Z から M*N/(Z*N) = M/Z に削減されます。さらに、ノード内通信とノード間通信をオーバーラップし、CUDA カーネル (テンソル スライスの並べ替え + ノード内量子化) と (ノード内逆量子化 + ノード内削減 + ノード間量子化) を融合することで、qgZ のエンドツーエンドのレイテンシを最適化します。

総通信量の最適化

上記の 3 つのコンポーネントをすべて組み合わせることで、ノード間の通信量を 3M から 0.75M に削減しました。具体的には、qwZ を使用して、モデルの重みの前方完全収集/ブロードキャストを M から 0.5M に削減します。 hpZ を使用して、バックプロパゲーション中のノード間の全収集を排除し、通信を M から 0 に削減します。最後に、qgZ を使用して、バックプロパゲーション中のノード間の Reduce-Scatter 通信を M から 0.25M に削減します。

ZeRO++ は大規模言語モデルのトレーニングを加速します

ここでは、384 個の Nvidia V100 GPU での実際の LLM トレーニング シナリオでの ZeRO++ のテスト結果を紹介します。

図 6: 384 個の V100 GPU 上のさまざまなモデル サイズでの ZeRO++ と ZeRO のスループット。ノードは 4 つの Infiniband (IB) リンクを使用して相互接続され、それぞれ 100 Gbps で動作します。

ZeRO++は、GPU上でバッチサイズが小さい場合に高いトレーニング効率を実現します。

高帯域幅クラスター: 図 6 に示すように、まず、4x Infiniband (IB) を使用して 400Gbps のクロスノード相互接続帯域幅を実現し、それぞれ 100Gbps で実行することで、さまざまなモデル サイズとマイクロバッチ サイズで ZeRO++ のスループットが ZeRO よりも向上することを示します。マイクロバッチ サイズが GPU あたり 1k トークンの場合、ZeRO++ は ZeRO-3 と比較してスループットを 28% ～ 36% 向上させます。 2k トークンのマイクロバッチ サイズの場合、ZeRO++ は ZeRO-3 よりも 24% ～ 29% のスループット向上を実現します。

図7: 384 V00 GPU 上の 100Gbps クロスノード帯域幅でのさまざまな LLM のスループット

低帯域幅クラスター: 100Gbps などの低帯域幅ネットワーク環境では、ZeRO++ は ZeRO-3 よりも大幅に優れたパフォーマンスを発揮します。図 7 に示すように、ZeRO++ は、ZeRO-3 と比較して、エンドツーエンドのスループットが最大 2.2 倍向上します。平均して、ZeRO++ は ZeRO-3 ベースラインに比べて約 2 倍の高速化を実現します。

図8: ZeRO++は大幅に削減された帯域幅で高帯域幅のクラスターパフォーマンスを実現します

高帯域幅のZeROクラスターと低帯域幅のZeRO++クラスター間でモデルトレーニング効率の同等性を実現

さらに、ZeRO++ は、はるかに高い帯域幅の環境における ZeRO と比較して、低帯域幅のクラスターで同等のシステム スループットを実現できます。図 8 に示すように、18B および 138B モデル サイズの場合、200Gbps のクロスノード帯域幅を備えた ZeRO++ は、800Gbps のクロスノード帯域幅を備えた ZeRO-3 と同様の TFLOP を達成できます。

ZeRO++ の優れたスケーラビリティを考慮すると、ZeRO++ は大規模な AI モデルをトレーニングするための次世代の ZeRO になると考えられます。

RLHF トレーニングのための ZeRO++ と組み合わせた DeepSpeed-Chat

RLHF トレーニングの紹介

ChatGPT クラス モデルは LLM を活用し、RLHF を使用して微調整されています。 RLHF は、生成 (推論) フェーズとトレーニング フェーズで構成されます。生成フェーズでは、アクター モデルはダイアログの一部を入力として受け取り、一連のフォワード パスを使用して応答を生成します。次に、トレーニング フェーズ中に、批評家モデルは生成された応答を品質に応じてランク付けし、アクター モデルに強化信号を提供します。これらのランキングを使用して参加者モデルを微調整し、後続の反復でより正確で適切な応答を生成できるようにします。

RLHF トレーニングでは、4 つのモデル (アクター、参照、コメント、報酬) を使用するため、メモリに大きな負荷がかかります。一般的な解決策は、低ランク適応トレーニング (LoRA) を採用して RLHF のメモリ圧力を解決することです。 LoRA は、事前トレーニング済みモデルの重みを固定し、トレーニング可能なランク分解マトリックスを Transformer アーキテクチャの各レイヤーに挿入することで、トレーニング可能なパラメータの数を大幅に削減します。 LoRA はメモリ使用量を削減し、バッチ サイズを大きくすることで RLHF を高速化し、スループットを大幅に向上させます。

RLHF トレーニングのための ZeRO++ との DeepSpeed チャット

図9: ZeRO++はRLHFトレーニングの生成とトレーニング段階を高速化します

ZeRO++ は、ほとんどのモデルの重みが固定されているため、RLHF + LoRA シナリオで独自のアプリケーションを備えています。つまり、ZeRO++ は、これらの凍結された重みを fp16 に保存して各通信操作の前に量子化するのではなく、INT4/8 に量子化して保存できます。計算用の重みを準備するために通信後の逆量子化は依然として実行されますが、逆量子化された重みは計算後に単に破棄されます。

このように RLHF トレーニングに ZeRO++ を使用すると、メモリ使用量と通信を削減できます。これは、通信が削減されることでトレーニングのスループットが向上し、メモリ使用量が削減されるためバッチ サイズが大きくできることを意味します。生成フェーズでは、ZeRO++ は hpZ を使用してすべての重み通信を各ノード内に保持し、より高いノード内通信帯域幅を活用して通信量を削減し、生成スループットをさらに向上させます。

ZeRO++ は、ChatGPT のようなモデルの RLHF トレーニングをサポートするために DeepSpeed-Chat に統合されました。図 9 では、さまざまなサイズのアクター モデルの RLHF 生成スループットを比較しています。テスト構成は 32 個の V100 GPU で、アクター モデルのサイズは 30B と 66B であり、ZeRO と ZeRO++ のパフォーマンスをテストします。結果は、ZeRO++ の RLHF 生成スループットが ZeRO の 2.25 倍高いことを示しています。また、16 個の V100 GPU でのトレーニング フェーズでの高速化も実証しました。ZeRO++ では、通信量が少なく、バッチ サイズが大きいため、ZeRO よりも 1.26 倍高いスループットが実現されています。

DeepSpeed ZeRO++がリリースされました！

DeepSpeed ZeRO++ をリリースし、AI コミュニティの皆さんにご利用いただけることを大変嬉しく思っております。 LLM トレーニング チュートリアルを入手するには、GitHub ページにアクセスしてください。 DeepSpeed-Chat 用の ZeRO++ は今後数週間以内にリリースされる予定です。 ZeRO++ の技術的な詳細については、arxiv の論文をご覧ください。

DeepSpeed-ZeRO++ は DeepSpeed エコシステムの一部です。詳細については、詳細なブログ投稿、チュートリアル、役立つドキュメントが掲載されている当社の Web サイトをご覧ください。

最新のDeepSpeedニュースは、英語版Twitter、日本語版Twitter、中国語版Zhihuでもご覧いただけます。

DeepSpeedはあなたの貢献を歓迎します！ DeepSpeed GitHub ページで問題を報告し、PR を投稿し、ディスカッションに参加することをお勧めします。詳細については、寄稿ガイドラインをご覧ください。当社は大学、研究室、企業とのコラボレーションを歓迎します。このようなリクエスト (および GitHub に適さないその他のリクエスト) については、[email protected] に直接メールでお問い合わせください。

寄稿者:

このプロジェクトは、DeepSpeed チームの貢献によって実現しました。

Guanhua Wang、Heyang Qin、Sam Ade Jacobs、Connor Holmes、Samyam Rajbhandari、Olatunji Ruwase、Ammar Ahmad Awan、Jeff Rasley、Michael Wyatt、Yuxiong He (チームリーダー)