PyTorch はどのようにしてデータ並列トレーニングを高速化するのでしょうか?分散型チートが明らかに

[[333298]]

現在、チップのパフォーマンスの向上は限られているため、分散トレーニングは超大規模なデータセットとモデルを処理するための主な方法となっています。この記事では、人気のディープラーニング フレームワーク PyTorch の最新バージョン (v1.5) 向けの分散データ並列パッケージの設計、実装、評価について紹介します。

PyTorch は、ディープラーニングの研究とアプリケーションで広く使用されている科学計算パッケージです。ディープラーニングの最近の進歩により、大規模なデータセットと大規模なモデルの価値が実証され、より多くのコンピューティング リソースを使用してモデル トレーニングを拡張する能力が求められています。

同時に、データ並列処理は、そのシンプルな原理と幅広い適用性により、分散トレーニングの一般的なソリューションとなっています。通常、分散データ並列技術では、各コンピューティング リソースでモデルを複製して勾配を個別に生成し、各反復でこれらの勾配を伝播してモデルのコピーの一貫性を保ちます。この手法は概念的には単純ですが、計算と通信の間に微妙な依存関係があるため、分散トレーニングの効率を最適化するのは簡単ではありません。

したがって、この論文では、Facebook AI とワルシャワ大学の研究者が、PyTorch での分散データ並列モデルの設計、実装、評価を紹介します。

バージョン 1.5 以降、PyTorch 自体は、勾配のバケット化、計算と通信の重複、勾配同期のスキップなど、分散データ並列処理を高速化するいくつかの手法を提供しています。関連する評価結果によると、正しく構成すると、PyTorch 分散データ並列モデルは 256 個の GPU でほぼ線形のスケーラビリティを実現できます。

次に、PyTorch分散データ並列トレーニングのモデル設計、具体的な実装、効果評価について見ていきます。

システム設計

PyTorch は、複数のプロセスとマシンで並列トレーニングを実装するのに役立つ分散データ並列 (DDP) モデルを提供します。分散トレーニング中、各モデルにはモデルの独自のローカル コピーとローカル オプティマイザーが存在します。分散データ並列トレーニングとローカルトレーニングは、エラー訂正の点では数学的に同等である必要があります。

下の図 1 は、Python API フロントエンドと C++ 勾配降下コア アルゴリズムを含み、c10d 集約通信ライブラリを使用する DDP ビルディング ブロックの構成を示しています。

Python API フロントエンド

API を設計する際、研究者は必要な機能を実現するために次の 2 つの設計目標を設定しました。

非侵入的: アプリケーションに提供される API は非侵入的である必要があります。

インターセプティブ: API はさまざまな信号のインターセプトを許可し、適切なアルゴリズムを直ちにトリガーする必要があります。

分散データ並列処理は、より多くのコンピューティング リソースを使用してトレーニングを高速化することを目的としています。

上記の要件に基づいて、研究者は nn.Module を使用して分散データ並列処理を実装しました。 nn.Module は、コンストラクター引数としてローカル モデルを受け取り、バックプロパゲーション中に勾配を透過的に同期します。次のコードは、DDP モデルの使用例です。

勾配降下法

研究者らは、PyTorch での分散データ並列トレーニングのためのいくつかの勾配削減手法について説明しています。 DDP の勾配降下アルゴリズムに新たな改良が加えられました。現在の実装の構造を紹介するために、単純な素朴なソリューションから始めて、徐々により複雑なバージョンを導入し、最終的に PyTorch v1.5.0 の現在のバージョンを使用します。

当初の計画

DDP はまずすべてのトレーニング プロセスを調整して、各プロセスが次の条件を満たしていることを確認します。

• 同じモデル状態から開始します。
• 各反復には同じ量の勾配がかかります。

2 番目のポイントを実現するために、初期スキームでは、ローカル バックプロパゲーションの後、ローカル パラメータを更新する前に勾配同期ステップを挿入します。幸いなことに、PyTorch の autograd エンジンはカスタム後方フックを受け入れることができます。 DDP は、各バックプロパゲーションの後に計算をトリガーするための autograd フックを登録できます。次に、AllReduce 集約通信を使用して、すべてのプロセスにわたる各パラメータの平均勾配の計算を呼び出し、その結果を勾配テンソルに書き戻します。

当初のソリューションは目的を達成するには十分でしたが、パフォーマンス上の欠陥が 2 つありました。集約通信は小さなテンソルではパフォーマンスが低下します。これは、多数の小さなパラメータを持つ大規模モデルの場合に特に当てはまります。両者の間には境界があるため、勾配計算と同期を別々に実行すると、通信重複計算の機会が失われます。

勾配バケット

勾配をバケット化するというアイデアは、大規模なテンソルでは集約通信がより効率的であるという事実に着想を得ています。

下の図 2(a) と (b) は、各 AllReduce でパラメータ数が異なる場合の AllReduce 60M torch float32 パラメータの完全な実行時間を示す定量的なビューを示しています。

これらの実験は、AllReduce を起動する前に各勾配テンソルが利用可能になるまで待つのではなく、より短い時間待機して 1 回の AllReduce 操作で複数の勾配を保存すると、DDP はより高いスループットとより低いレイテンシを実現できることを示しています。

通信重複計算

バケット化の場合、DDP は通信を開始する前に、すべてのコンテンツが同じバケット内に入るまで待機するだけです。このような設定では、バックプロパゲーションの最後に AllReduce をトリガーするだけでは不十分です。そのため、より頻繁な信号に応答し、AllReduce をより迅速に開始する必要があります。したがって、DDP は各勾配累積器に対して autograd フックを登録します。

下の図3(a)の例では、2つの縦軸は時間を表し、点線は勾配が準備できた時間を表しています。プロセス 1 では、4 つの勾配が順番に計算されます。プロセス2では、g_3とg_4の後にg_2が計算されます。図3(b)の例では、勾配g_3に対応するパラメータが1回の反復でスキップされ、g_3の準備完了信号が失われます。

この問題を解決するために、DDP はフォワード パスの出力テンソル内の autograd グラフを走査して、関連するすべてのパラメーターを見つけます。関係するテンソルの準備状態は、バックプロパゲーションが完了したことを示すのに十分です。

次のアルゴリズム 1 は DDP の疑似コードを示します。

下の図 4 は、前方伝播および後方伝播中に DDP がローカル モデルとどのように相互作用するかを示しています。

勾配累積

さらに、DDP は、アプリケーションがバックプロパゲーションの直後に optimizer.step() を呼び出す予定なのか、それとも複数の反復にわたって勾配を蓄積する予定なのかを判断できません。したがって、研究者はこのユースケースに別のインターフェース（つまり、同期なし）を導入する必要があります。以下はサンプルコードスニペットです。

集約コミュニケーション

DDP は、NCCL、Gloo、MPI の 3 つのオプションを含む集合通信ライブラリ上に構築されています。 DDP はこれら 3 つのライブラリから API を取得し、同じ ProcessGroup API にラップします。

すべての通信は集約操作であるため、ProcessGroup インスタンスに対する後続のすべての操作は、そのタイプと一致し、同じ順序に従う必要があります。すべてのライブラリに同じ ProcessGroup API を使用すると、研究者は同じ DDP 実装でさまざまな通信アルゴリズムを試すことができます。

単一の NCCL、Gloo、または MPI ProcessGroup でリンク容量を飽和させることができない場合、DDP はラウンドロビン ProcessGroup を使用してより高い帯域幅使用率を実現できます。

具体的な実装

DDP 実装は以前のバージョンで何度か改善されてきました。研究者らは、PyTorch v1.5.0 の現状を紹介した。 DDP は Python と C++ の両方で実装されており、Python は API を公開してパフォーマンスが重要でないコンポーネントを構成し、C++ はコアとなる勾配降下アルゴリズムを提供します。 Python API は、Pybind11 API を介して C++ カーネルを呼び出します。

Python フロントエンド

Python フロントエンドの実装の詳細によって、DDP の動作が決まります。構成可能なノブは、DDP コンストラクター API で公開されます。具体的には以下が含まれます:

• DDP で AllReduce を実行しているプロセス グループ インスタンスを見つけるためのグループ処理。これにより、デフォルトのプロセス グループとの混乱を避けることができます。
• bucket_cap_mb は AllReduce のバケット サイズを制御します。アプリケーションはノブを調整してトレーニング速度を最適化する必要があります。
• 未使用のパラメータを見つけて、DDP が autograd グラフをトラバースして未使用のパラメータを検出する必要があるかどうかを確認します。

ローカル モデルのモデル デバイス アフィニティは、特にモデルが非常に大きく、複数のデバイスで実行する必要がある場合に、DDP の動作を制御することもできます。大規模なモデルの場合、モデルの各レイヤーを異なるデバイスに配置し、中間出力を Tensor.to(device) API を使用してあるデバイスから別のデバイスに転送できます。 DDP は複数のモデルで実行することもできます。

モデル バッファーは、レイヤー (BatchNorm など) が実行中の分散や平均などの状態を追跡する必要がある場合に必要です。 DDP は、ランク 0 のプロセスを承認できるようにすることで、モデル バッファーをサポートします。

コア勾配降下法

開発プロセスにおける主な作業は勾配削減であり、これは DDP のパフォーマンスを決定する重要なステップでもあります。このreducer.cppの実装には、パラメータとバケットのマップの構築、autogradフックのインストール、バケットAllReduceの開始、未使用のグローバルパラメータの検出という4つの主要コンポーネントがあります。

パラメータとバケットのマッピングは、DDP の速度に大きな影響を与えました。各バックプロパゲーションでは、すべてのパラメータ勾配からのテンソルがバケットにコピーされ、平均勾配が AllReduce 後のテンソルにコピーされます。

Autograd フックは、DDP バックプロパゲーションのエントリ ポイントです。構築中、DDP はモデル内のすべてのパラメータを反復処理し、各パラメータの勾配アキュムレータを見つけ、各勾配アキュムレータに同じポストフック関数をインストールします。勾配アキュムレータは、対応する勾配が準備できたときにポストフックを有効にし、DDP は、バケット全体が AllReduce 操作を開始する準備ができたときにポストフックを有効にします。

Bucket Allreduce は、DDP における通信オーバーヘッドの主な原因です。デフォルトでは、バケット サイズは 25 MB です。

実験的評価

研究者らは、専用の 32 GPU クラスターと共有権限を使用した PyTorch DDP の評価結果を発表しました。GPU は 4 台のサーバーに展開され、Mellanox MT27700 ConnectX-4 100GB/s ネットワーク カードを介して接続されています。各サーバーには 8 個の NVIDIA Tesla V100 GPU が搭載されています。

下の図 5 は、サーバー上で 8 つの GPU がどのように相互接続されているかを示しています。

次に、研究者らは、ResNet50 と BERT という 2 つの一般的なモデルを使用して、各反復における PyTorch DDP のレイテンシとスケーラビリティを測定しました。ほとんどの実験では、各反復におけるレイテンシを比較するのに十分な、ランダムに生成された合成入力とラベルを使用しました。

遅延の失敗

通信オーバーラップ計算の有効性を検証するために、図 6 は、それぞれ NCCL と Gloo を逆方向伝送に使用した場合の ResNet50 モデルと BERT モデルのレイテンシ障害を示しています。すべての実装では、4 台のサーバーで 32 個の GPU を使用しました。

結果は、AllReduce 通信 (つまり、勾配同期) がこのプロセスで完了するため、後方パスが PyTorch DDP トレーニングで最も長いステップであることを示しています。

バケットサイズ

バケット サイズは重要な構成オプションです。経験則として、ベストエフォートの見積もりでは、bucket_cap_mb のデフォルト値は 25 MB です。研究者らは 2 台のマシンで 16 個の GPU を使用して、異なるバケット サイズの各反復のレイテンシを比較しました。もう 1 つの極端な方法は、すべての勾配を短時間で転送することです。その結果は、下の図 7 に示されています。

下の図 8 は、同じ設定で 32 個の GPU を使用した実験結果を示しています。この場合、外れ値の範囲が広くなるのは当然のことです。参加者が増えると、必然的に同期に時間がかかり、絞殺者の効果がより顕著になるからです。

スケーラビリティ

DDP のスケーラビリティを理解するために、研究者は最大 256 個の GPU 上で NCCL と Gloo バックエンドを使用して、ResNet50 と BERT の反復ごとのトレーニング待ち時間を測定しました。結果を下の図9に示します。

下の図 10 は、勾配降下法を 1、2、4、8 回の反復ごとに実行した場合の反復あたりの平均レイテンシを示しています。

反復ごとの遅延に加えて、収束速度を測定して、収束を遅くすることで加速が排除されているかどうかを確認することが重要です。この実験では、MNIST データセットを使用して ResNet をトレーニングします。学習率は 0.02 に設定され、バッチ サイズは 8 です。結果を下の図11(a)に示します。図11(b)はバッチサイズを256、学習率を0.06に設定した場合の測定結果です。

ラウンドロビンプロセスグループ

PyTorch 分散パッケージは、ラウンドロビン プロセス グループと複数の NCCL または Gloo プロセス グループを組み合わせて、ロビンロビン順序で各プロセス グループ インスタンスに集約通信を分散することをサポートします。

下の図 12 は、1、3、および 5 つの NCCL または Gloo プロセス グループを使用したラウンドロビン プロセス グループの反復ごとのレイテンシを示しています。最も顕著な高速化は、NCCL バックエンドを使用する BERT モデルで実現されます。