LLMの技術：並列単語埋め込み（VocabParallelEmbedding）

0 概要

大規模言語モデル（LLM）の指数関数的なスケール拡大に伴い、高度な分散コンピューティングパラダイムの利用が不可欠となっています。パイプライン並列（PP: Pipeline Parallelism）は、モデルの層を複数のプロセッシングユニットに分散させるために広く採用されている技術です。しかし、標準的なPPの実装において、単語埋め込み層と最終出力層という語彙関連層に起因する計算およびメモリ負荷の不均衡は、重大でありながら見過ごされがちな問題です。本稿では、この不均衡を解消するために設計された専門技術であるVocabParallelEmbedding（語彙並列埋め込み）の技術的概要を解説します。パイプライン内の全デバイスに語彙層を分割配置することで、本手法は作業負荷とメモリのより均等な分散を実現し、パイプラインバブルを効果的に削減、ピークメモリ消費量を低減させ、訓練スループットを大幅に向上させます。

1 LLM訓練における並列化技術序論

数百億から数兆のパラメータで構成される現代のLLMの訓練は、単一のアクセラレータでは計算上不可能です。そのため、GPUクラスタなどの専用ハードウェアを用いた分散訓練が標準的な手法となっています。この巨大な計算・メモリ要求を管理するために、いくつかの主要な並列化戦略が、しばしば組み合わせて用いられます。

• データ並列（DP: Data Parallelism）: 各デバイスがモデル全体のレプリカを保持し、それぞれがデータバッチの異なるサブセットを処理します。実装は容易ですが、各デバイスがフルモデルを格納するのに十分なメモリを持つ必要があります。

• テンソル並列（TP: Tensor Parallelism）: 個々のモデルパラメータ（テンソル）を複数のデバイスに分割します。これにより、単一デバイスに収まらない巨大な層の訓練が可能になりますが、頻繁な同期操作（例: all-reduce）による大きな通信オーバーヘッドを伴います。

• パイプライン並列（PP: Pipeline Parallelism）: モデルの層を複数のステージに分割し、各ステージを異なるデバイスに割り当てます。データはこれらのステージをパイプラインのように流れていきます。PPはTPと比較して演算強度が高く、通信コストが低いという特徴があります。これは、通信がステージ間の境界でのみ発生するためです。

PPはその利点にもかかわらず、「パイプラインバブル」として知られる非効率性の影響を受けやすいです。これは、デバイスが依存関係の完了を待っている間のアイドル時間です。パイプラインバブルを悪化させる主な原因は、パイプラインステージ間での作業負荷の不均等な分散です。

2 パイプライン並列における語彙のボトルネック

典型的なTransformerベースのLLMアーキテクチャでは、初期のトークン埋め込み層と、隠れ状態を語彙ロジットに射影する最終的な言語モデルヘッドが、最もメモリを消費するコンポーネントです。この傾向は、多様な言語やドメインに対応するために語彙サイズが拡大するにつれて顕著になります。標準的なパイプライン並列の構成では、最初のステージ（デバイス0）が入力埋め込み層を管理し、最終ステージ（デバイスN-1）が出力射影層を処理します。

この従来型の配置は、顕著な不均衡を生み出します。

• メモリの不均衡: 最初と最後のデバイスは、巨大な単語埋め込み行列全体を保持する必要があり、中間ステージと比較してピークメモリ使用量が著しく高くなります。

• 計算の不均衡: これら埋め込み層に関連する計算は、中間にあるTransformerブロックの計算よりも負荷が高くなる可能性があり、最初と最後のステージが計算上のボトルネックとなります。

この不均衡は、他のステージをアイドル状態に追い込み、パイプラインバブルを拡大させ、全体的な訓練効率を低下させます。語彙サイズが大きくなるにつれて、この問題はますます深刻になります。

3 方法論：Vocabulary Parallelism

前述の不均衡に対処するため、ターゲットを絞った最適化としてVocabulary Parallelismが導入されました。この技術の基本原則は、語彙関連層を語彙次元に沿って分割し、これらの分割片をパイプラインに参加するすべてのデバイスに均等に分散させることです。

語彙サイズがV、隠れ次元がHの埋め込み層を考えます。対応する重み行列のサイズはV x Hです。パイプラインにN個のデバイスがある場合、VocabParallelEmbeddingはこの行列をN個のシャードに分割し、各シャードのサイズは(V/N) x Hとなります。これにより、パイプライン内の各デバイスは、割り当てられた語彙のシャードのみを保持し、計算する責任を負います。

このアプローチにより、語彙のメモリフットプリントと計算負荷が均等に分散され、パイプラインが不均衡な状態から均衡した状態へと変換されます。

この分割スキームを訓練プロセスに統合するため、計算は明確な「語彙パス」に構造化され、既存のパイプラインスケジュール（例: 1F1B - 1 forward, 1 backward pass）にシームレスに挿入されます。これにより、パイプライン並列の基本構造を維持しつつ、作業負荷の再分散が行われます。語彙の分散した性質を計算中に管理するため、特殊な通信アルゴリズムが採用され、同期バリアを削減することで、追加のアクティベーションメモリのオーバーヘッドを最小限に抑えます。

4 実行メカニズム

前述の語彙分割を実現するためには、モデルのフォワードパスおよびバックワードパスにおいて、一連の精密な操作を実行する必要があります。その核心は、各デバイスが担当する語彙部分のみを処理し、効率的な通信を通じてグローバルな結果を統合することにあります。以下に、フォワードパスにおける主要なステップを解説します。

1. マスクの生成 (Mask Generation):
   • まず、入力されたミニバッチのグローバルな語彙インデックスに対し、各デバイスはどのインデックスが自身の管轄範囲に属するかを識別する必要があります。
   • 各デバイスは、自身が担当する語彙インデックスの範囲（例：デバイス i は [i * (V/N), (i+1) * (V/N) - 1] の範囲を担当）をあらかじめ把握しています。
   • 入力されたグローバルインデックスを自身の担当範囲と比較することで、ブール型のマスクを生成します。マスクの形状は入力テンソルと同じで、値が True の位置が、そのインデックスが現在のデバイスにとっての「ローカルな語彙」であることを示します。
2. ローカルインデックスへの変換 (Conversion to Local Indices):
   • ローカルの埋め込み層行列（サイズ (V/N) x H）が期待するインデックス範囲は [0, V/N - 1] であり、グローバルな [0, V - 1] ではありません。
   • そのため、マスクで True とマークされたグローバルインデックスをローカルインデックスに変換する必要があります。これは通常、現在のデバイスが担当する語彙パーティションの開始オフセットを減算することで実現します（すなわち local_index = global_index - partition_start_offset）。
   • 非ローカルなインデックス（マスクの値が False）は、後続のステップで計算結果がクリアされるため、一時的なローカルインデックス（例: 0）にマッピングすることができます。
3. ローカルEmbeddingルックアップ (Local Embedding Lookup):
   • 前のステップで生成されたローカルインデックスを使用して、現在のデバイスに格納されている語彙パーティション行列（(V/N) x H）上で、標準的な埋め込みルックアップ操作を実行します。
   • この操作により、入力バッチ内のすべての語彙に対して埋め込みベクトルが生成されます。
4. 無関係な出力のゼロ化 (Zeroing Out Irrelevant Outputs):
   • 前ステップのルックアップ結果には、非ローカルな語彙（一時的にローカルインデックスにマッピングされたものも含む）に対する埋め込みベクトルが含まれていますが、これらは無効です。
   • ステップ1で生成したマスクを利用し、ローカルのルックアップ結果とマスク（埋め込みベクトルに合うように次元を拡張したもの）を要素ごとに乗算します。これにより、すべての非ローカルな語彙に対応する埋め込みベクトルがゼロに設定されます。
   • この時点で、各デバイスは「部分的」な埋め込みテンソルを保持しており、その中では自身が担当する語彙部分のみが非ゼロの埋め込みベクトルを持っています。
5. グローバル結果の集約 (Aggregation of Global Results):
   • 最後に、パイプライン並列に参加しているすべてのデバイス間で all-reduce 通信操作を実行します。
   • 各デバイスは自身の担当部分の有効な埋め込みのみを計算し、他の部分はすべてゼロであるため、全デバイスの「部分的」なテンソルを加算（all-reduce の中核操作）することで、各デバイス上で完全かつグローバルに一貫した埋め込み結果を正確に再構築することができます。

5 パフォーマンスへの影響と実証分析

Vocabulary Parallelismの実装は、大幅なパフォーマンス向上をもたらします。広範な評価により、この手法が語彙サイズに関わらず計算とメモリを効果的に均衡させることが実証されています。主な成果は以下の通りです。

• スループットの向上: 作業負荷の不均衡を緩和し、パイプラインバブルを削減することで、VocabParallelEmbeddingは単純なパイプライン並列アプローチと比較して、訓練スループットを5%から51%向上させることが示されています。
• ピークメモリ使用量の削減: 埋め込み行列を全デバイスに分散させることで、システムのメモリボトルネックになりがちな最初と最後のステージにおけるピークメモリ消費量が大幅に削減されます。
• スケーラビリティの向上: この技術により、特定のデバイスでメモリや計算のボトルネックを生じさせることなく、はるかに大きな語彙を持つモデルの効率的な訓練が可能になります。

6 テンソル並列との違い

Vocabulary Parallelismとテンソル並列はどちらも埋め込み行列の分割に使用できますが、その目的と運用コンテキストは異なります。

• テンソル並列は、通常、埋め込み行列を隠れ次元（H）に沿って分割します。これはあらゆる巨大なテンソルを分割するためのより一般的な技術であり、その単一層の前進・後進パス内で頻繁かつ高帯域な通信（例: all-reduce）を必要とします。その主な目的は、単一GPUに収まらない一つの演算子を複数のGPUに収めることです。
• 対照的に、Vocabulary Parallelismは語彙次元（V）に沿って分割し、特にパイプラインのステージを均衡させるためのメタ戦略として設計されています。単一層のためにデバイスのサブグループに負荷を分散するのではなく、パイプライン全体にわたって作業負荷を分散させます。通信はパイプラインスケジュール自体に統合され、ステージ間のアイドル時間を削減することを目的としています。

要するに、テンソル並列は単一の層のサイズの問題に対処し、Vocabulary Parallelismはパイプライン内の複数層にわたる作業負荷の分散の問題に対処します。

7 結論

VocabParallelEmbeddingは、大規模言語モデルの分散訓練を最適化するための、極めて重要かつ効果的な技術です。巨大な語彙を持つモデルにパイプライン並列を適用する際に内在する計算上およびメモリ上の不均衡を特定し、体系的に解決することで、より均衡の取れたハードウェア利用を可能にします。その結果として得られるスループットの向上とピークメモリ消費量の削減は、このような専門的な並列化戦略が、モデルのスケールと訓練効率の限界を押し上げる上で不可欠であることを示しています。