テンソル並列処理は、単一のニューラルネットワーク層の計算を複数のハードウェアデバイスに分散する技術です。このアプローチでは、層の重み行列がシャードに分割され、各デバイスがパラメータの一部を保持し、行列乗算の対応する部分を実行します。部分的な結果は、all-reduceなどの集合通信操作を介して結合され、層の出力が生成されます。この方法はモデル並列処理の一形態であり、各デバイスがモデルの完全なコピーを保持し異なるデータサンプルを処理するデータ並列処理とは区別されます。テンソル並列処理は、単一のアクセラレータのメモリ容量を超える大規模言語モデルなどの非常に大きなモデルのトレーニングと推論の実行に不可欠です。
テンソル並列処理の主な動機は、現代の深層学習モデルのメモリと計算の要求です。例えば、数千億のパラメータを持つモデルは、通常16から80ギガバイトの範囲である単一のGPUやTPUのメモリに収まりません。重み行列をシャーディングすることで、テンソル並列処理はモデルを多くのデバイスに分散させ、より大きなモデルの使用とより高速なトレーニング時間を可能にします。また、デバイスごとのメモリフットプリントを削減し、より大きなバッチサイズを可能にします。しかし、テンソル並列処理は通信オーバーヘッドを導入します。デバイスは部分的な結果を同期する必要があり、特に遅いインターコネクトではボトルネックになる可能性があります。
テンソル並列処理は、パイプライン並列処理やデータ並列処理などの他の並列処理戦略と組み合わせて、数千のデバイスにわたるトレーニングをスケーリングするために一般的に使用されます。これは、OpenAI、Anthropic、Google DeepMindによって開発されたものを含む、多くの最先端モデルのトレーニングインフラストラクチャの重要なコンポーネントです。この技術は推論にも関連しており、複数のアクセラレータに計算を分散することで、低レイテンシで大規模モデルをサービス提供できます。
歴史的背景
行列演算を複数のプロセッサに分割する概念は、初期の並列コンピューティング研究に遡ります。1980年代と1990年代には、研究者はニューラルネットワークの並列実装を探求しましたが、ハードウェアとソフトウェアのツールは限られていました。テンソル並列処理の現代的な形態は、2010年代のGPUベースの深層学習の台頭とともに登場しました。最も初期の注目すべき使用例の1つは、大規模な畳み込みネットワークのトレーニングであり、トロント大学や他の機関の研究者がモデル並列処理を実験しました。しかし、テンソル並列処理を主流の技術にしたのは、2017年に導入されたTransformerアーキテクチャなどのトランスフォーマーベースのモデルの登場でした。Transformerのアテンション層とフィードフォワードネットワークは大きな行列乗算で構成されており、シャーディングの自然な候補です。
2019年には、NVIDIAや他の組織の研究者が、大規模言語モデルのトレーニングのためのテンソル並列処理の詳細な説明を発表しました。NVIDIAによって開発されたMegatron-LMフレームワークは、トランスフォーマーモデルの効率的なテンソル並列処理を実証し、複数のGPUでほぼ線形のスケーリングを達成しました。この作業は、GoogleのGShardやMicrosoftのDeepSpeedなど、テンソル並列処理をコア機能として組み込んだその後の大規模トレーニングシステムの基盤を築きました。
テンソル並列処理の仕組み
典型的なニューラルネットワーク層では、フォワードパスは行列乗算を計算します: Y = XW、ここでXは入力アクティベーション、Wは重み行列、Yは出力です。テンソル並列処理では、重み行列Wは1つ以上の次元に沿って分割されます。完全結合層の場合、一般的なアプローチはWを列方向（つまり出力次元に沿って）に分割することです。各デバイスはWの列のサブセットを保持し、部分的な出力を計算します。部分的な出力は、all-reduce操作を使用して結合され、最終的なYが生成されます。これは列並列分割として知られています。あるいは、Wを行方向（入力次元に沿って）に分割することもでき、入力に対するall-gather操作と出力に対するall-reduceが必要です。
トランスフォーマーモデルでは、テンソル並列処理はマルチヘッドアテンションとフィードフォワードネットワークの両方に適用されます。マルチヘッドアテンションでは、クエリ、キー、バリューの重み行列がデバイス間でシャーディングされ、アテンションヘッドが分散されます。出力射影はその後結合されます。フィードフォワードネットワークでは、2つの線形層が通常分割されます: 最初の層は列方向に分割され、2番目の層は行方向に分割されるため、中間アクティベーションがシャーディングされ、出力が削減されます。この設計は、特定のポイントでのみall-reduce操作を必要とすることで通信を最小限に抑えます。
通信パターンは重要です。各all-reduce操作は、テンソル並列グループ内のすべてのデバイス間でデータを交換する必要があります。通信量は、重みではなくアクティベーションのサイズに比例し、これはしばしば小さくなります。しかし、非常に大きなモデルでは、通信は依然として重要になる可能性があります。これを軽減するために、テンソル並列処理は通常、NVIDIAのNVLinkやAMDのInfinity Fabricなどの高速インターコネクトを持つ単一ノード内で使用され、データ並列処理はノード間で使用されます。
他の並列処理技術との比較
テンソル並列処理は、いくつかのモデル並列処理戦略の1つです。例えば、パイプライン並列処理は、モデルを層ごとに分割し、各デバイスが連続した層のセットを担当します。これにより、層境界でのアクティベーションのみが交換されるため、テンソル並列処理と比較して通信が削減されます。しかし、パイプライン並列処理は、パイプラインバブルによるアイドル時間に悩まされる可能性があります。一方、データ並列処理は、各デバイスでモデルを複製し、異なるデータサンプルを処理し、各ステップ後に勾配同期が必要です。テンソル並列処理は両方と補完的です: デバイスごとのメモリを削減するためにパイプライン並列処理と組み合わせ、スループットを向上させるためにデータ並列処理と組み合わせることができます。
もう1つの関連技術はシーケンス並列処理であり、トランスフォーマーモデルのシーケンス次元をシャーディングします。これは、メモリ使用量をさらに削減するために、テンソル並列処理と組み合わせて使用されることがよくあります。エキスパート並列処理は、mixture-of-expertsモデルで使用され、エキスパートモジュールをデバイス間でシャーディングします。これはモデル並列処理の一形態ですが、厳密にはテンソル並列処理ではありません。
アプリケーションと使用例
テンソル並列処理は、大規模言語モデルのトレーニングと推論で広く使用されています。例えば、OpenAIによって開発されたGPT-3モデルは、数千のNVIDIA GPUにわたるテンソル並列処理を含むモデル並列処理の組み合わせを使用してトレーニングされました。同様に、AnthropicのClaudeモデルとGoogle DeepMindのGeminiモデルは、数千億のパラメータにスケーリングするためにテンソル並列処理に依存しています。推論では、テンソル並列処理により、NVIDIAのTriton Inference ServerやHugging FaceのText Generation Inferenceによって行われるように、複数のGPUに計算を分散することで、GPT-4などのモデルを低レイテンシでサービス提供できます。
言語モデル以外にも、テンソル並列処理は、大規模なモデルや大きなバッチサイズが分散計算を必要とするコンピュータビジョンや科学計算などの他のドメインでも使用されます。例えば、高解像度画像での大規模なビジョントランスフォーマー（ViT）のトレーニングは、テンソル並列処理の恩恵を受けることができます。
課題と制限
その利点にもかかわらず、テンソル並列処理にはいくつかの課題があります。主な問題は通信オーバーヘッドです。all-reduce操作には、高帯域幅、低レイテンシのインターコネクトが必要です。イーサネットなどの遅いネットワークを備えたシステムでは、テンソル並列処理は非効率になる可能性があります。したがって、通常は単一サーバーまたは緊密に結合されたクラスター内で使用されます。もう1つの課題は負荷分散です: 不均一なシャーディングは、一部のデバイスが十分に活用されない原因となる可能性があります。さらに、テンソル並列処理は、各デバイスが重みのシャードと中間アクティベーションを格納する必要があるため、慎重なメモリ管理が必要です。最後に、テンソル並列処理を手動で実装することは複雑です。モデルコードの変更と通信プリミティブの処理が必要です。幸いなことに、PyTorchやTensorFlowなどのフレームワークは、テンソル並列処理の組み込みサポートを提供し、複雑さの多くを抽象化しています。
ソフトウェアとフレームワークのサポート
いくつかの深層学習フレームワークとライブラリがテンソル並列処理をサポートしています。PyTorchは、テンソルのシャーディングのためにtorch.distributedモジュールとtensor_parallel APIを提供しています。TensorFlowは、モデル並列処理をサポートするtf.distributeを提供しています。Megatron-LM、DeepSpeed、Hugging FaceのTransformersなどの専門ライブラリは、トランスフォーマーモデル用のテンソル並列処理を実装しています。これらのツールにより、研究者やエンジニアは、多くの場合、デバイス数とシャーディング戦略を指定するだけで、最小限のコード変更でテンソル並列処理を適用できます。
将来の方向性
モデルが成長し続けるにつれて、テンソル並列処理は重要な技術であり続けるでしょう。将来の開発には、階層的all-reduceなどのより効率的な通信アルゴリズムや、AMD GPU、Intelアクセラレータ、AWS TrainiumやGroqのLPUなどのカスタムチップを含む異種ハードウェアとのより良い統合が含まれる可能性があります。さらに、システムが最適なシャーディング戦略を決定する自動並列処理発見は、活発な研究分野です。大規模AIへの需要が高まる中、テンソル並列処理は、スケールと効率の課題に対応するために進化し続けるでしょう。