FLINT: Efficiently Leveraging High Bandwidth Flash for Capacity-Scalable LLM Inference Acceleration

人工知能の世界は現在、物理的な限界へと突き進んでいます。コードを書き、物語を起草し、複雑な問題を解決する最も高度な言語モデルは、あまりにも巨大になりすぎて、それらを動かすコンピュータ・プロセッサに付随する標準的なメモリチップには収まりきらなくなっています。10億冊の本が入った図書館を、数十冊分しか入らないバックパックに入れて運ぼうとする場面を想像してみてください。これが、これらのモデルを動かしているハードウェアが直面している現状です。これらのシステムを機能させるために、エンジニアはモデルを多くの別々のコンピュータに分割することを余儀なくされてきました。しかし、コンピュータ同士が負荷を共有するために常に通信し合わなければならないため、この解決策は高価で、遅く、そしてエネルギーを大量に消費します。ハイバンド幅フラッシュと呼ばれる新しい技術は、これら巨大なライブラリをプロセッサのすぐ隣に保存する方法を提供し、小さなパッケージの中にテラバイト級の容量を実現します。しかし、この新しいストレージはプロセッサ自身のメモリよりも読み取り速度がはるかに遅いため、コンピュータがデータの到着を待つためにほとんどの時間を費やしてしまうというボトルネックを生み出しています。

ファーウェイ（Huawei）とチューリッヒ工科大学の研究者たちは、この「待ち」の問題を解決するために、FLINTと呼ばれるシステムを設計しました。最近の論文で詳述された彼らの研究は、以前のデザインが必要としていた複雑な事前計画なしに、この大容量フラッシュストレージをプロセッサの要求に追いつくほど高速に動作させる方法を提案しています。次にプロセッサが必要とするデータを予測する代わりに、FLINTはリアルタイムでプロセッサを監視し、要求が発生した瞬間にデータを整理します。これらの要求をグループ化し、ストレージチップの内部バッファをより効率的に使用することで、このシステムはプロセッサ自身のメモリに近い速度でデータをストリーミングできます。チームは、現在存在する最も大規模で複雑なモデルを含む6つの異なる大規模言語モデルを用いて、このシステムをシミュレーションしました。その結果、このアプローチにより、以前は数十台のマシンを必要としていたタスクを、単一のコンピュータ・パッケージで処理でき、エネルギー消費を大幅に抑えつつ、より高速に結果を出すことができることが示されました。

問題の核心は、これらのモデルがどのように構築されているかにあります。推論（モデルが応答を生成する行為）を実行するには、コンピュータはモデルの「重み」、つまりその知識を定義する数十億もの数値に絶えずアクセスしなければなりません。かつて、これらの重みはプロセッサに直接付随する高速メモリに完全に収まるほど小さかったものです。現在、数百億あるいは数兆のパラメータを持つモデルでは、重みは高速メモリには大きすぎます。代替案としては、ソリッドステートドライブ（SSD）を使用する方法がありましたが、これは容量は大きいものの安価ですが、あまりにも遅すぎます。あるいは、複数のプロセッサを使用する方法もありますが、これは作業を同期させる際に遅延を引き起こします。ハイバンド幅フラッシュは、中間の選択肢として提案されました。これは、プロセッサと同じ小さなパッケージに収まり、膨大な容量を提供しながらも、プロセッサが直接使用して停止してしまうほどではない読み取り速度を持つストレージ技術です。

フラッシュストレージを使用しようとするこれまでの試みは、読者が次にどの本を必要とするかを司書が予測するような手法に頼っていました。システムは、プロセッサが次に必要とするであろうモデルのレイヤーを予測し、プロセッサが要求する前にそれらの重みをフェッチ（取得）しようとします。しかし、研究者たちは、この「予測ゲーム」が惨めな失敗に終わることを発見しました。モデルは非常に複雑であり、その挙動は質問の内容によって変化するため、予測はしばしば外れます。システムは、プロセッサがまだ必要としていないデータをフェッチしてしまい、限られたバッファ空間を埋めてしまい、実際に必要だったデータを破棄せざるを得なくなります。その結果、プロセッサは正しいデータが再びフェッチされるのを待つために、ほとんどの時間を浪費することになりました。

FLINTは、この予測を取り除くことで、このアプローチを変更しました。司書が予測するのではなく、このシステムはリアルタイムで車の流れを監視する、非常に効率的な交通管制官のように機能します。プロセッサが特定のデータへの要求を行うと、FLINTは単にその一つのピースをフェッチするだけではありません。FLINTはプロセッサからのリクエストのストリームを観察し、それらをグループ化します。もしプロセッサが、フラッシュチップ内の同じ物理領域に格納されているいくつかのデータへの要求を行った場合、FLINTはそれらを単一の大きなリクエストへとまとめます。これにより、フラッシュチップは一度に大量のデータを読み取ることができ、そのフルスピードを活用できます。システムはこのデータをチップの内部バッファに保持し、プロセッサが必要とするまさにその瞬間に提供できるよう準備します。このダイナミックなグルーピングにより、システムは不要なデータをフェッチして時間を無駄にすることがなくなり、データパイプラインを常に満たし、流し続けることができます。

フラッシュストレージにおけるもう一つの大きな障害は、メンテナンスが必要です。時間の経過とともに、データの読み取り行為によってチップ内部の電気的な電荷が劣化し、エラーが発生することがあります。これを修正するために、チップは定期的に停止してデータを新しい場所に書き換えなければなりません。以前のデザインでは、このメンテナンスがプロセッサがデータを読み取ろうとしている最中に行われる必要があり、システムが長時間フリーズする原因となっていました。研究者たちは、「ファントム・プレーン（幻影平面）」システムを作成することで、この問題を解決しました。彼らはチップに、予備のレーンとして機能する少量の追加スペースを加えました。チップの特定のセクションでメンテナンスが必要になったとき、システムはメインのデータフローを止めることなく、バックグラウンドで静かにデータをこの予備レーンにコピーします。これにより、メンテナンス作業がプロセッサの作業を妨げることがなくなり、チップが経年劣化してもシステムをスムーズに稼働させ続けることができます。

チームはまた、プロセッサのリクエストをチップ上の物理的な位置へと変換する方法を簡素化しました。標準的なストレージシステムは、絶え間ない書き込みと書き換えを処理するように構築されており、管理のために複雑なマップと重厚な機構を必要とします。言語モデルの重みは一度書き込まれたら変更されないため、FLINTは読み取りにのみ最適化された、よりシンプルで軽量なマップを使用します。この複雑さの軽減により、システムはリクエストをほぼ瞬時に変換でき、もう一つの遅延レイヤーを取り除いています。

研究者たちが詳細なシミュレーションでこれらのアイデアをテストした際、その結果は驚くべきものでした。彼らは、低速なソリッドステートドライブを使用した標準的なシステム、多くのプロセッサを使用して連携するシステム、そしてフラッシュストレージを使用する以前の試みの3つのセットアップと比較しました。FLINTシステムは、テキストのデコード（モデルが出力を生成するプロセス）において、ソリッドステートドライブを使用するシステムよりも最大1,205倍高速でした。多くのプロセッサが協力して働くシステムと比較しても、FLINTは平均して2倍以上高速でした。おそらく最も重要なことは、FLINTを備えた単一のコンピュータが、以前は4台から8台のコンピュータを必要としていた作業を、より少ないエネルギーを消費しながら実行できたことです。

研究では、システムの寿命についても調査が行われました。フラッシュストレージは絶えず読み取られるため、研究者たちはチップがどの程度早く摩耗するかを懸念していました。彼らは、エラーが発生する前にデータを書き換える彼らのメンテナンスシステムが、激しい使用条件下でもストレージの寿命を数年間にわたって延ばすことを発見しました。システムは、高密度なモデルから、問題を解決するために異なる専門家（エキスパート）の間を切り替える複雑なモデルまで、6つの異なるモデルでテストされ、そのすべてにおいて良好なパフォーマンスを示しました。研究者たちは、システムがチップにわずかな物理的スペースを追加するものの、そのコストは速度と容量の莫大な利点に比べれば無視できるものであると指摘しました。

この研究は、人工知能の限界が、モデルがいかに賢いかだけでなく、いかにデータを移動させるかにあることも証明しています。プロセッサとストレージの相互作用のあり方を再考することで、研究者たちは、膨大な容量と高速なスピードの両立が可能であることを示しました。FLINTシステムは、かつてはリアルタイムでの使用には遅すぎると考えられていた技術を、次世代の人工知能のための実用的なソリューションへと変え、強力なモデルをデータセンター全体ではなく、単一のマシンで実行することを可能にします。この知見は、適切なエンジニアリングがあれば、メモリの物理的な制約を克服できることを示唆しており、将来のより有能で効率的なAIシステムへの扉を開いています。