Amazon FSx for Lustre 上の GPUDirect と TurboQuant で LLM モデル読み込みを高速化し、コンテキストウィンドウを拡大
本記事では、Amazon FSx for Lustre、NVIDIA GPUDirect Storage (GDS)、およびシャーディング並列読み込みを組み合わせることで、大規模言語モデルのコールドスタート時の最初のトークン生成時間を数分から数秒に短縮する方法と、TurboQuant KV キャッシュがコンテキストウィンドウサイズを大幅に拡大する効果について解説します。
AWS GPU インスタンス上で大規模言語モデル (LLM) のデプロイを反復していると、GPU 高帯域幅メモリ (HBM) にモデルを読み込む時間がモデル規模の増大に伴って長くなり、コールドスタート時の最初のトークン生成時間 (TTFT) が悪化することに気づくでしょう。本記事では、Amazon FSx for Lustre と NVIDIA GPUDirect Storage (GDS) の組み合わせと、適切な計画によってコールドスタート TTFT を根本的に変革し、モデル読み込み時間を数分から数秒に短縮する方法を解説します。また、最近発表された TurboQuant KV キャッシュがコンテキストウィンドウサイズを大幅に拡大する効果についても触れます。
背景：AWS 上の NVIDIA Blackwell アーキテクチャ AWS は最近、NVIDIA の Blackwell アーキテクチャを搭載した Amazon EC2 P6e および P6 インスタンスファミリーをリリースしました。フラッグシップの P6e UltraServer は 72 個の NVIDIA Blackwell GPU を単一の NVLink ドメインに統合し、130 TB/s の双方向帯域幅、13.4 TB の HBM3e、360 petaflops の FP8 計算能力を提供します。これらの UltraServer は、通常、数兆パラメータ規模のフロンティアモデルの大規模分散トレーニングに使用されます。本記事では、単一の P6 または P5en インスタンスのコールドスタート TTFT の改善に焦点を当てます。
モデル読み込みのボトルネック 従来の CPU ベースのモデル読み込みでは、チェックポイントが CPU メモリを経由し、各 GPU に PCIe 経由で順次コピーされます。一方、シャーディングされた GDS 読み込みでは、チェックポイントがテンソル並列ランクごとに FSx for Lustre 上に事前分割され、全 8 GPU が EFA を介して直接 HBM にシャードを並列読み込みします。Llama 3.1 405B モデル（約 800 GB の BF16 チェックポイント）の場合、従来の読み込みには 10～20 分かかりますが、GDS はこのボトルネックをほぼ排除します。長い読み込み時間は、コールドスタートレイテンシだけでなく、自動スケーリングの応答性、障害回復、コスト効率にも悪影響を及ぼします。
直接パス：FSx for Lustre と GPUDirect Storage Amazon FSx for Lustre は、高パフォーマンスの並列ファイルシステムであり、NVIDIA GDS と組み合わせることで、CPU とシステムメモリを完全にバイパスして GPU メモリへの直接データパスを複数確立します。この統合は、Amazon Elastic Fabric Adapter (EFA) と NVIDIA GPUDirect Storage の 2 つの主要技術に依存しています。EFA は SRD プロトコルを使用して OS のオーバーヘッドを回避し、P5en インスタンスは 16 の 200 Gbps EFA インターフェースを備え、合計 3,200 Gbps のネットワーク帯域幅を提供します。GDS はネットワークインターフェースから GPU HBM への DMA 転送を可能にします。テストでは、Persistent_2 EFA ファイルシステム（1000 MBps/TiB、20 OST）を使用し、約 94 GiB/s のファイルシステムスループットを実現しました。
シャーディング並列読み込み P5en インスタンス（8 基の H200 GPU、各 141 GB HBM3e）を例にとると、Llama 3.1 405B モデル（FP8 で約 400 GB）は 1 つの GPU に収まらないため、テンソル並列処理が必要です。このアプローチは 4 つの段階で機能します。段階 0：インフラのプロビジョニング。同一 VPC とアベイラビリティゾーンに、EFA 対応の FSx for Lustre ファイルシステムと GDS 設定済みの GPU インスタンスを用意します。CloudFormation テンプレートとセットアップスクリプトが用意されています。段階 1：重みの事前シャーディングと量子化。オフラインで vLLM を使用し、モデルを 8 つのテンソル並列シャードに分割し、FP8 に量子化して FSx for Lustre に保存します。チェックポイントサイズは ~800 GB から ~400 GB に削減されます。このステップはモデル更新時のみ繰り返します。段階 2：GDS 並列読み込み。推論サーバー起動時、各 GPU が GDS 経由で FSx for Lustre から自身のシャードを並列読み込みします。段階 3：推論。モデル準備完了後、直ちに推論要求を処理します。
TurboQuant KV キャッシュ：TurboQuant は KV キャッシュを 4 ビットまたは 2 ビット精度で保存する量子化技術であり、品質を損なうことなく KV キャッシュメモリ使用量を 4～8 倍削減し、より長いコンテキストウィンドウを可能にします。
これらの技術を組み合わせることで、AWS 上で大規模モデルのコールドスタートをほぼ瞬時に実現し、より長いコンテキストをサポートできるようになります。