今回の論文
今回取り上げるのは、Ying Sheng、Lianmin Zheng、Binhang Yuan らによる論文「FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU」です。2023年3月に arXiv で公開され、その後 ICML 2023 で発表された論文です。研究分野は LLM 推論最適化、オフロード実行、高スループットサービングです。URL は https://doi.org/10.48550/arXiv.2303.06865 です。
この論文を選んだ理由は、モデルの中身を変えるのではなく、推論基盤の設計だけで「大きすぎて載らないモデルをどう動かすか」に正面から答えているからです。特に、GPU が潤沢ではない環境でもバッチ処理系のLLM活用を成立させる考え方が整理されており、開発や事業の観点でも学びが多い技術です。
どんな技術か
FlexGen は、単一の比較的小さな GPU しかなくても、大規模 LLM をできるだけ高いスループットで実行するための推論エンジンです。
発想の中心は、モデル重み・中間活性・KV キャッシュを全部 GPU に置こうとしないことです。代わりに、GPU、CPU メモリ、SSD という三層のメモリ階層をまとめて使い、どのテンソルをどこに置くか、どの順番で読み書きするかを最適化します。
つまり FlexGen は、「GPU に入り切る小さなモデルを速く動かす」技術ではなく、「GPU に載らない巨大モデルを、遅延をある程度許容しつつ、できるだけ多くのトークンをまとめて処理する」ための技術です。チャットの1リクエストを最速化するというより、評価バッチ、文書処理、社内データ変換のようなバックグラウンド処理向けの推論基盤だと考えると理解しやすいです。
課題
FlexGen が解こうとしている課題は、大規模 LLM の推論が GPU メモリに強く縛られていることです。
まず難しいのは、モデル重みだけでも巨大だという点です。論文では GPT/OPT 175B 級モデルの重み読み込みに数百 GB の GPU メモリが必要になると説明しています。これでは一般的な 16GB クラスの GPU に載りません。
次に問題になるのが、推論時には重みだけでなく活性や KV キャッシュも必要になることです。しかも高スループットを狙ってバッチサイズを大きくすると、KV キャッシュが非常に大きくなります。論文では、OPT-175B、入力長 512、出力長 32、実効バッチサイズ 256 の条件で、KV キャッシュが 1.2TB に達し、モデル重みの 3.8 倍のボトルネックになると報告しています。
既存手法にも限界がありました。量子化や圧縮だけでは、175B 級モデルを単一GPUで扱うにはまだ厳しい場面があります。逆に CPU やディスクへのオフロードだけに頼ると、どのテンソルをどのタイミングで動かすかの設計が悪いと I/O が支配的になり、バッチサイズを大きくできません。結果として、せっかくオフロードで「動く」ようになっても、「実務で回る速度」には届かないという問題がありました。
なぜこの課題を解く必要があるのかというと、現実の AI 活用はチャットUIだけではないからです。大量の文書に対する抽出、評価ベンチマーク、夜間バッチの分類や要約、ログ整形、フォーム処理などでは、1件の応答速度より、限られたマシンで何件さばけるかのほうが重要です。FlexGen はその前提で設計された点に価値があります。
用語解説
- オフロード
- GPU に載り切らない重みやキャッシュを CPU メモリや SSD に逃がし、必要なときだけ読み込む実行方法です。FlexGen は単なる退避ではなく、どのテンソルをどこへ置くかを戦略として最適化する点が重要です。
- KVキャッシュ
- 自己回帰生成で、過去トークンの Key と Value を保存して再利用する仕組みです。推論高速化に必須ですが、バッチサイズと系列長に応じて膨らむため、FlexGen では重みと並ぶ主要な配置対象になります。
- スループット指向推論
- 1リクエストの待ち時間より、単位時間あたりに何トークンまたは何件処理できるかを重視する推論設定です。FlexGen はこの前提に立つため、低遅延チャット最適化とは設計目標が少し異なります。
- バッチサイズ
- 同時にまとめて処理する入力数です。大きいほど計算やI/Oを償却しやすくなりますが、活性やKVキャッシュが増えるためメモリ要求も跳ね上がります。FlexGen はこのバランスを取りながら実効バッチサイズを大きくするのが狙いです。
- グループ単位量子化
- テンソル全体ではなく、小さな要素グループごとに最小値・最大値を取って量子化する方法です。FlexGen では演算高速化より圧縮率とI/O削減を重視し、重みとKVキャッシュを4bit化するために使われます。
技術の仕組み
FlexGen の面白さは、単に「CPUに逃がす」ではなく、推論全体を配置問題とスケジューリング問題として整理している点です。
基本アイデア
FlexGen の基本アイデアは、GPU、CPU、SSD をひとつの統合メモリ資源として見なし、その上で生成処理をブロック単位に分けて回すことです。
LLM 推論では、各レイヤーの重み、各ステップで生じる活性、そして生成が進むほど増える KV キャッシュを扱います。FlexGen はこれらを別々に最適化するのではなく、全部まとめて「どこに置くと最も高スループットになるか」を探索します。論文では、計算順序、テンソル配置、CPU 側の計算委譲まで含めて設計空間を定義し、線形計画法ベースで最適方針を探します。
三層メモリを前提にした配置設計
FlexGen では GPU が最速だが最小、CPU メモリは中間、SSD は最遅だが大容量という前提で扱います。ここで重要なのは、テンソルごとに別の置き方があり得ることです。
論文では、重みを GPU・CPU・ディスクに何パーセント置くかを wg, wc, wd、活性を hg, hc, hd、KV キャッシュを cg, cc, cd で表しています。つまり「重みはCPU中心、KVはCPU、活性はGPU寄り」のような方針を、モデルサイズやハードウェアに応じて自動で変えられます。
この見方は実務的にも重要です。多くの実装では「モデル重みをどこに置くか」だけに注目しがちですが、高スループット設定では KV キャッシュのほうが支配的になることがあります。FlexGen はそこを明示的に扱っています。
Zig-zag block schedule
FlexGen は計算順序も工夫しています。論文では、レイヤー、トークン、GPU バッチが絡む依存関係を考えた上で、zig-zag block schedule というブロック型の処理順序を採用しています。
考え方としては、あるレイヤーの重みを読み込んだら、その重みをできるだけ多くの GPU バッチに使い回してから次へ進む、という流れです。これにより、毎回小刻みに重みを読み直す無駄を減らせます。一方で、活性や KV キャッシュは各バッチごとに保存・再読込が必要なので、CPU や SSD の容量上限を見ながら、何バッチ分を1ブロックとして回すかを決めます。
論文では、このブロックスケジュールは理論上の最適 I/O 複雑度に対して 2 倍以内だと示しています。完全最適スケジュールより実装しやすく、それでも十分に良いという立て付けです。
重ね合わせ実行
FlexGen のもうひとつの重要な工夫は、読み込み・書き込み・計算を重ねることです。
具体的には、ある GPU バッチを計算している間に、次のレイヤーの重みをロードし、前のバッチの活性やキャッシュを書き戻し、次のバッチのキャッシュや活性を先読みします。論文の擬似コードでも、これらを依存関係のない並列タスクとして扱っています。
ここで効くのは、I/O を完全に消すことではなく、待ち時間を計算の裏に隠すことです。オフロード実行では I/O をゼロにはできないので、いかに GPU が遊ぶ時間を減らすかが勝負になります。
線形計画法によるポリシー探索
FlexGen を単なるハック集ではなく技術として面白くしているのが、方針を探索問題として解いている点です。
論文では、GPU メモリ容量、CPU メモリ容量、各デバイスの帯域、演算量、テンソルサイズなどを使ってコストモデルを作り、スループット最大化を目的に配置率やブロックサイズを決めています。変数数はそこまで多くなく、論文では配置に関する主要変数は9個として整理されています。
この発想から得られる示唆は、推論最適化は固定の正解があるわけではないということです。同じモデルでも、GPU が 1 枚なのか 2 枚なのか、CPU RAM が多いのか少ないのか、SSD 帯域が速いのかで最適方針は変わります。FlexGen はそこを実行時ポリシーとして扱います。
4bit量子化でI/Oを下げる
FlexGen は正確なオフロードだけでなく、近似手法も組み合わせます。中心は重みと KV キャッシュの 4bit 量子化です。
論文では、重みと KV キャッシュをグループサイズ 64 のグループ単位非対称量子化で 4bit 化し、計算直前に FP16 に戻します。狙いは演算器を整数化して速くすることではなく、保存量と転送量を減らすことです。つまり「計算の前に小さくして運ぶ」という思想です。
加えて論文では、Value キャッシュに 10% の疎性を入れる sparse attention も試しています。ただし FlexGen の本筋は、疎化よりも、オフロード設計と 4bit 圧縮の組み合わせだと捉えるほうが本質をつかみやすいです。
パイプライン並列への拡張
論文は単一GPUを主対象にしていますが、複数 GPU に対する拡張も示しています。レイヤーを複数 GPU に等分し、各 GPU で同じオフロード方針を使いながらマイクロバッチを流す形です。
ここで大事なのは、FlexGen の本質が「単一GPU専用の特殊技法」ではなく、「メモリ階層とI/Oを前提に推論を設計する」枠組みだという点です。GPU が増えても、I/O と配置の問題が消えるわけではない、という理解につながります。
実験と結果
論文では、FlexGen が本当に大規模モデルを単一GPUで回せるのか、既存オフロード方式と比べてどれくらい有利か、近似による精度劣化は小さいのかを検証しています。
何を検証したのか
主な検証項目は4つあります。1つ目は、単一 T4 GPU で 30B や 175B 級のモデルをどれだけ高スループットで動かせるかです。2つ目は、DeepSpeed Zero-Inference や Hugging Face Accelerate など既存オフロード系手法との比較です。3つ目は、4bit 量子化や疎化を入れても精度が保てるかです。4つ目は、HELM ベンチマークのような実際のバッチ用途をどの程度の時間で回せるかです。
どんな条件で評価したのか
論文の代表的な実験では、NVIDIA T4 16GB GPU 1枚、CPU メモリ 208GB、SSD 1.5TB という構成を使っています。対象モデルは OPT-6.7B、OPT-30B、OPT-175B で、入力長 512、出力長 32 など複数条件で throughput を測定しています。
評価指標の中心は generation throughput、つまり総生成トークン数を総処理時間で割った値です。加えて、Lambada の accuracy、WikiText の perplexity で近似手法の精度劣化も見ています。
175Bモデルを単一16GB GPUで実行
最も印象的なのは、OPT-175B を単一 16GB T4 上で動かしている点です。論文では、圧縮なしでも実効バッチサイズを 256 まで拡張して 0.69 token/s を達成し、4bit 圧縮ありでは実効バッチサイズ 144 で 1 token/s に到達したと報告しています。
この結果の意味は、「1件の応答が高速」というより、「巨大モデルを安価な構成でも夜間バッチや評価用途に現実的な形で回せる可能性を示した」というところにあります。推論基盤としての発想転換が効いた例です。
既存オフロード手法より大幅に高スループット
論文では、同じレイテンシ 5000 秒の条件で、OPT-175B に対して FlexGen が DeepSpeed Zero-Inference より 40 倍超高い throughput を出したとしています。さらに高いレイテンシを許す設定では、最大 throughput が 69 倍まで伸び、4bit 圧縮を使うと 100 倍に達したと報告されています。
ここで効いているのは、単に量子化したことではなく、大きな実効バッチサイズを取れることです。ベースラインはメモリ制約のためバッチサイズ 1 や 2 から増やせない一方、FlexGen は配置最適化で大きくバッチを積めます。スループット最適化では、ここが本質です。
ポリシー探索と重ね合わせの寄与
アブレーションでも、FlexGen の改善要因がはっきりしています。OPT-30B では全最適化ありで 7.32 token/s、175B では 0.69 token/s でしたが、175B でポリシー探索を外すと 0.27 token/s まで落ち、DeepSpeed 系の方針をそのまま持ち込むと 0.01 token/s まで落ちています。
また、overlap を無効にしたり CPU 側計算を使わなかったりすると throughput が落ちるため、配置だけでなく実行スケジュール自体が重要だとわかります。FlexGen は「何をどこに置くか」と「どう回すか」を両方設計して初めて効く技術です。
4bit圧縮の精度劣化は小さい
近似手法の評価では、OPT-30B の Lambada accuracy が FP16 で 0.725、4bit で 0.724、WikiText perplexity が 12.72 から 12.90 です。OPT-175B でも Lambada accuracy は 0.758 から 0.756、WikiText perplexity は 10.82 から 10.94 と小さな差に収まっています。
この結果から、少なくとも論文の設定では、4bit 化は「使えるか使えないか」を分けるほどの大きな精度崩壊を起こしていません。3bit は精度を保てなかったとも報告されており、4bit が実用上の落としどころだったと読めます。
HELMの実行時間も現実的
論文では、OPT-IML-30B を使った HELM の代表的 7 シナリオを 21 時間で完走できたと報告しています。これは速いというより、これまで多GPU前提だった評価作業を、限られた計算資源でも回せるようにした、という意味合いが大きいです。
開発現場では、推論そのものだけでなく、モデル比較、評価、データ変換の反復コストが重くなりがちです。FlexGen はそうした裏方処理に向いたシステムだと理解すると、この結果の価値が見えやすいです。
何に使える？
FlexGen の使い道は、低遅延チャットよりも、資源制約のあるバッチ推論や裏方ワークロードにあります。
大規模文書の夜間バッチ処理
社内文書の分類、要約、項目抽出、契約レビュー補助のような処理では、リアルタイム性より総処理量が重要です。FlexGen 型の設計なら、GPU を大量に並べなくても、大きなモデルをバッチで回して1晩で処理する構成を考えやすくなります。
モデル評価とベンチマーク
新しいモデルやプロンプトを比較するとき、評価セットをまとめて流す作業が発生します。こうした用途では 1 件ずつの応答速度より、総サンプル数をどれだけ早く消化できるかが大事です。FlexGen はまさにこの「評価のための推論基盤」に向いています。
GPUが限られた社内ツール
スタートアップや小規模チームでは、高価な大容量 GPU を常時使うのが難しいことがあります。その場合、対話サービスの本番ではなくても、社内向けレポート生成、ログ要約、データ整形などの内部バッチ用途で、大きめのオープンモデルを安価なマシンで回す選択肢が出てきます。
RAG前処理や知識ベース整備
RAG の本番応答は低遅延が求められますが、前処理側は比較的遅くても構いません。たとえば大量文書から要約メタデータを作る、FAQ 候補を抽出する、項目正規化を行うといった処理は FlexGen の前提と相性がよいです。大きなモデルを前処理にだけ使うという設計は十分あり得ます。
高価なGPUを常時確保したくない実験環境
研究やPoCの段階では、「常に最速」より「今ある計算機で試せる」ことのほうが重要です。FlexGen の考え方は、巨大モデルを試すための敷居を下げます。とくに比較実験やデータ生成では、スループット最適化がそのまま試行回数の増加につながります。
開発や事業へのヒント
FlexGen から得られるヒントは、モデル選定だけでなく、推論基盤の設計がプロダクトの可能性をかなり左右するということです。
使いたいモデルと提供形態を分けて考える
AIアプリを作るとき、すぐ「本番チャットにそのまま載るか」でモデルを判断しがちです。しかし FlexGen が示しているのは、大きいモデルには大きいモデルの役割があり、それが必ずしもリアルタイム提供とは限らないということです。前処理、評価、バックオフィス自動化なら、多少遅くても高性能モデルを使う価値があります。
インフラの制約をアルゴリズムで吸収する発想
この論文は、GPU が足りないから小さいモデルに妥協する、という単純な話ではありません。GPU・CPU・SSD の使い分けを設計することで、使えるモデルの幅を広げています。小規模プロダクトでも、モデル性能だけでなく、配置・キャッシュ・I/O 設計で勝てる余地があるとわかります。
KVキャッシュを主役として設計する
実務では重みサイズばかり話題になりますが、高バッチ推論では KV キャッシュがボトルネックになることがあります。FlexGen はその事実をかなりわかりやすく示しています。長文処理や大量同時実行を扱うなら、重みより先に KV の増え方を見積もるべき、というのは開発上の重要な学びです。
最適化は固定レシピではなく探索問題
どのテンソルをどこに置くべきかは、マシン構成で変わります。これは事業面でも重要で、同じソフトを別の顧客環境に導入したとき、GPU 構成やCPUメモリ量に応じて最適化の余地が残るということです。将来的には、推論エンジン自体がハードウェアを見て自動チューニングする方向が強くなると考えられます。
低遅延プロダクトとバッチ基盤を分離する
FlexGen を読むと、チャット向けのオンライン推論と、集計・前処理・評価向けのオフライン推論は、同じ基盤で無理に最適化しないほうがよいとわかります。事業としても、ユーザー向けAPIは低遅延重視、内部処理は高スループット重視と分けたほうが、コスト構造を素直に最適化しやすいです。
限界
FlexGen には明確な限界もあります。まず最大の注意点は、低レイテンシ用途向けではないことです。論文自体が throughput-oriented generative inference を対象にしており、何千秒という総レイテンシを許容して大きなバッチを流す前提です。リアルタイムチャットにそのまま使う設計ではありません。
次に、CPU メモリや SSD を大量に使うため、GPU が安くてもシステム全体ではそれなりのマシン構成が必要です。単一 16GB GPU で動くとはいえ、論文の代表設定は 208GB RAM と 1.5TB SSD を使っています。したがって「手元PCで巨大モデルが軽快に動く」と誤解しないほうがよいです。
また、I/O と量子化に強く依存するため、ストレージ帯域や CPU 側の性能が弱い環境では期待どおりに伸びない可能性があります。ポリシー探索である程度吸収できても、ハードウェア制約そのものは消えません。
実装面でも、単なるモデルロードよりかなり複雑です。テンソル配置、先読み、書き戻し、圧縮、スケジューリングが絡むため、保守性やデバッグ性は下がりやすいです。プロダクトに取り込むなら、推論品質だけでなく運用の複雑さも評価する必要があります。
さらに、論文の比較対象は当時の DeepSpeed Zero-Inference、Accelerate、Petals であり、現在の推論サービング実装と完全に同じ土俵で比較した結果ではありません。この点は論文を読む際の注意点で、FlexGen の価値は「当時の最終勝者かどうか」より、「メモリ階層を前提に高スループット推論を設計する考え方を示したこと」にあります。
よくある質問
Q. FlexGen はチャットボットの応答速度を上げる技術ですか？
A. 主目的はそこではありません。FlexGen は1件の応答を速くするより、限られたハードウェアで大量の推論をまとめて回す技術です。チャットのような低遅延用途より、評価、夜間バッチ、前処理のほうが適しています。
Q. 単一GPUで 175B モデルを動かせるなら、誰でも簡単に使えますか？
A. そこまで単純ではありません。論文では単一 16GB GPU を使っていますが、同時に 208GB の CPU メモリと 1.5TB の SSD を使っています。GPU 以外の資源もかなり重要です。
Q. FlexGen の本質は量子化ですか？それともオフロードですか？
A. 本質は両方ですが、中心はオフロード方針の最適化です。4bit 量子化は重みと KV キャッシュの転送量を減らしてくれますが、それだけで性能が出るわけではありません。配置、ブロックスケジュール、重ね合わせ実行が組み合わさって効きます。
Q. RAG システムでも使えますか？
A. 本番応答の低遅延化にはそのまま向かないことが多いですが、RAG の前処理や評価には十分使えます。大量文書の要約、メタデータ作成、検索対象の正規化のようなオフライン処理とは相性がよいです。
Q. この論文から今でも学ぶ価値はありますか？
A. あります。比較対象や実装環境は時代とともに変わりますが、重み・活性・KVキャッシュを分けて考え、GPU/CPU/SSD の全体で推論を設計する発想は今でも有効です。特に、資源制約のある環境で何を優先すべきかを考える材料になります。
今日の学び
この論文は、GPU メモリに入り切らない巨大 LLM を、限られた計算資源でどう実用的に動かすかという課題を扱いました。そこに対して、GPU・CPU・SSD へのオフロード配置、ブロックスケジュール、I/O の重ね合わせ、4bit 圧縮を組み合わせた FlexGen で解こうとしています。
ここから得られるヒントは、AI プロダクトの価値はモデルそのものだけでなく、どのハードウェア階層をどう使うかという実行設計でも大きく変わるということです。大きなモデルを低遅延で出せなくても、前処理や評価の裏側で使う設計なら、まだ活かせる余地は十分あります。