インテル® Gaudi® 2 アクセラレーター上で 3D 並列処理を活用した 130 億パラメーター LLM 「Llama」と「Bloom」のトレーニング

大規模言語モデル (LLM) の学習処理における重要な課題として、多くのケースで単一ノードに収めるには大規模すぎることや、仮に適合したとしても学習処理に時間がかかる点が挙げられます。この問題を解決するには、複数のインテル® Gaudi® アクセラレーター (HPU) 間で学習処理を並列化します。

これは、データかモデルまたはその両方を並列化すれば、複数のデバイスにわたって分散できるということです。並列化にもさまざまな形態があり、いくつかを組み合わせることで、LLM の効率的なトレーニングが可能になります。DeepSpeed[2] はディープラーニング用のソフトウェア・ライブラリーとして広く利用されており、大規模言語モデルの演算処理を効率化して、メモリー効率の高い学習処理を可能にします。Megatron-LM は、モデル並列化によって大規模言語モデルの学習処理を実行する、高度に最適化された高効率のライブラリーです。

[12]DeepSpeed の学習処理エンジンは、データとパイプラインのハイブリッド並列化を可能にし、これをさらに Megatron-LM のようなモデル並列化と組み合わせることもできます。このように、データ並列、パイプライン並列、Tensor 並列 (水平モデル並列) を組み合わせて、LLM トレーニングの 3D 並列処理を実現します。これらのモデルは、Habana DeepSpeed フレームワーク、Megatron LM、Habana SynapseAI ソフトウェアを使用した、インテル® Gaudi® 2 アクセラレーター上での学習処理が可能です。このブログでは、3D 並列処理とは何か、LLM のトレーニングにどれほど便利なのかを説明したいと思います。まず、インテル® Gaudi® 2 アクセラレーター上での 3D 並列処理による、LLama 13B [10] と BLOOM 13B [11] それぞれの LLM のトレーニング方法について、簡単に概要を紹介します。Habana SynapseAI ソフトウェアでの DeepSpeed のサポートについて、詳しくは Habana DeepSpeed User Guide を参照してください。
https://docs.habana.ai/en/latest/PyTorch/DeepSpeed/DeepSpeed_User_Guide.html

それでは、並列化手法の違いと、これらの並列モードを組み合わせて、DeepSpeed、Megatron LM、Habana SynapseAI ソフトウェアを使用した LLM のトレーニング方法について、詳しく見ていきましょう。

LLM の規模は拡大の一途をたどっていますが、こうした大規模モデルを効率的にトレーニングするにはどのようにすればいいでしょうか？その答えはもちろん「並列化」です。トレーニングの演算とメモリーを効率化するために適用できる並列化にはさまざまな手法があり、並列処理は、複数デバイス間にデータまたはモデルを分割することで成り立ちます。

主に用いられる手法は、大きく分けるとデータ並列とモデル並列の 2 タイプです。データ並列トレーニングの場合、各データ並列プロセス / ワーカーはモデル全体のコピーを保持しますが、データセットは Nd 個のパーツに分割されます。Nd はデータ並列処理のプロセス数です。つまり Nd 個のデバイスがあるとすると、小サイズのバッチを Nd 個のパーツに分割し、それぞれのパーツを 1 つのデバイスに割り当てることになります。次に、モデルを通して対応するバッチを各デバイスに渡し、小バッチの分割ごとに勾配を取得します。その後、すべての勾配を収集し、全体の平均でパラメーターを更新します。すべての勾配を取得したら (ここで同期が必要)、勾配の平均を算出し、その平均を使用してモデルやパラメーターを更新します。更新が終わると、モデル・トレーニングの次のイテレーションに進みます。データ並列性のレベルは分割の数であり、Nd. で表されます。

【Tensor 並列とパイプライン並列】

モデル並列処理では、データを分割する代わりに、モデルを Nm 個のパーツに分割します。通常、Nm はデバイス数です。モデルの並列処理には相当な通信オーバーヘッドが発生するため、単一ノード内では効果的ですが、ノード間では通信オーバーヘッドによるスケーリングの問題が生じる可能性があります。また、実装が複雑です。大規模言語モデル 1 つは複数レイヤーで構成され、各レイヤーが Tensor で動作します。モデル並列処理の場合、各レイヤー内でモデルを分割し、各モデルレイヤーを複数デバイスにわたって分割します。これはレイヤー内モデル並列、または Tensor 並列とも呼ばれる処理です。Tensor 並列では、各レイヤーがデバイス間を横断して並列演算されるように、複数のデバイス間で Tensor をシャードに分割します。

Tensor 並列とは対照的に、レイヤー間モデル並列と同様に、1 つ以上のモデルレイヤー境界でモデルを分割することも可能です。こちらはパイプライン並列とも呼ばれます。パイプライン並列は当初 GPipe 論文で提案されました。[4]パイプライン並列は、直感的に、標準的な工場で見られる組み立てパイプラインの運用と非常によく似た方法です。この場合、モデルは複数のステージ (1 ステージは 1 つ以上のモデルレイヤーで構成) に分割され、各ステージが別々のインテル® Gaudi® アクセラレーター・デバイスに割り当てられます。1 ステージの出力は、次のステージの入力として渡されます。簡単な例として、4 レイヤー構成の多層パーセプトロン (MLP) モデルを考えてみましょう。ここでは 1 レイヤーをそれぞれインテル® Gaudi® アクセラレーター・デバイスに割り当てます。ただし下図に示すとおり、このナイーブ分割では 1 度に 1 つのデバイスしかアクティブにならないため、結果的にデバイス使用率が十分ではなくなります。

図 1. デバイス使用率が十分ではないナイーブ・パイプライン並列

上記の図 1[6] は、4 つの異なるデバイスに配置した (縦軸) 4 レイヤー構成のモデルを表したものです。横軸は、このモデルの 1 度に 1 つのデバイスしか使用できない学習処理を時間経過で示しています。この問題を軽減するために、パイプライン並列によって小バッチの入力を複数のマイクロバッチに分割し、これらのマイクロバッチの実行を複数のデバイス間でパイプライン化します。このパイプライン並列処理を図示したものが、下の図 2 です。

図 2. マイクロバッチのパイプライン並列処理

上記の図 2[6] は、4 つの異なるデバイスに配置した (縦軸) 4 レイヤー構成のモデルを表したものです。横軸は、デバイスがより効率的に使用されている学習処理を時間経過で示しています。ただし、(図からも分かるとおり) 一部のデバイスが活用されない非効率な「バブル」部分は依然として残ったままです。パイプライン並列は理解しやすく実装もシンプルですが、デバイス使用率が十分ではないことから、このような演算上の非効率性が生じかねません。また、前ステージの出力を次ステージに渡す必要があるため、同一ノード内のデバイス間ではなく、ノード間で通信が発生した際、通信オーバーヘッドにつながる可能性があります。
パイプライン並列ではレイヤーの境界でモデルを分割するのに対し、Tensor 並列ではレイヤーごとに Tensor をシャード化することによってレイヤー内でモデルを分割します。Tensor 並列については Megatron[3] の論文で詳述されていますが、実装は極めて複雑です。とはいえ、パイプライン並列と比べて効率性の高い処理と言えます。簡単な例として、 2 レイヤーの MLP モデルを考えてみましょう。レイヤー 1 のパラメーターを行列 A、レイヤー 2 のパラメーターを行列 B とし、それぞれ入力データをバッチ X、MLP の出力を Z で表しました。通常、Y = f(XA) の演算が必要になり、この式で f は活性化関数であり、MLP の出力は Z = YB となります。Tensor 並列の場合、以下の図 3 に示すとおり、両方の行列を 2 つのピースに分割して、これらのピースを 2 つのデバイスに割り当てることができます。ここで行列 A は列方向に 2 等分、行列 B は行方向に 2 等分です。

図 3. Tensor 並列を用いた多層パーセプトロン (MLP) の前方パス

上記の演算は、以下のようにシャード化 Tensor に書き換えることができます。

上記の演算式では Tensor 並列を使用して前方パスを記述していますが、後方パスの演算式も同様に、シャード化 Tensor に書き換えることが可能です。パイプライン並列と Tensor 並列の詳細については、[6] を参照してください。

3D 並列処理とは？
すでに説明したとおり、並列処理には、データ並列、Tensor 並列 (レイヤー内モデル並列)、パイプライン並列 (レイヤー間モデル並列) の 3 タイプがあります。DeepSpeed も Megatron LM も、LLM の学習処理を効率化する、次のような特長を持つライブラリーです。

•メモリー効率: まずパイプライン並列でモデルを分割してから、パイプライン化された各ステージを Tensor 並列でさらに分割。この 2D の組み合わせによって、モデル、オプティマイザー、アクティベーションのメモリー消費を同時に抑えることが可能。ただし、極端なパーティション分割は通信オーバーヘッドにつながり、演算効率に影響する。

•演算効率: Tensor 並列やパイプライン並列のアプローチを超えてスケーリングするには、データ並列を活用。[1]
•トポロジー認識 3D マッピング: 3 タイプの並列処理を入念に組み合わせて、ノード内とノード間どちらの場合も通信オーバーヘッドを最適化。

3D 並列処理の詳細については、[6,7,8] を参照してください。以下の図では、DeepSpeed と Megatron LM の手法を取り入れ、いかに 3D 並列処理を活用するか説明しています。32 ワーカーを 8 ノードにマッピング、各ノードが 4 台のインテル® Gaudi® アクセラレーター・デバイスをホストする例です。ニューラル・ネットワークのレイヤーは、4 つのパイプライン・ステージに分割します。各パイプライン・ステージ内のレイヤーは、モデル並列の 4 つのワーカー間でさらにパーティション分割します。最後に各パイプラインが 2 つのデータ並列インスタンスへ複製します。

図 4. 各ノードに 4 デバイスを実装した 8 ノード構成 DeepSpeed 3D 並列処理の例 ([1] から抜粋)

【インテル® Gaudi® 2 アクセラレーターで動かす 3D 並列処理による Llama-13B のトレーニング】

それでは、3D 並列処理の仕組みを理解できたところで、この特長を活用した LLM のトレーニング方法に注目していきましょう。LLM のサンプルには LLama-13B を使用します。Llama は、Transformer アーキテクチャーをベースとする大規模言語モデルの 1 つです。絶対位置 / 相対位置埋め込みを組み合わせた Rotary Position Embeddings (RoPE)、SwiGLU 活性化関数、事前正規化、AdamW オプティマイザーの使用といったさまざまな拡張機能とあわせて、1 兆を超えるトークンでトレーニングします。Llama モデルの完全な技術仕様については、[10] を参照してください。

ここでは Habana SynapseAI ソフトウェア v1.10.0、PyTorch 2.0.1、DeepSpeed 0.7.7 ソフトウェア・ライブラリーを使用して LLama-13B モデルをトレーニングしました。トレーニング実装は、https://github.com/microsoft/Megatron-DeepSpeed/ (英語) に基づいています。Llama-13B や同等サイズの大規模モデルの具体的なトレーニング手順は、こちらを参照してください。トレーニングの実行コマンドは以下のとおりです。
https://github.com/microsoft/Megatron-DeepSpeed/
https://github.com/HabanaAI/Model-References/tree/master/PyTorch/nlp/DeepSpeedExamples/Megatron-DeepSpeed

HL_HOSTSFILE=scripts/hostsfile HL_NUM_NODES=8 HL_PP=2 HL_TP=4 HL_DP=8
scripts/run_llama13b.sh

HL_NUM_NODES には実行に関与するノード数を指定、HL_PP はパイプライン並列係数の設定、HL_TP 変数は Tensor 並列係数の設定に使用します。HL_DP はデータ並列係数です。HL_HOSTSFILE には、各 HPU ノードの IP アドレスが含まれます。トレーニングの実行スクリプトは、以下のとおり、重要な変数を設定します。

NUM_NODES=${HL_NUM_NODES}

DP=${HL_DP}

TP=${HL_TP}

PP=${HL_PP}

インテル® Gaudi® 2 アクセラレーター x64 デバイスを BF16 精度で実装することで、SynapseAI 1.11 を使用した Llama-13B モデルのトレーニング実行で 1 秒当たり 55.48 センテンスのスループットを達成しました。

【インテル® Gaudi® 2 アクセラレーターで動かす 3D 並列処理による BLOOM-13B のトレーニング】

ここまで LLama-13B モデルのトレーニングを行う 3D 並列処理について説明しましたが、こちらに記載されている Bloom-13B [11] モデルなど、ほかの大規模モデルにも同様の設定を使用することができます。
https://github.com/HabanaAI/Model-References/tree/master/PyTorch/nlp/DeepSpeedExamples/Megatron-DeepSpeed

マルチカード構成の場合、例えば Bloom を 32 基のインテル® Gaudi® 2 アクセラレーター・デバイスを使用し BF16 精度で実行するには、HL_HOSTSFILE に各 HPU ノードの IP アドレスを含む、以下のような実行コマンドになります。

HL_HOSTSFILE=scripts/hostsfile HL_NUM_NODES=4 HL_PP=2 HL_TP=4 HL_DP=4
scripts/run_bloom13b.sh

64 デバイスでトレーニングを実行するために、
このコマンドの HL_NUM_NODES と HL_DP を以下のように変更しています。

HL_HOSTSFILE=scripts/hostsfile HL_NUM_NODES=8 HL_PP=2 HL_TP=4 HL_DP=8
scripts/run_bloom13b.sh

128 デバイス、BF16 精度の場合のスループットは、1 秒当たりのサンプル数 114.1 となりました。これはトレーニング実行を完了し、収束するまでに 15.6 日かかることに相当します。以下の図は、学習処理の反復回数に対するトレーニング損失 (lm-loss) を描いたグラフです。青色の曲線は HPU (bf16 精度) で測定したトレーニング損失、オレンジ色の曲線は FP16 精度の BigScience Bloom リファレンス実装 [11] によるトレーニング損失を示しています。

図 5. 学習処理の反復回数に対するトレーニング損失 (lm-loss)

Megatron-DeepSpeed を使用して大規模モデルをトレーニングしたいけれども、そのモデルが実装に含まれていない場合は、Megatron-DeepSpeed フレームワークに移植することも可能です。独自モデルを Megatron-DeepSpeed フレームワークに移植する詳しい手順については、こちらのドキュメントを参照してください。
https://developer.habana.ai/blog/porting-a-model-to-megatron-deepspeed-with-habana-gaudi/

インテル® Gaudi® プラットフォームを使用して生成 AI の開発を今すぐ開始しましょう。LLM のトレーニングにインテル® Gaudi® 2 アクセラレーターへのアクセスが必要な場合は、こちらの手順に従ってインテル® Tiber™ デベロッパー・クラウドに登録、またインテル® Gaudi® 2 アクセラレーター搭載サーバー・インフラストラクチャーについては、Supermicro までお問い合わせお願いします。
インテル® Gaudi® 2 アクセラレーター、SynapseAI、DeepSpeed を活用し、快適かつスピーディーにモデルの学習処理を行ってください。
https://developer.habana.ai/intel-developer-cloud/
https://www.supermicro.com/en/products/system/ai/8u/sys-820gh-tnr2

参考資料
1.DeepSpeed: Extreme-scale model training for everyone
https://www.microsoft.com/en-us/research/blog/deepspeed-extreme-scale-model-training-for-everyone/

2.PyTorch blog on pipeline parallelism
https://pytorch.org/docs/stable/pipeline.html

3.Megatron LM paper
https://arxiv.org/pdf/1909.08053.pdf

4.Gpipe paper.
https://arxiv.org/abs/1811.06965

5.PipeDream paper
https://arxiv.org/abs/1806.03377

6.Tensor parallelism vs pipeline parallelism

7.https://huggingface.co/docs/transformers/v4.15.0/parallelism/

8.Deepspeed tutorial on pipeline parallelism
https://www.deepspeed.ai/tutorials/pipeline/

9.What Language Model to Train if You Have One Million GPU Hours?
https://intel-my.sharepoint.com/personal/sandya_mannarswamy_intel_com/Documents/What Language Model to Train if You Have One Million GPU Hours?

10.LLaMA:Open and Efficient Foundation Language Models
https://arxiv.org/pdf/2302.13971.pdf

11.BLOOM: A 176B-Parameter Open-Access Multilingual Language Model
https://arxiv.org/abs/2211.05100

12.https://github.com/NVIDIA/Megatron-LM/