LLM推論にGPUは1つでいい？FlexGenについてまとめてみた

FlexGenとは

FlexGenは2023年2月にリリースされた、LLMの処理を行うための生成エンジンです。FlexGenには大きな特徴が2点あります。

• GPUを1つ積んだPCでLLMの処理が行えるようになる。
• 高いスループットを持つ。スループットとは、FlexGenの論文上では生成スループットと呼んでおり、入力として与えられた文章やフレーズから生成されるトークン数(単語や句読点などの意味を持つ最小単位)を、入力テキストを処理しトークンを生成するまでの時間で割った値と定義しています。つまり、1秒あたりどの程度のトークンを生成できるかをスループットとしています。

特に一点目のインパクトが強いです。パラメータ数が特に多いLLMの処理は、GPUを並列したPCや大規模なレンタルサーバー・クラウド上で行うことが主でした。しかし、FlexGenによって一般的なPCでもそのようなLLMの処理を行えるようになりました。

FlexGenのメカニズム

FlexGenの主なメカニズムとして以下の3つを解説していきます。

1. GPU、CPUメモリ、ディスクを利用した分散処理 (オフローディング)
2. LLMの重み、キー、バリューを4ビット整数に圧縮する処理 (量子化)
3. 従来の処理手法と異なるジグザグ処理

1. GPU、CPUメモリ、そしてディスクを利用した分散処理 (オフローディング)

FlexGenはPCの主なリソースであるGPU、CPUメモリ、そしてディスクを使ってLLMを処理することができます。このような分散処理をオフローディングと呼びます。利用するリソースはよく聞くPCリソースであると思いますが、簡単にそれぞれのリソースについてまとめます。

• GPU: 画像の処理を行う装置です。行列計算が得意であり、ディープラーニングを行う際にも利用されます。
• CPU: コンピュータの頭脳に当たります。コンピュータ内の計算処理や制御を行います。
• メモリ: CPUが一時的にデータを記録するための装置です。例えば、コンピュータ上で文章や画像をコピーした際、コピーされたデータはメモリに保存されていることになります。CPUに保存されたデータは、コンピュータの電源を落とした時に保持されず消去されます。
• ディスク: データを保管する装置になります。HDDやSSDがディスクに当たります。例えば、画像やドキュメント等をダウンロードした際、これらはディスクに保存されていることになります。ディスクに保存されたこのファイルは、コンピュータの電源を落としても保持されます。

LLMはGPUに全て処理させる方法が一般的です。しかし、LLMの多くは莫大なメモリ量を必要とします。例えば、Meta社からリリースされたOPT-175B*3というモデルをGPUのメモリにロードしたい時、必要なメモリ量は325GBとなります。このモデルを利用するには、例えばNVIDIA A100 (80GB) といった高性能なGPUが5枚必要となります。しかしFlexGenではオフローディングを採用しているため、LLMのサイズがGPUのメモリを超えていたとしても、LLMの処理を行うことが可能です。

2. LLMの重み、キー、バリューを4ビット整数に圧縮する処理 (量子化)

LLMには重み、キー、そしてバリューというパラメータが含まれます。これらは、入力された文章に対する回答を出力する際に必要となるパラメータです。FlexGenはこれらのパラメータの精度を維持しつつ、4ビットの整数に圧縮する手法*4を用いています。このように、データやパラメータをより小さいビット数に圧縮することでデータ容量を下げる手法を量子化と呼びます。

この手法を用いた場合、出力に悪影響を及ぼすことが危惧されます。FlexGenの論文では、OPT-175Bの重み、キー、バリューを4ビットに圧縮した場合と圧縮しない場合の出力結果の比較を行っており、圧縮の影響は無視できるレベルであることが示されています。さらに、3ビット量子化も試験しています。しかし、3ビットまで圧縮すると精度を保てないことが示されています。

3. 従来の処理手法と異なるジグザグ処理

FlexGenのもう一つの特徴として、従来の生成エンジンの処理の順序と異なる順序を用いることで、PCにかかる負担を減らしています。

図1は従来の生成エンジンの処理の順序(a)と、FlexGenの処理の順序(b)を表します。この図は、ある入力テキストからトークンを3つ生成する状態を考えています。それぞれの色の四角はGPUの層とバッチごとの計算を示します。
従来の方法 (図1 (a)) は、全トークンのバッチを行ごとに処理を行います。この手法のメリットは、1バッチ分の計算を最も速く終わらせることができることと、行単位で保持する必要があるキー・バリュー等のパラメータの容量をバッチの計算終了後に解放できることです。したがって、行の計算を完了させメモリを解放し、次の行へ移動するというプロセスを繰り返し行う手法となります。しかし、行方向に連続する2つのバッチは重みを共有しないため、この手法では重みを繰り返しロードする必要があり、PCに負担がかかります。

このロードの繰り返しを削減するために、FlexGenは複数行のバッチ (ブロック) を列ごとに処理します(図1 (b))。列内のすべての正方形は重みを共有しているため、重みをGPUに残して再利用することができます。しかし、行ごとのパラメータの容量を保持する必要があるため、列の最後まで処理を行おうとすると全行のパラメータをメモリに残す必要が出てきます。そのため、CPUメモリとディスクを溢れない程度の行数に絞って処理を行います。FlexGenの論文上ではこの手法をZig-zag block scheduleと呼んでいます。この手法により、PCにかかる負担を減らしています。

以上の内容が、FlexGenのメカニズムとなります。

FlexGenと他生成エンジンのベンチマーク比較

LLMに利用できる生成エンジンはFlexGen以外も存在します。オフローディングを採用している生成エンジンの例として、DeepSpeed Zero-Inference*5と Hugging Face Accelerate*6があります。これらは、お手元のPCでLLMを処理したい際に利用する生成エンジンの候補となることが考えられます。

では、FlexGenは他の生成エンジンと比較してどの程度優れているのでしょうか。FlexGenの論文では、FlexGen(圧縮を使う場合(c)と使わない場合)、DeepSpeed Zero-Inference、そして Hugging Face Accelerate の生成スループットとレイテンシ (処理を開始してから応答までかかる時間) を比較しています (図2)。なお、比較した際に用いられたコンピュータのリソースは、GPU: NVIDIA T4 (16GB)、メモリ: DRAM 208GB、ディスク: SSD 1.5TB となっています。

OPT-175Bの結果が顕著であり、同等のレイテンシである場合、既存の生成エンジンよりFlexGenのスループットが向上していることがわかります。複数データ点がある理由としては、用いたオプションや生成しようとしたトークン数等が異なるためです。圧縮を利用し、ディスクのオフロードを取り除いた場合、スループットは最も高くなり既存の生成エンジンの最大スループットの約100倍ほど向上します (青線、レイテンシが2¹²秒 (~ 4000秒))。

つまり、利用したコンピュータリソース要件やモデル、そしてスループットにおける比較では、FlexGenが最も優れていると言えます。

まとめ

今回は、大規模言語モデルを1つのGPUを搭載したPCで扱えるFlexGenという技術について紹介させていただきました。
FlexGenのメカニズムとして、オフローディング、パラメータの量子化、そして従来とは異なるLLMの処理方法を採用していることがわかりました。
また、FlexGenの性能として、LLMであるOPT-175Bを用いた際に、スループットが従来の生成エンジンに対し最大100倍まで向上することがわかりました。

参考文献

arxiv.org

github.com