7
6

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

記事投稿キャンペーン 「2024年!初アウトプットをしよう」

無料枠のGoogle ColabでGPT-3.5を超えるLLMを動かす!Mixtral offloadingについて

Posted at

はじめに

2023年12月DeepMindやMeta出身の研究者が設立したMistral AIが、大規模言語(LLM)モデル「Mixtral 8x7B」をリリースしました。多くのベンチマークでGPT-3.5やLlama 2 70Bを上回る性能を持つと言われています。

Mixtral 8x7BはMixture-of-Experts(MoE)モデルを採用しています。
MoEとはモデルの中に複数のブロック(Expert)を用意しておき、入力によって適したブロック(Expert)のみを用いて推論をする手法です。Mixtral 8x7Bでは8つのExpertから2つのExpertのみを使用して推論します。

上記のMixtral 8x7Bは、通常の実装だと8つのExpertをGPU上に置いておく必要があるのでGPUのコストがかかります。

この問題を解決すべく、限られたGPU環境(例:無料枠のGoogle ColabのGPU)でMixtral 8x7Bを動かす手法を紹介した論文「Fast Inference of Mixture-of-Experts Language Models with Offloading」を紹介します。

  • arXiv

  • GitHub

すぐ動かしたい方はこちら

内容

本論文では以下3つの手法でMixtral 8×7Bの効率化をしています。

  • Expertのキャッシュ保持
  • Expertのプリロード
  • 量子化

Expertのキャッシュ保持

下図はExpertがどのように使われているかを示した図であり、縦軸がExpertで横軸が入力のトークンです。青色で示したExpertは実際に使われているExpertです。(灰色は後述)

スクリーンショット 2024-01-05 11.56.24.png

図より、各入力トークンによって異なるExpertが使われていることがわかると同時に、隣り合ったトークンで同一のExpertを使っていることが多いことがわかる。

このパターンを利用するために、本論文ではアクティブなExpertを将来のトークンのための「キャッシュ」としてGPUメモリに保持します。これにより同じエキスパートが再びアクティブになれば、即座に利用できるようになります。

キャッシュの保持の仕方としては、k個の直近使用されたExpertを常に保持するというLRU キャッシュを採用しています。

図の灰色の部分は、k=2の保持されたキャッシュを表しています。

スクリーンショット 2024-01-05 11.55.35.png

上図はLRU キャッシュのヒット率を表していてk=2,3で大体5割程度キャッシュがヒットしていることがわかります。

Expertのプリロード

上記キャッシュによってExpertのローディング時間を短縮できますが、依然キャッシュで対応できないExpertはローディングの必要があり、推論時間の大半を占めています。

MoEでは、前の層の最終出力を使って次回用いるExpertを選択することから、事前に確実にExpertをプリフェッチすることは不可能です。

しかし、何らかの手段で次のExpertを推測できればExpertをプリロードすることができます。(投機的なロード)

本論文では、次の層のゲーティング関数を前の層の隠れ状態に適用することで、次の層のExpertを正確に推測できることを発見しました。

スクリーンショット 2024-01-05 12.20.48.png

上図はプリロードのリコール(取りこぼしなくロードできたか)を示しており、実線が1層前の出力、点線が2層前の出力、ドット線が10層前の出力からExpertを予測しています。

層が手前であればあるほど、よりモデルの最終出力に近くなるので推測が容易になるためリコールが高くなります。(その分、並列処理内でプリロードできる時間が限られてしまいます。)

どの程度プリロードできるかはGPUサイズによりますが、プリロードできるExpertが1個(緑線)かつ2層前の出力を使ってもリコールが4,5割程度で、ランダムにプリロードする期待値よりは高いので効果がありそうです。

期待値の計算の仕方は2通りあって、キャッシュを保持するかしないかによって計算方法が異なります。

プリロードできるExpertを1個としたとき、
(i)キャッシュを保持しない場合だと、2/8 = 25%が期待値であり、(ii)k=2でキャッシュを保持する場合だと、キャッシュヒット率を50%として(2-1)/(8-2) = 16%が期待値となります。

つまり前の層の出力からのExpertの推測は、期待値よりも高い精度で推測できていることがわかります。(これも論文に書けば良いんじゃなかろうかと思っています。)

キャッシュの保持とプリロード1つ->キャッシュ保持をやめて、プリロードのエキスパートを3つにした方が精度が良く、時間を削減できそうですが、キャッシュ保持の方がメモリ的に嬉しそうです。(おそらく)

量子化

すべての非Expertを4bit、Expertをハーフ2次量子化(HQQ:Half Quadratic Quantization)(Badri & Shaji, 2023)で量子化しています。これによりMoEモデルの品質とサイズのトレードオフが改善されます。また1bit量子化では性能が著しく低下すると記述されています。

おわりに

いかがだったでしょうか。

2023年12月11日に公開されたMixtral 8x7bを、即座にGPUコストを小さくするという変更を加えて実験を行い、2023年12月28日に初稿をarXivへ出すという仕事の速さが怖いです。(裏で実はやっていたのかもしれませんが、)

手法自体はとてもシンプルであり、今後応用が十二分にできそうなので参考にしたいと思います。

7
6
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
7
6

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?