vLLM-Omni 分散レイヤーワイズオフロード

vLLM-Omni の分散レイヤーワイズオフロード（DLO）により、1デバイスのHBM容量を超える動画生成モデル（例：64B Cosmos3-Super）を、ホストメモリのオーバーヘッドを最小限に抑えながら複数のNPUsまたはGPU上で実行できます。このシステムは、メタデバイス初期化、重みシャーディング、およびダブルバッファドプリフェッチパイプラインを組み合わせることで、大規模なDiffusion Transformer（DiT）モデルに伴うメモリボトルネックを解決します。

HBMおよびホストメモリのボトルネックを解決する

大規模なDiTモデルはしばしばデバイスのHBMに収まりません。従来のオフロードや並列化戦略は、大きなトレードオフを伴います。データ並列（DP）構成での純粋なレイヤーワイズオフロードでは、通常各ランクがモデルの完全なコピーをホストメモリに保持する必要があり、ホストRAMの要件がデバイス数に線形に増加します（O(dp_size × model_size））。

分散レイヤーワイズオフロードは、以下の4つの主要な技術的メカニズムにより、これらの制限を克服します。

1. メタデバイス初期化とmmap重みロード

モデル作成時の巨大なResident Set Size（RSS）のスパイクを防ぐために、vLLM-Omniはメタデバイス初期化とメモリマップ（mmap）ロードを使用します。

メカニズム：システムは to_empty(device="meta") を使ってDiTモジュールをメタデバイスに変換し、パラメータのストレージを解放しながらメタデータを保持します。その後、これらのパラメータを共有OSページキャッシュを直接指すmmapビューに置き換えます。
結果：すべてのランクが同じsafetensorsファイルをmmapするため、OSはキャッシュにファイルページを1つだけ保持します。Cosmos3-Nano DP4では、冷起動時のcgroup可視ピークメモリが178 GBから47 GBに低下（73%削減）しました。

2. 重みシャーディングとAllGather再構成

各ランクがピンされたCPUメモリにモデルの完全なコピーを保持する必要を排除するために、vLLM-Omniは重みシャーディングを実装しています。

メカニズム：各ランクはモデル重みの1/dp_sizeのみを保持します。実行時、専用の通信ストリーム上で all_gather_into_tensor を使用して、現在のレイヤーの完全な重みをデバイス上で再構成します。
結果：ピンされたホストメモリの合計は dp_size × model_size から単純に model_size に削減されます。Cosmos3-Super DP4では、ランクあたりのピンメモリが124 GBから31 GBに低下しました。

3. ダブルバッファドプリフェッチパイプライン

データ移動中にGPUがアイドル状態になるのを防ぎ、モデルの深さに関係なくHBM使用量を一定に保つために、vLLM-Omniはダブルバッファ方式を使用しています。

メカニズム：システムはモデルで最大のブロックサイズに合わせた正確に2つのデバイスバッファを維持します。計算ストリームがスロット0でレイヤーNを処理している間、バックグラウンドストリームはスロット1にレイヤーN+1のH2D（ホストからデバイス）転送とAllGatherを処理します。
結果：重みのHBM使用量は2 × max_block_sizeで上限に抑えられます。Cosmos3-NanoおよびCosmos3-Superでのテストでは、モデルサイズが3.8倍に増加しても、ピークHBMはわずか22%（23.1 GBから28.1 GB）しか増加しませんでした。

4. DPマルチコンカレントによるスループット向上

AllGatherはリクエストに依存しない（アクティベーションではなく重みを集約する）ため、vLLM-Omniは異なるDPランクが並行して異なるリクエストを処理しつつ、重み再構成において同期を維持できます。

メカニズム：スケジューラは最大 dp_size 個のリクエストをバッチ化します。各DPランクはリストから1つのリクエストを選択し、パイプラインの単一リクエストフォワードパスを通じて実行します。
結果：これによりAllGatherのオーバーヘッドが均等化されます。測定結果では、4つの並行リクエストでHSDP単一リクエストベースラインの3.3倍のスループットを達成し、理想的な線形スケーリングの約83%に達しました。

パフォーマンスと検証結果

プラットフォーム非依存性

分散レイヤーワイズオフロードは、NVIDIA GPU（CUDA/NCCL）とAscend NPU（CANN/HCCL）の両方で動作します。

NVIDIA B300 GPUでは、DLO+AG DP4で4つの並行リクエストを実行した場合、1024×1024 T2IタスクにおいてHSDP+USP4に比べて1.39倍のスループットを達成しながら、HBM使用量は30%（12.6 GiB vs 42.0 GiB）に抑えられました。720p 10秒の動画生成では、DLO+AG+USP4はHSDPの遅延に2.13%以内で近づきながら、HBM使用量は47%に抑えられました。

トポロジーアウェアポリシー（MiniMax-H3研究）

8× NVIDIA B300 GPU上でMiniMax-H3モデルを用いた研究では、最適なDLOモードはハードウェアトポロジーに依存することが示されました。

DP1×SP8：最低遅延（P50 34.55秒）のためにAllGatherが推奨されます。
DP4×SP2：スループットのバランスの取れた点（151.89動画/時間）を提供します。
DP8×SP1：最高スループット（183.78動画/時間）と最低エネルギー消費（43.97 Wh/動画）のためにランクローカルDLOが推奨されます。

Ascend NPUにおけるメモリアカウンティング

Ascendハードウェアでは、pin_memory() は /dev/davinci_manager を通じて割り当てられ、シャードはCPUカーネルDMAメモリに配置されます。このメモリはcgroupメモリコントローラーからは見えません。したがって、cgroup可視メモリは O(model_size + dp_size × constant) に比例しますが、合計物理RAMにはこれらのピンされたDMAシャードも含まれます。

メモリスケーリングの要約（外挿）

測定されたメモリモデルに基づき、vLLM-Omniは2 TBのシステムRAM制限内で非常に大きなモデルにスケーリング可能です：

モデルサイズ	dp_size	推定cgroupピーク	推定合計RAM	2 TBに収まるか？
33 GB	4	47 GB	~80 GB	はい
124 GB	4	172 GB	~296 GB	はい
185 GB	4	~220 GB	~405 GB	はい
400 GB	4	~423 GB	~823 GB	はい
400 GB	8	~443 GB	~843 GB	はい