テンソル並列処理は、複数のGPUにわたってローカルAI推論をどのように分割するのか？

テンソル並列処理は、大規模な重みテンソルを各トランスフォーマー層内で複数のGPUに分割し、それぞれの部分的な結果を統合することで、ローカルAI推論を分散します。

これは、各GPUに別々のリクエストを割り当てたり、異なる層を別々のデバイスに割り当てたりする方法とは異なります。テンソル並列処理では、すべてのGPUが同じ層の計算に参加し、多くの場合、生成されるすべてのトークンで処理を行います。この方法により、複数の家庭用GPUにまたがってモデルを配置できますが、GPU間通信が推論処理のクリティカルパスに組み込まれます。

テンソル並列処理はモデル全体をコピーせず、1つの層を分割する

データ並列処理では、各GPUにモデルのコピーを持たせ、リクエストやバッチを分散します。一方、テンソル並列処理は、1つのモデル内の各層にある大規模なパラメータテンソルをデバイス間で分割します。

NVIDIA NeMoでは、TPを1つの層のパラメータテンソルをGPU間に分散する処理と定義しています。各GPUが保持するのは、対象となる行列の一部だけです。

これは、1つのモデル、あるいは1つの大きな層でさえ、単一の家庭用GPUに無理なく収まらない場合に役立ちます。

列方向と行方向のシャードで行列乗算の処理を分割する

トランスフォーマー層には、大規模な線形射影が含まれています。列方向の分割では異なる出力列をGPUごとに割り当て、行方向の分割では異なる入力行や特徴量の範囲を割り当てます。

PyTorchのテンソル並列チュートリアルでは、トランスフォーマー層に行方向と列方向の並列方式を適用しています。各ランクは、ローカルに保持する重みシャードを使って、行列積の一部を計算します。

モデルは依然として、論理的には1つの層を表します。分割によって変わるのは、計算の各部分を実行する場所と、部分的な結果を統合する方法です。

集合通信によって層の論理的な出力を再構成する

各GPUがテンソルの一部しか参照できないため、次の計算に必要な表現を得るには、処理の前にall-reduce、all-gather、reduce-scatter、または同等の集合通信が必要になることがあります。

Open MPIでは、AllReduceを、プロセス間で値を結合し、その結果をすべての参加者に再配布する処理と定義しています。テンソル並列ランタイムは、この種の集合通信に加えてall-gatherやreduce-scatterを使用し、シャード化された層の結果を再構成または再分配します。

家庭用ワークステーションでは、GPU間接続の品質によって、分割が処理時間の短縮につながるのか、それとも単に容量を拡張するだけなのかが決まる場合があります。

Featured

ZimaCube 2 パーソナルクラウドホームNAS

ZimaCube2

$799.00 - $2,499.00

Buy Now

通信は層単位の頻度で発生するため、高速なインターコネクトが重要

テンソル並列処理では、トランスフォーマーブロックごと、さらに生成されるトークンごとに複数回の集合通信が必要になることがあります。PCIeのみのシステムは、大規模な分散処理向けに設計されたハイエンドアクセラレータファブリックと比べて、GPU間のピア帯域幅が大幅に低くなります。

AMD RCCLでは、PCIe接続GPU向けのピアツーピア転送について説明しています。使用する集合通信ライブラリはプラットフォームによって異なりますが、同じトポロジー上の制約が存在します。

2台のGPUに十分な合計VRAMがあれば、より大きなモデルを正常に実行できる場合があります。しかし、各層で同期を待つ必要があるため、トークン処理速度が期待を下回ることもあります。

性能の異なるGPUでは、最も遅いシャードが全体のペースを決める

テンソル並列化された層は、必要な部分結果が到着してからでなければ次の処理に進めません。1台のGPUの計算スループット、メモリ帯域幅、または接続速度が低い場合、より高速なランクが待機することになります。

DeepSpeedの自動テンソル並列推論機能は、推論プロセスグループ全体でのモデルシャーディングを前提に設計されています。実用上の効率は、参加するデバイスがバランスの取れた処理量を担えることを前提としています。

家庭にある性能の異なるGPUを組み合わせる場合でも、モデルを収めるという点では役立ちます。ただし、ハードウェアの性能差が大きい場合、シャードのサイズを均等にすることが必ずしも最適とは限りません。

大規模な層と単一GPUの限られたメモリにはテンソル並列処理を選ぶ

テンソル並列処理が特に有効なのは、大きな隠れ層の次元や層のテンソルをGPU間で分割する必要があるモデルです。特に、デバイス間に高速なローカルインターコネクトがある場合に適しています。複数の独立した小規模リクエストを処理する方法として、必ずしも最適とは限りません。

ZimaSpaceの記事ローカルAIアクセラレータのメモリ計画では、必要容量の基準を説明しています。テンソル並列処理では、1つのモデルの層を複数のデバイスに分散することで、この基準が変わります。

可能であればまず1台のGPUでベンチマークを行い、次に同じモデル、プロンプト、コンテキスト、バッチを使って2台以上でも測定します。GPUごとのメモリ使用量、1秒あたりのトークン数、レイテンシ、集合通信にかかった時間、リンク使用率を記録してください。

ZimaSpaceの記事マルチユーザー環境におけるローカルAIの負荷は、サービング側からの比較材料です。TPでは1つのモデルを複数のデバイスにまたがって配置し、リクエストの同時実行数によって、その分散モデル上で競合する独立したコンテキストの数が決まります。