最近の清華大学の研究者による研究は、大規模言語モデル(LLM)の計算とハードウェア構成を再編成することで、推論コストを大幅に削減できることを示しています。彼らは「アテンションオフロード」と呼ばれる手法を導入し、コスト効率の良いGPUをメモリ集約型タスクに使用することで、高性能アクセラレーターが計算重視の操作に集中できるようにしています。
高級AIアクセラレーターは高価で希少、かつ需要も高いため、アテンションオフロードは企業がLLMを大規模に展開する際にハードウェア資源を最適化する機会を提供します。
2種類の計算
LLMの推論には、利用可能なメモリや処理能力を最大限に活かすために戦略的に整理する必要があるさまざまな操作が含まれます。これらの操作は、計算重視型とメモリ重視型の2つの主なタイプに分類できます。計算重視型の操作は、A100やH100のような高速アクセラレーターの恩恵を受ける一方、特に各トークンによって起動される自己注意メカニズムなどのメモリ集約型操作は、十分なビデオRAM(VRAM)を必要とします。
研究者たちは、「このメモリ集約型ワークロードは現代のアクセラレーターの強みと対立し、メモリコントローラーが過負荷になる一方で計算コアがアイドル状態になる」と指摘しています。このリソースの不均衡は、ユーザーのプロンプトが長くなったり、モデルとの会話が延びるとさらに悪化します。
革新的な解決策:アテンションオフロード
現在のアプローチは、推論のための高級アクセラレーターの均一アーキテクチャを拡張することに焦点を当てています。企業はH100プロセッサに多額の投資を行い、推論能力を拡大しようとしますが、これによりコストが膨らんでしまい、ハードウェアの利用効率が低下します。
研究者たちは「LLM生成フェーズの独自の要求には、効率向上とコスト削減のために異種アーキテクチャが必要」と主張しています。彼らの研究によれば、さまざまなタイプのアクセラレーターはLLM推論の特定の側面に適しています。たとえば、消費者向けGPUはメモリ集約型タスクにおいて経済的な選択肢を提供し、高級モデルに比べてドルあたりのメモリ容量と帯域幅が3倍になります。ただし、これらの低コストオプションに完全に依存することは、計算パワーが限られているため、非効率的になる可能性があります。
アテンション計算は非常に並列化可能であり、複数のコスト効率の良いメモリ効率GPUに分散することができます。
異種アーキテクチャの実装
アテンションオフロード技術には、計算能力に特化したアクセラレーターのプールとメモリ帯域幅に最適化されたプールの2つの異なるグループを作成することが含まれます。これにより、アテンションタスクは低コストのGPUによって処理され、高級アクセラレーターが他の操作を管理します。
研究者たちは、「この異種アーキテクチャにより、計算力、メモリ容量、帯域幅を効率的に組み合わせてLLM推論を強化し、過剰なコストをかけずに実現するサービスシステムが可能になる」と説明しています。このハードウェアの強みを操作要件と戦略的に整合させることで、企業はメモリと計算に最適化されたアクセラレーターのバランスの取れたミックスに投資することで予算を最大限に活用できます。
アーキテクチャの課題に対処
この研究は、特に2つのアクセラレーターグループを接続するために必要な帯域幅に関連する異種アーキテクチャの課題も評価しています。研究結果は、標準のシステムバスであるPCIe 4.0が十分であるだけでなく、AIデータセンターで一般的に使用されているネットワーキング技術(200Gb InfinibandやEthernet)も適切であることを示しています。
高度なスケジューリングとパイプライニング技術を活用することで、非均一アーキテクチャによって引き起こされるレイテンシを軽減し、メモリと計算リソースが順次計算に妨げられることなく同時に機能できるようになります。
Laminaの導入
研究者たちは、アテンションオフロードを採用した分散型異種LLM推論システムであるLaminaを開発しました。Laminaは、消費者向けGPUを使用して計算されたアテンション値(「KVキャッシュ」)を保存し、アテンション操作を実行し、高級アクセラレーターがモデルパラメータやその他の推論タスクを管理します。これらのデバイスは、同じ物理マシン内で機能することも、複数のノードに分散することもできます。
KVキャッシュのストレージとアテンション計算をメモリ効率の良いGPUにオフロードすることで、Laminaは、広く使用されているLLMサービングプラットフォームであるvLLMによって管理されるバッチの10.7倍から64倍のサイズを処理できます。この効率性は、高価な計算最適化アクセラレーターを効果的に活用するために重要です。
実験評価の結果、Laminaは13Bおよび33Bモデルに対して、既存のソリューションに比べてコストあたり1.48倍から12.1倍のスループット率を達成しています。
LLMが主流になるにつれて、企業はコスト効果の高い推論とアクセラレーターへの資本支出の削減のための革新的な戦略が必要です。アテンションオフロードはこの目的に成功裏に応えています。研究者たちはまだLaminaのコードを公開していないものの、基本的な概念は明確に示されており、オープンソースコミュニティによる迅速な実装が期待されます。