メモリー効率の高い学習処理をインテル® Gaudi® プロセッサーと DeepSpeed で実行

https://www.intel.com/content/www/us/en/developer/articles/training/memory-efficient-training-on-gaudi-with-deepspeed.html

＜インテル® Gaudi® AI アクセラレーターを活用した最適化＞
 •新たなディープラーニング・モデルの作成も既存コードの移植も数分で完了
 •シンプルな開発と生産性の向上を両立しながら生成 AI のパフォーマンスを実現

詳細情報
https://www.intel.com/content/www/us/en/developer/platform/gaudi/overview.html　

大規模言語モデル (LLM) の学習処理では、コンピューティングや通信の効率とモデル精度を犠牲にすることなく、トレーニングに必要なメモリー要件を抑えることが大きな課題の 1 つとなっています。
https://developer.habana.ai/resources/generative-ai-and-large-language-models/　

Microsoft DeepSpeed は、大規模言語モデルでメモリー効率の高い学習処理を可能にする、広く利用されている ディープラーニング用ソフトウェア・ライブラリーです。DeepSpeed には分散トレーニングにおけるメモリー効率化の手法の 1 つ「ゼロ重複オプティマイザー (ZeRO)」が含まれています。 ZeRO のメモリー効率最適化は複数段階で構成され、インテル® Gaudi® ソフトウェアでは現在 ZeRO-1 と ZeRO-2 をサポートしています。

https://habana.ai/blog/release-of-synapseai-version-1-6-0/

このブログでは ZeRO について、LLM のトレーニングにどれほど有効かを、ZeRO テクノロジーの概要とあわせて、ZeRO-1 と ZeRO-2 のメモリー効率最適化の段階ごとに説明します。それでは、LLM にメモリー効率の高い学習処理が不可欠とされる理由と、ZeRO がこの課題を解決する方法について、深く掘り下げていきましょう。

【LLM の登場】
LLM は巨大化を続け、図 1 からも分かるとおり、モデル規模はわずか数年の間に 10 倍にも拡大しました。モデル規模が大きくなるとともに、モデルの精度も大幅に向上します。GPT-2* (15 億パラメーター)、Megatron-LM (83 億パラメーター)、T5 (110 億パラメーター)、Turing-NLG (170 億パラメーター)、Chinchilla (700 億パラメーター)、GPT-3* (1,750 億パラメーター)、OPT-1 (750 億パラメーター)、BLOOM (1,760 億パラメーター) といった大規模言語モデルが続々とリリースされ、その用途は自然言語の理解、Q&A、要約、翻訳、自然言語生成と多種多様です。LLM の規模は拡大の一途をたどっていますが、こうした大規模モデルを効率的にトレーニングするにはどのようにすればいいでしょうか？ その答えは「並列化」です。

![1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3490396/e7f319d5-6505-4602-b5ee-0b592df69d94.png)
図 1. 拡大を続けるモデル規模

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

モデル並列処理では、データを分割する代わりに、モデルを Nm 個のパーツに分割します。通常、Nm はデバイス数です。モデルの分割は垂直方向または水平方向で実行します。層内モデル並列は、各モデル層を複数デバイスにわたって分割する垂直分割。パイプライン並列とも呼ばれる層間モデル並列は、1つまたは複数の層境界でモデルを分割する水平分割です。モデルの並列処理には相当な通信オーバーヘッドが発生するため、単一ノード内では効果的ですが、ノード間では通信オーバーヘッドによるスケーリングの問題が発生します。また、実装が複雑です。

データ並列トレーニングのほうがシンプルで実装しやすいうえに、演算と通信の効率がよいものの、モデル全体のコピーを各デバイスに保持するため、1 デバイスのメモリー・フットプリントを縮小できません。モデル規模が拡大を続けると、モデル全体を 1 つのデバイスメモリーに保持するには限界があります。多くの場合、パラメーター数が 10 億を超えてくると、広帯域幅メモリー (HBM) がわずか 32GB の初代インテル® Gaudi® プロセッサーでは、メモリーを使い切ってしまう可能性があります。HBM 96GB の第 2 世代インテル® Gaudi® 2 プロセッサーは、メモリーをフルに使ってはるかに規模の大きいモデルに適応できますが、モデル規模が指数関数的に拡大し続け 5,000 億パラメーターを超えるようになると、メモリー・フットプリントが学習処理の主なボトルネックになりかねません。では、どうすればこのメモリー・フットプリントを縮小して、大規模モデルのデータ並列トレーニングを実行できるでしょうか？ この問題を解決するには、ディープラーニング用ハードウェア上で実行する大規模モデルの学習処理でメモリー・オーバーヘッドを発生させる主な要因を知る必要があります。まずは、このポイントに着目してみましょう。

【いったいメモリーはどこに行ってしまったのか？】
メモリーを最も消費するのはモデルステートです。モデルステートには、オプティマイザーのステート、勾配、パラメーターを表す Tensor が含まれます。次にメモリー消費が大きいのが、アクティベーション、一時バッファー、使用できない断片化したメモリーなどの常駐ステートです。学習処理中のメモリー消費に関しては、次のように、主に 2 つの見解があります。
　• 多くの場合、モデルステートは学習処理中に最大量のメモリーを消費し、データ並列トレーニングではデータ並列プロセス全体を通してすべてのモデルステートを複製するため、これが過剰なメモリー消費につながっている。
　• 学習処理中、モデルステートの全部が常に必要というわけではないとしても、学習処理プロセス全体を通して必要なモデルステートはすべて静的に保持する。
これらの見解から、モデルステートの複製による過剰なメモリー消費を抑えて、反復学習の段階ごとに必要な値のみを保持するにはどうすればよいかを検討することにしました。

ところで、モデルステートはどうしてこれほど大量のメモリーを必要とするのでしょうか？ その原因の 1 つが、オプティマイザー・ステートです。ディープラーニングの学習処理に広く用いられている Adam オプティマイザーを見てみましょう。Adam は、標準的な確率的勾配降下アルゴリズムの拡張です。Adam オプティマイザーでは、勾配の 1 次モーメントの平均 (平均値) と 2 次モーメントの平均 (非中心分散) に基づいて学習率を適用します。そのため、1 つのパラメーターにつき、時間平均モーメンタムと勾配の分散という 2 つのオプティマイザー・ステートを記憶して、パラメーターの更新を計算する必要があります。混合精度トレーニングを実装した場合、メモリー要件はさらに厳しくなりますが、これについては次のセクションで説明することにしましょう。

大規模モデル学習の最新アプローチは、パラメーターとアクティベーションを bfloat16 で保持する、混合精度 (bfloat16 と FP32) トレーニングです。混合精度トレーニング中、BF16 重みとアクティベーションを用いて、前方伝播と後方伝播の両方を実行します。ただし、効率的に演算を実行して後方伝播の最後に更新内容を適用するために、混合精度オプティマイザーは 全オプティマイザー・ステートの FP32 コピーを保持しなければならず、これにはパラメーターとほかのオプティマイザー・ステートの FP32 コピーも含まれます。オプティマイザーの更新が完了すると、このオプティマイザーの FP32 パラメーター・コピーから、BF16 モデルのパラメーターが更新されます。
インテル® Gaudi® 2 プロセッサーと初代のインテル® Gaudi® プロセッサーどちらも、BF16 データ型をサポートしています。2 つのデータ並列ユニットによるシンプルなデータ並列トレーニングのセットアップ構成を表している図 2 のとおり、以下の要素を保持できるメモリー容量が必要です。

https://habana.ai/products/gaudi2/
 • BF16 パラメーター
 • BF16 勾配
 • FP32 オプティマイザー・ステート (勾配の FP32 モーメンタム、勾配の FP32 分散、FP32 パラメーターを含む)

![2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3490396/00abd5c2-bda6-4207-a961-71d5a5f3095a.png)
図 2. 2 ユニットのデータ並列トレーニング

例えば、N 個のパラメーターがあるモデルを、Adam オプティマイザーを使用し混合精度で学習処理する場合を考えてみましょう。bfloat16 パラメーターに 2N バイト、bfloat16 勾配に 2N バイト必要になります。オプティマイザー・ステートには、FP32 パラメーターに 4N バイト、勾配分散に 4N バイト、勾配モーメンタムに 4N バイトと、合計で 12N バイトのメモリーが必要です。ここでは Adam オプティマイザーの例を挙げましたが、別のオプティマイザーを使用すれば必要なメモリー容量も違ってきます。オプティマイザー・ステートに関係するメモリー要件を定式化するために、N をモデルのパラメーター数、K をオプティマイザー固有の定数とし、KN と表しました。Adam オプティマイザーの場合、K = 12 です。

10 億パラメーターのモデルを想定すると、オプティマイザー・ステートは最終的に 12GB 消費することになり、bfloat16 モデル・パラメーターの重みを保持するため必要な 2GB と比べると、膨大なメモリー容量と言えます。標準的なデータ並列トレーニングでも、データ並列プロセスごとに、モデルの重み、勾配、オプティマイザー・ステートを含む、完全なモデルのコピーを保持する必要があるということを思い出してください。つまり、各データ並列プロセスでオプティマイザー・ステートのコピーを重複して持つことになり、それぞれが 10 億パラメーターに対して 12GB のメモリーを占有することになります。

このように、オプティマイザー・ステートが過剰な冗長メモリーを占有していることが分かりましたが、最適化によってこのオーバーヘッドを解消するにはどうすればいいでしょうか？ ここで出番となるのが DeepSpeed ZeRO です。次のセクションで詳しく説明します。

【ZeRO によるメモリー効率最適化】
標準的なデータ並列トレーニングの場合、各データ並列プロセスでモデルステートによりメモリーを過剰に占有することは、ここまで説明してきたとおりです。ZeRO では、データ並列プロセス全体でモデルステートを分割することで、このメモリー消費を削減します。ZeRO は、モデルステートの分割レベルによって、複数ステージのメモリー最適化で構成されます。ZeRO-1 は、オプティマイザー・ステートのみを分割。ZeRO-2 は、オプティマイザー・ステートと勾配の両方を分割。ZeRO-3 は、3 つのモデルステートすべてを分割、つまりデータ並列プロセス全体を通してオプティマイザー・ステート、勾配、パラメーターのすべてを分割するということです。インテル® Gaudi® ソフトウェアでは DeepSpeed ZeRO-2 ステージまでをサポートしているため、オプティマイザー・ステートと勾配の両方を分割できます。

ZeRO-1 ステージは、データ並列プロセス全体でオプティマイザー・ステートのみの分割です。Nd 回のデータ並列トレーニング構成を想定し、オプティマイザー・ステートを Nd 等分のグループに分けて、i 番目のデータ並列プロセスは i 番目の分割に対応するオプティマイザー・ステートのみを更新するようにします。これで各データ並列プロセスが Nd 分の 1 の容量のみ保持 / 更新するだけで済むようになります。各学習ステップの最後に all-gather 操作を 1 回実行すれば、データ並列プロセス全体から情報を収集して、全プロセスで更新されたパラメーターすべてを受け取ることができます。各データ並列プロセスごとに全オプティマイザー・ステートの Nd 分の 1 を保持 / 更新すればよいので、全パラメーターの Nd 分の 1 のみ更新、したがって各データ並列プロセスのオプティマイザー・ステートに必要なメモリー要件が元の要件の Nd 分の 1 に抑えられるという仕組みです。標準的なデータ並列トレーニング構成でのメモリー要件が 4N + KN であるのに対して (K = オプティマイザー用のメモリー定数、N = モデルのパラメーター数)、ZeRO ステージ 1 ではこのメモリー要件が 2N + 2N + (KN)/Nd. に抑えられ、かなりの容量となる Nd と比べて、全体で 4N 近くまで削減できます。10 億パラメーターのモデルで K=12 の Adam オプティマイザーを考えると、標準的なデータ並列プロセスではモデルステートに 16GB のメモリーを消費するのに対し、ZeRO ステージ 1 では 4GB に抑えられ、その結果、メモリー・オーバーヘッドは 4 分の 1 に削減できるということです。

ZeRO-2 ステージは、データ並列プロセス全体でオプティマイザー・ステートと勾配の両方を分割します。各データ並列プロセスでは、分割ごとに対応するパラメーターだけを更新するという事実を思い出してください。各データ並列プロセスは、担当する分割の勾配だけを必要とし、すべての勾配が必要なわけではありません。ほかレイヤーの勾配は後方伝播で使用できるようになるため、担当するパラメーター分割のデータ並列プロセスでのみ、勾配を削減します。標準的なデータ並列プロセスでは all-reduce 操作となりますが、こちらの場合はデータ並列プロセスごとの reduce 操作です。この reduce 操作が終われば、勾配はもう必要なくなるため、メモリーを解放できます。これで勾配に必要なメモリーは標準的データ並列プロセスの 2N から、ZeRO ステージ 2 では 2N/Nd まで抑えられます。

オプティマイザー・ステートの分割のみを適用すると、パラメーター数 N 個のモデルでのメモリー要件は 2N (パラメーター数) + 2N (勾配) + KN/Nd (オプティマイザー・ステート) でした。オプティマイザーと勾配の両方を分割する ZeRO ステージ 2 では 2N + 2N/Nd + KN/Nd に削減され、Adam オプティマイザーの場合は 2N + 14N/Nd です。Nd, が大きいほど、2N 近くまで削減できます。Adam オプティマイザーで 10 億パラメーターのモデルをトレーニングする場合、標準的なデータ並列プロセスではモデルステートに 16GB のメモリーを消費するのに対し、ZeRO ステージ 2 では 2GB と、8 分の 1 に抑えることが可能です。

【インテル® Gaudi® ソフトウェアの DeepSpeed ZeRO-2 サポート状況】
メモリー効率最適化のサポートにより、インテル® Gaudi® プラットフォームに実装する大規模モデルの効率的な学習処理が可能になります。DeepSpeed ZeRO-2 は、インテル® Gaudi® ソフトウェア Release 1.6 で利用可能です。インテル® Gaudi® プロセッサー対応のアップデートが含まれたインテル® Gaudi® ソフトウェア DeepSpeed ライブラリーを使用する必要があります。詳細については、DeepSpeed User Guide を参照してください。インテル® Gaudi® ソフトウェア GitHub* ページでは、15 億 / 50 億パラメーターの BERT モデルを DeepSpeed ZeRO-1 と ZeRO-2 によりインテル® Gaudi® プロセッサー上でトレーニングする実装サンプルを公開しています。
https://github.com/HabanaAI/deepspeed
https://docs.habana.ai/en/latest/PyTorch/DeepSpeed/index.html

________________________________________
参考資料
•ZeRO & DeepSpeed: New system optimizations enable training models with over 100 billion parameters
https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/

•ZeRO-2 & DeepSpeed: Shattering barriers of deep learning speed & scale
https://www.microsoft.com/en-us/research/blog/zero-2-deepspeed-shattering-barriers-of-deep-learning-speed-scale/

•ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

•DeepSpeed User Guide
https://docs.habana.ai/en/latest/PyTorch/DeepSpeed/index.html

•Use DeepSpeed and Megatron to Train Megatron-Turing NLG 530B
https://www.microsoft.com/en-us/research/blog/using-deepspeed-and-megatron-to-train-megatron-turing-nlg-530b-the-worlds-largest-and-most-powerful-generative-language-model/


【製品およびパフォーマンスに関する情報】
性能は、使用状況、構成、その他の要因によって異なります。詳細については、https://www.Intel.com/PerformanceIndex/ (英語) を参照してください。

関連情報：
インテル® Gaudi® 3 アクセラレーター
https://www.intel.co.jp/content/www/jp/ja/products/details/processors/ai-accelerators/gaudi.html