ボキャブラリー並列化によるパイプライン並列化のバランス調整

核心概念

Transformerベースの大規模言語モデルの学習において、ボキャブラリー層に起因するパイプラインステージ間の計算およびメモリ使用量の不均衡を解消することで、スループットを向上させ、ピークメモリ使用量を削減できる。

要約

ボキャブラリー並列化によるパイプライン並列化のバランス調整

本論文は、Transformerベースの大規模言語モデル（LLM）の学習における、パイプライン並列化におけるボキャブラリー層に起因する非効率性に対処する新しい手法を提案しています。

パイプライン並列化における課題

LLMの学習には、モデル並列化とパイプライン並列化（PP）という2つの主要な並列化手法が用いられます。PPは、モデルを複数のステージに分割し、各ステージを異なるデバイスに配置することで、大規模なモデルの学習を可能にします。しかし、PPには、パイプラインバブルと高いメモリ消費という2つの課題があります。

• パイプラインバブル: パイプラインステージ間での計算の依存関係により、一部のステージがアイドル状態になり、計算リソースの利用効率が低下することがあります。
• 高いメモリ消費: PPでは、パイプラインバブルを削減するために、複数のマイクロバッチのアクティベーションをメモリに保持する必要があるため、メモリがボトルネックになることがよくあります。

ボキャブラリー層の不均衡問題

従来のPPの実装では、Transformer層はパイプラインステージ間で均等に分散されますが、入力層と出力層（ボキャブラリー層）はそれぞれ最初のステージと最後のステージに配置されます。この不均衡な構成は、以下の2つの理由から、計算とメモリの両方においてパフォーマンスを大幅に低下させます。

1. 計算の不均衡: ボキャブラリー層の計算コストは、Transformer層よりも大幅に高くなる可能性があります。その結果、ボキャブラリー層が配置されたステージの計算時間が長くなり、他のステージでパイプラインバブルが発生します。
2. メモリの不均衡: 入力層は、モデルパラメータとアクティベーションの両方を大量に消費するため、最初のステージでメモリボトルネックが発生しやすくなります。

ボキャブラリー並列化

本論文では、PPにおける計算とメモリの不均衡問題に対処するために、ボキャブラリー並列化という新しい手法を提案しています。この手法では、ボキャブラリー層をボキャブラリー次元で分割し、すべてのパイプラインデバイスに均等に分散します。これにより、ボキャブラリー層の計算とパラメータメモリが効果的に分散され、パイプラインステージ間のバランスが向上します。

ボキャブラリーパスの構築とスケジューリング

ボキャブラリー並列化を実現するために、ボキャブラリー層の計算を複数のパイプラインパスに分割し、既存のパイプラインスケジュールに統合する手法を提案しています。具体的には、以下の3つのステップで実現します。

1. ボキャブラリーパスの構築: ボキャブラリー層の計算を、前方パス、後方パス、および通信バリアを含む複数のパスに分割します。
2. 通信バリアの最適化: ソフトマックス計算で必要となるAll-reduce通信などの通信バリアの数を削減するアルゴリズムを提案しています。
3. パイプラインスケジューリング: ボキャブラリーパスを既存のパイプラインスケジュール（1F1B、V-Halfなど）に統合する手法を提案しています。

評価

提案手法を評価するために、最大32個のNVIDIA A100 GPUを用いて、様々なモデルサイズとボキャブラリーサイズのGPT系モデルの事前学習実験を行いました。実験の結果、提案手法は、従来手法と比較して、スループットを最大51%向上させ、ピークメモリ使用量を削減できることが確認できました。

まとめ

本論文では、LLMの学習におけるパイプライン並列化の効率性を向上させるボキャブラリー並列化という新しい手法を提案しました。提案手法は、ボキャブラリー層の計算とメモリを効果的に分散することで、パイプラインステージ間のバランスを向上させ、スループットの向上とピークメモリ使用量の削減を実現します。

統計

Gemma2 9B (ボキャブラリーサイズ256k) では、出力層の計算量とパラメータメモリはTransformer層の約5倍である。
10Bモデル（シーケンス長2048）の場合、ボキャブラリーサイズを64kから128kに増やすと、MFUが9.7%低下する。
21Bモデルの場合、Vocabulary Parallelismは、MFUにおいて、インターレースドパイプラインよりも6.7%から8.2%優れている。
ベースラインと比較して、Vocab-1は、MFUにおいて、7.2%から143%の間で常に優れている。
ベースラインでは、パイプラインデバイス間で最大45GBのメモリ 불균형 が発生する。

引用

抽出されたキーインサイト

Balancing Pipeline Parallelism with Vocabulary Parallelism

by Man Tsung Ye... 場所 arxiv.org 11-11-2024

https://arxiv.org/pdf/2411.05288.pdf

深掘り質問

ボキャブラリー並列化は、他のモデル並列化手法（テンソル並列化、ZeROなど）とどのように統合できるでしょうか？

ボキャブラリー並列化は、テンソル並列化やZeROなどの他のモデル並列化手法と統合することで、さらに大規模な言語モデルの学習を効率化できます。

テンソル並列化との統合: ボキャブラリー並列化は、主にボキャブラリー層の重み行列を分割しますが、テンソル並列化は、他のすべての層の重み行列やテンソル演算を分割します。これらを組み合わせることで、巨大なボキャブラリーとモデルサイズを持つLLMを、複数のデバイスに効率的に分散できます。具体的には、各デバイスは、ボキャブラリーの一部と、テンソル並列化によって分割されたモデルパラメータの一部を保持します。これにより、メモリ使用量と計算をさらに削減できます。

ZeROとの統合: ZeROは、データ並列化、モデル並列化、パイプライン並列化などの並列化戦略と組み合わせて使用できるメモリ最適化技術です。ボキャブラリー並列化とZeROを組み合わせることで、メモリ使用量をさらに削減できます。例えば、ZeRO-3はモデルの重みをすべてのデバイスに分散しますが、ボキャブラリー並列化と組み合わせることで、各デバイスはボキャブラリーの一部と対応する重みのみを保持できます。
これらの統合は、大規模なモデルのメモリフットプリントと通信オーバーヘッドを削減し、学習の高速化に貢献します。ただし、最適なパフォーマンスを得るためには、モデルアーキテクチャ、ハードウェア構成、通信帯域幅などの要素を考慮して、並列化戦略を慎重に調整する必要があります。

ボキャブラリー層の計算のオーバーヘッドをさらに削減するために、どのようなハードウェア/ソフトウェアの最適化が可能でしょうか？

ボキャブラリー層の計算オーバーヘッドをさらに削減するために、以下のようなハードウェア/ソフトウェアの最適化が考えられます。
ソフトウェア最適化:

CUDAカーネルの融合: 論文中にも言及されているように、ボキャブラリー層のForwardパスとBackwardパスを融合したCUDAカーネルを実装することで、softmax結果のメモリへの書き込み/読み込みを回避できます。これは、巨大なボキャブラリーと長いコンテキストを持つ設定では、大幅な高速化につながります。
疎な勾配の活用: 巨大なボキャブラリーを持つLLMでは、多くの場合、勾配は疎です。この疎性を活用したアルゴリズムやデータ構造（Compressed Gradient Descentなど）を用いることで、通信量と計算量を削減できます。
混合精度演算: ボキャブラリー層の一部を低精度演算 (FP16など) で実行することで、計算と通信のオーバーヘッドを削減できます。ただし、精度低下によるモデルの性能への影響を考慮する必要があります。
ハードウェア最適化:

高帯域幅メモリ: HBM2eなどの高帯域幅メモリを使用することで、ボキャブラリー層の重み行列へのアクセスを高速化できます。
計算専用ハードウェア: ボキャブラリー層の計算に特化したASICやFPGAなどのハードウェアを開発することで、大幅な高速化と電力効率の向上が見込めます。
光通信: デバイス間の通信を高速化するために、NVLinkやInfinibandなどの高速インターコネクトだけでなく、将来的には光通信技術の活用も期待されます。
これらの最適化を組み合わせることで、ボキャブラリー層の計算オーバーヘッドを大幅に削減し、より大規模で高性能なLLMの学習が可能になります。

ボキャブラリー並列化は、画像や音声などの他のモダリティを含むマルチモーダルLLMの学習にどのように適用できるでしょうか？

ボキャブラリー並列化は、テキストベースのLLMだけでなく、画像や音声などの他のモダリティを含むマルチモーダルLLMの学習にも適用できます。
マルチモーダルLLMでは、各モダリティごとに巨大な埋め込み層が存在することが一般的です。例えば、画像を扱う場合、画像をトークン化し、各トークンを埋め込みベクトルに変換します。この埋め込み層は、テキストのボキャブラリー層と同様に巨大になる可能性があり、ボトルネックとなる可能性があります。
ボキャブラリー並列化を適用する場合、画像や音声の埋め込み層も、テキストのボキャブラリー層と同様に、複数のデバイスに分割できます。これにより、メモリ使用量と計算量を削減し、学習を高速化できます。
具体的には、以下のような方法が考えられます。

画像埋め込み: 画像をパッチに分割し、各パッチをトークンとして扱うVision Transformer (ViT)のようなモデルでは、パッチ埋め込み層をボキャブラリー並列化できます。各デバイスは、パッチ埋め込み行列の一部を保持し、対応するパッチの埋め込みベクトルを計算します。
音声埋め込み: 音声認識などで用いられる音声埋め込み層も同様に分割できます。音声信号をフレームに分割し、各フレームの特徴量を抽出する埋め込み層を、複数のデバイスに分散して計算します。
ただし、マルチモーダルLLMの学習では、各モダリティのデータ特性やモデルアーキテクチャが異なるため、最適な並列化戦略はテキストベースのLLMとは異なる場合があります。各モダリティの特性を考慮し、ボキャブラリー並列化と他のモデル並列化手法を適切に組み合わせることで、効率的な学習を実現できます。