パイプライン並列処理は、ホームAIアクセラレーター間でモデルのレイヤーをどのように移動させるのか?

エヴァ・ウォン は テクニカルライター であり ZimaSpaceの常駐ティンカーでもあります。 生涯のオタクであり、 ホームラボとオープンソースソフトウェアに情熱を持っています。彼女は複雑な技術的概念をわかりやすく、 実践的なガイドに翻訳することを専門としています。エヴァはセルフホスティングは楽しくあるべきで、怖がるものではないと信じています。彼女のチュートリアルを通じて、コミュニティが ハードウェアのセットアップを解明する手助けをしています。初めてのNAS構築からDockerコンテナの習得まで。

パイプライン並列化は、異なるレイヤーグループを順番に実行するデバイスステージに割り当て、各ステージ間でアクティベーションを渡すことで、モデルを家庭用AIアクセラレーター間に分散します。

この手法は、深いモデルが1台のアクセラレーターには大きすぎるものの、複数のデバイスに収まるレイヤーグループへ分割できる場合に役立ちます。テンソル並列化とは異なり、各ステージは同じ行列の断片ではなく、完全なレイヤーを担当します。主なトレードオフは、デバイスごとの重み常駐量を減らせる一方で、アクティベーション転送、ステージ間の負荷の偏り、パイプラインのアイドル時間が発生することです。

パイプライン並列化は深いモデルを順番に実行するレイヤーステージへ分割する

トランスフォーマーは本質的に順序構造を持っています。後段のレイヤーを実行する前に、前段のレイヤーで入力を変換する必要があります。パイプライン並列化では、この深さ方向に連続する部分をグループ化し、各部分を異なるアクセラレーターに割り当てます。

PyTorchのパイプライン実行ランタイムは、モデルを異なるデバイス上で順番にステージとして実行されるパーティションに分割します。各デバイスには、割り当てられたステージのパラメーターだけが保存されます。

これにより、デバイスごとのモデルメモリ使用量を削減し、1台の家庭用GPUには収まらないモデルを実行できる場合があります。

アクティベーションはステージから次のステージへ順番に進む

GPU 0が担当するレイヤーの処理を終えると、隠れ状態のアクティベーションをGPU 1へ送信します。GPU 1は次のレイヤー群を適用し、さらに別のアクティベーションを転送します。これを繰り返し、最終ステージがロジットなどのモデル出力を生成します。

vLLM-Omniは、パイプライン並列化をトランスフォーマーブロックをGPU間で順番に実行するステージへ分割する手法として説明しています。基本的なフォワードパスでは、隣接するステージ間を越えて送る必要があるのは境界部分のアクティベーションだけです。

テンソル並列化とは異なり、各デバイスが同じレイヤー内の1つの行列を共同で計算するわけではありません。それぞれが異なるレイヤーを担当します。

1つのリクエストでもステージを順番に通過する

1つのトークンまたはマイクロバッチでは、後段のステージは前段から必要なアクティベーションが届くまで処理を開始できません。4つのレイヤーグループを4台のGPUに配置しただけで、1つの順次処理リクエストが4倍速くなるわけではありません。

NVIDIAの並列化ガイドでは、PPをモデルのレイヤーを深さ方向に分割する手法と定義しています。深さ方向の直列依存性は残ります。

そのため、すぐに得られる効果はレイテンシーの短縮よりも、処理可能なモデル容量の拡大である場合があります。ランタイムがステージ間で異なる処理をオーバーラップできれば、スループットは向上します。

マイクロバッチでパイプラインを埋め、複数のステージを同時に稼働させる

入力をマイクロバッチに分割するか、複数の推論項目を異なるステージで同時に処理できる場合、パイプラインは効率的になります。ステージ2が1つのマイクロバッチを処理している間に、ステージ1は次のマイクロバッチの処理を開始できます。

Hugging Faceの並列化ガイドでは、パイプライン実行においてパイプラインのアイドル時間を削減するためにマイクロバッチを使用すると説明しています。推論ランタイムが独立した処理をオーバーラップできる場合は、同じ稼働率向上の原則が適用されます。

インタラクティブな単一ユーザーのトークン生成では、新しいトークンごとに前のトークンの出力が必要になるため、自然に利用できる並列性は低くなります。バッチ処理されたリクエストや複数ユーザーからのトラフィックがあれば、各ステージを稼働させ続ける機会を増やせます。

ステージごとの処理時間が不均一だとバブルが生じ、最も遅いステージが全体を左右する

あるステージにより多くのレイヤー、低速なハードウェア、または特に処理コストの高い演算が含まれていると、次のステージは入力を待ち、前のステージもやがて空きを待つことになります。これらのアイドル時間がパイプラインバブルです。

DeepSpeedのパイプラインエンジンは、モデルのレイヤーをパイプラインステージへ分割することを重視しています。そのため、パラメーター容量のバランスだけでなく、ステージごとの計算量を均等にすることも重要です。

異なる家庭用アクセラレーターを組み合わせる構成では、低速なGPUに割り当てるレイヤーを意図的に少なくすることもできます。ただし、メモリ容量、アクティベーションサイズ、対応するデバイス間転送が配置を制約します。

パイプライン並列化は重みの収容能力と引き換えにアクティベーション転送とスケジューリングの負荷を受け入れる

PPは、モデルが深く、各デバイスに収まるレイヤーグループへ分割できる場合に魅力的です。テンソル並列化で一般的なレイヤー内の頻繁な集合通信ではなく、境界部分のアクティベーションを通信しますが、その代わりに直列依存性とバブルが発生します。

ZimaSpaceの複数ユーザーによる家庭内AIの同時実行に関するガイドも関連します。同時に処理する作業が増えると、パイプラインを埋めやすくなる一方で、KVキャッシュとスケジューリングの負荷も高まるためです。

ステージごとのメモリ使用量、ステージの実行時間、アクティベーション転送時間、エンドツーエンドのレイテンシー、スループット、デバイスのアイドル率を測定してください。モデルを収容できても、アクセラレーターの半分がアイドル状態になる構成であれば、別の分割方法が必要かもしれません。

ZimaSpaceのローカルAIハードウェアガイドが容量の基準を示しています。パイプライン並列化は、1台のデバイスのメモリ上限を超えるための方法の1つですが、成否はGPUの台数だけでなく、ステージのバランスに大きく左右されます。

テック&AIハブ

もっと読む

Get More Builds Like This

Stay in the Loop

Get updates from Zima - new products, exclusive deals, and real builds from the community.

Stay in the Loop preferences

We respect your inbox. Unsubscribe anytime.