パイプライン並列処理、テンソル並列処理、および 3D 並列処理

パイプライン線形並列処理、テンソル並列処理、3D 並列処理という 3 つの分散トレーニング手法を、原理から具体的な手法のケースまで詳細に解釈します。

パイプライン線形並列処理とテンソル並列処理はどちらも、限られたシングルカードグラフィックスメモリを使用して大規模なモデルをトレーニングすることを目的として、モデル自体を分割します。簡単に言えば、パイプライン並列処理はモデルを水平方向に分割します。つまり、モデルをレイヤーに従って分割します。テンソル並列処理はモデルを垂直方向に分割します。3D 並列処理は、一般的なライン並列処理、テンソル並列処理、およびデータ並列処理をモデルのトレーニングに同時に適用します。

パイプラインパラレル

パイプライン並列処理の目標は、より大きなモデルをトレーニングすることです。このセクションでは、まず直感的な単純層並列手法を紹介し、その制限を分析します。次に、パイプライン並列アルゴリズム GPipe および PipeDream が導入されます。

1. 単純な層の並列処理

モデルが大きすぎて単一の GPU でトレーニングできない場合、最も簡単なアイデアは、モデルレイヤーを分割し、分割された部分を異なる GPU に配置することです。以下では、単純な層の並列処理を導入する例として 4 層のシーケンスモデルを取り上げます:アニメーション上記の紹介に基づいて、単純な層の並列処理の欠点を見つけることができます。

GPU 使用率が低い。常に 1 つの GPU のみが動作し、他の GPU はアイドル状態になります。
計算と通信には重複がありません。GPU は、順方向伝播 (FWD) または逆方向伝播 (BWD) の中間結果を送信するときにもアイドル状態になります。
メモリ使用量が多い。GPU1 は、最後のパラメーター更新が完了するまで、ミニバッチ全体のすべてのアクティベーションを保存する必要があります。バッチサイズが大きい場合、ビデオメモリに大きな問題が生じます。

GPipe

GPipeの原理

GPipe は、ミニバッチをより小さい同じサイズのマイクロバッチに分割することで効率を向上させます。具体的には、各マイクロバッチに順方向伝播と逆方向伝播を個別に計算させ、各マイクロバッチの勾配を加算してバッチ全体の勾配を取得します。各レイヤーは 1 つの GPU 上にのみ存在するため、マイクロバッチの勾配の合計はローカルでのみ実行する必要があり、通信は必要ありません。

GPU が 4 つあり、モデルがレイヤーごとに 4 つの部分に分割されているとします。単純なレイヤー並列処理のプロセスは次のとおりです。

タイムステップ	0	1	2	3	4	5	6	7
GPU3				前輪駆動	BWD
GPU2			前輪駆動			BWD
GPU1		前輪駆動					BWD
GPU0	前輪駆動							BWD

ご覧のとおり、ある時点では 1 つの GPU のみが動作しています。また、GPU はミニバッチ全体の順伝播を実行する必要があるため、各 timesep には比較的長い時間がかかります。

GPipe はミニバッチを 4 つのマイクロバッチに分割し、順番に GPU0 に送信します。GPU0 での順伝播の後、結果は GPU1 に送信され、以下同様に続きます。全体のプロセスは次のとおりです

タイムステップ	0	1	2	3	4	5	6	7	8	9	10	11	12	13
GPU3				F1	F2	F3	F4	B4	B3	B2	B1
GPU2			F1	F2	F3	F4			B4	B3	B2	B1
GPU1		F1	F2	F3	F4					B4	B3	B2	B1
GPU0	F1	F2	F3	F4							B4	B3	B2	B1

F1 は、現在 GPU 上にある層を使用した microbatch1 の順伝播を意味します。GPipe のスケジューリングでは、各 GPU がマイクロバッチを処理するだけで済むため、各タイムステップにかかる時間は単純な層の並列処理よりも短くなります。

GPipe バブルの問題

バブルとは、パイプライン内で有効な作業が行われていないポイントを指します。これは、操作間の依存関係が原因です。たとえば、GPU4 は、GPU3 が F1 の実行を完了するまでしか待機できません。パイプラインプロセス全体のバブルを次の図に示します。したがって、マイクロバッチ m の数を増やすと、気泡の割合を減らすことができます。

GPipe のグラフィックメモリ要件

バッチサイズを増やすと、キャッシュによってアクティブ化する必要があるメモリ要件が直線的に増加します。GPipe では、GPU は順伝播から逆伝播までのアクティベーションをキャッシュする必要があります。GPU0 を例にとると、microbatch1 のアクティブ化をタイムステップ 0 からタイムステップ 13 まで保存する必要があります。

ビデオメモリの問題を解決するために、GPipe は勾配チェックポイントを使用します。この手法では、すべてのアクティベーションをキャッシュする必要はありませんが、代わりにバックプロパゲーション中にアクティベーションを再計算します。これにより、ビデオメモリの必要性が減りますが、計算コストが増加します。

すべての層はほぼ等しいと想定されます。各 GPU キャッシュのアクティブ化に必要なビデオメモリはPipeDreamです

GPipe は、バックプロパゲーションを開始する前に、すべてのマイクロバッチフォワードプロパゲーションが完了するまで待つ必要があります。PipeDream は、マイクロバッチの前方伝播が完了した後、すぐに逆方向伝播ステージに入ります。理論的には、バックプロパゲーションの完了後に、対応するマイクロバッチキャッシュのアクティブ化を破棄できます。PipeDream のバックプロパゲーションは GPipe よりも早く完了するため、必要なビデオメモリも削減されます。

下の図は、PipeDream、4 つの GPU、8 つのマイクロバッチのスケジューリング図です。青い四角は順伝播を表し、緑色は逆伝播を表し、数字はマイクロバッチの ID です。PipeDream はバブルという点では GPipe と変わりませんが、PipeDream はビデオメモリをより早く解放するため、ビデオメモリの需要が減ります。

データ並列処理とパイプライン並列処理のマージ

データ並列処理とパイプライン並列処理は直交しており、同時に使用できます。

パイプライン並列処理用。各 GPU は、次のパイプラインステージ (順方向伝播) または前のパイプラインステージ (逆方向伝播) と通信する必要があります。
データ並列処理用。各 GPU は、同じレイヤーが割り当てられた GPU と通信する必要があります。すべてのレイヤーのコピーでは、勾配を平均化するために AllReduce が必要です。

これにより、すべての GPU 上にサブグループが形成され、サブグループ内で集合的な通信が使用されます。特定の GPU には 2 つの通信部分があり、1 つはすべて同じレイヤーを含む GPU との通信 (データ並列処理)、もう 1 つは異なるレイヤーを持つ GPU との通信 (パイプライン並列処理) です。以下の図は、パイプライン並列度が 2、データ並列度が 2 の図の例です。水平方向は完全なモデル、垂直方向は同じレイヤーの異なるコピーです。

テンソル並列

Transformer の主なコンポーネントは完全に接続されたレイヤーとアテンションメカニズムであり、その中核となるのは行列の乗算です。テンソル並列処理の核心は、行列の乗算を分割することにより、単一カード上でのモデルのメモリ要件を削減することです。( 1) 行列乗算角度

(2) アクティベーション機能と通信

明らかに、上記のデータ式を観察するだけで、行並列であっても列並列であっても、各部分の計算後に実行する必要があるのは 1 回だけです。列並列処理では通信結果が結合され、行並列処理では通信結果が追加されるだけです。

ここで、非線形活性化 GeLU を追加し、完全にリンクされた 2 つの層をシミュレートします。X を入力、A と B を 2 つの完全に接続された層の重みとします。 したがって、複数のフルリンク層がスタックされている場合、(上の図に示すように) 最終出力中に必要な通信は 1 つだけです 。 (3) 複数の注意を並行して行う

マルチヘッドアテンション並列処理は 1D テンソル並列処理ではありませんが、Megatron-LM で 1D テンソル並列処理と同時に提案されたものですので、ここで簡単に紹介します。マルチヘッドアテンションの各ヘッドは基本的に独立しているため、各ヘッドは並行して動作できます。

注: テンソル並列処理 (TP) には非常に高速なネットワークが必要なため、複数のノードにわたるテンソル並列処理は推奨されません。実際には、ノードに 4 つの GPU がある場合、最高のテンソル並列処理は 4 です。

2D、2.5D テンソル並列処理

1D テンソル並列処理の後、2D、2.5D、および 3D テンソル並列処理が徐々に提案されました。ここでは、2D および 2.5D テンソル並列処理について簡単に説明します。

2D テンソル並列

1D テンソル並列処理ではアクティベーション (つまり、モデルの中間層の出力結果) が分割されないため、大量のビデオメモリが消費されます。ブロードキャストが実行され、各プロセッサが個別に計算します。最後に、2 番目の行列の結果が得られます。

2 つの行列の結果を加算します。

(2) 2.5D テンソル並列

3Dパラレル

一般に、3D 並列処理は、データ並列処理 (DP)、テンソル並列処理 (TP)、およびパイプライン並列処理 (PP) で構成されます。TP と PP はそれぞれ以前に紹介されましたが、ZeRO-DP はメモリ効率の高いデータ並列戦略です。

以下では、3 つの並列テクノロジを組み合わせて 3D 並列テクノロジを形成する方法について説明します。

3D 並列の例

Node1とNode2という 2 つのノードがあり、各ノードに 8 つの GPU、合計 16 の GPU があるとします。16 個の GPU の番号は Rank0、Rank1、...、Rank15 です。さらに、ユーザーがパイプライン並列度を4に、テンソル並列度を2に設定するとします。

パイプラインの並列処理。パイプライン並列処理では、モデル全体が 4 つの部分 (ここでは sub_model_1 から sub_model_4 と呼ばれます) に分割されます。連続する 4 つの GPU はそれぞれ 1 つの sub_model を担当します。つまり、上図の右上隅では、同じ色の GPU が同じ sub_model を担当します。

テンソル並列処理。テンソル並列処理は、パイプライン並列処理の sub_model のテンソルを分割します。つまり、Rank0 と Rank1 は sub_model_1 の 1 つのコピーを担当し、Rank2 と Rank3 は sub_model_1 の別のコピーを担当し、Rank4 と Rank5 は sub_model_2 を担当し、Rank6 と Rank7 は sub_model_2 の別のコピーを担当する、というようになります。上図の右下隅で、緑色の線は単一のテンソル並列グループを表しており、各テンソル並列グループは共同で特定の sub_model を担当します。

データの並列処理。データ並列処理の目的は、同じモデルパラメーターが同じデータを並行して読み取ることを保証することです。パイプライン並列処理とテンソル並列処理の後、Rank0 と Rank2 は同じモデルパラメーターを担当するため、Rank0 と Rank2 は同じデータ並列グループになります。上図の左上の赤い線はデータ並列グループを表します。ワオソフトアイオット http://143ai.com

3D並列解析

3D 並列処理では、なぜ上記のように GPU を分割する必要があるのでしょうか? まず、モデル並列処理は 3 つの戦略の中で通信オーバーヘッドが最も大きいため、より大きなノード内帯域幅を利用するにはモデル並列グループを 1 つのノードに配置することが望ましいです。第 2 に、パイプラインの並列通信は最小限であるため、パイプラインは異なるノード間でスケジュールされ、通信帯域幅によって制限されません。最後に、テンソル並列処理がノードをまたがない場合、データ並列処理もノードをまたぐ必要はありません。そうでない場合は、データ並列グループもノードをまたぐ必要があります。

パイプライン並列処理とテンソル並列処理により、1 枚のグラフィックスカードのメモリ消費量が削減され、メモリ効率が向上します。ただし、モデルを分割しすぎると通信オーバーヘッドが増加し、計算効率が低下します。ZeRO-DP は、オプティマイザの状態を分割することでメモリ効率を向上させるだけでなく、通信オーバーヘッドを大幅に増加させません。

参考文献

https://siboehm.com/articles/22/pipeline-Parallel-training

https://arxiv.org/pdf/1811.06965.pdf

https://huggingface.co/docs/transformers/v4.18.0/en/Parallelism

https://www.colossalai.org/zh-Hans/docs/features/1D_tensor_Parallel

https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/?utm_source=wechat_session&utm_medium=social&utm_oi=738477582021832704