M2 Mac で FlexGen を使ってみる

彗星のごとく現れた FlexGen を Apple Silicon M2 の Mac で使ってみたいと思います。

GitHub - FMInference/FlexGen: Throughput-oriented systems for large language models on commodity GPUs.
https://github.com/FMInference/FlexGen

この情報はFlexGenが発表されたばかりの2023å¹´3月10日現在の情報であるため、随時情報をアップデートしてください。

追記 : 下記の方法でするのがおすすめです。
大規模言語モデルOPTをM1/M2 Mac上のFlexGenで動かしてチャットする｜りん こうたつ｜note 

1. 必要環境を整える

※MacOS 13以上での実施推奨。
それより古いとtorch.cumsum の MPS サポートが無い為です。

Pythonやライブラリが既に入っている人は以下は必要ありません。

1-1. python のインストール

Homebrewで入れる方法ですが他の方法でも構いません。

$ brew install python

$ python3 -V

-> Python 3.11.2

1-2. 必要なライブラリのインストール

$ pip3 install --upgrade pip

$ pip3 install numpy

$ pip3 install attrs

$ pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cpu

pytorch (torch) をインストールするコマンドは下記から取得して pip を pip3 に変えています。

Start Locally | PyTorch
https://pytorch.org/get-started/locally/

PyTorch Build は Stable ではなく Preview (Nightly) を選択してください。

2. インストール

FlexGen を Git から clone してm1ブランチにきりかえてインストール。

$ git clone https://github.com/FMInference/FlexGen.git

$ git switch m1

$ cd FlexGen

$ pip3 install -e .

3. 試しに実行

小さい1.3bで実施してみます。

$ python3 -m flexgen.flex_opt --model facebook/opt-1.3b

<中略> minutes. If it seems to get stuck, you can monitor the progress by checking the memory usage of this process. Downloading pytorch_model.bin: 100%|█████████████████████████████████████████████| 2.63G/2.63G [03:41<00:00, 11.9MB/s] Fetching 1 files: 100%|████████████████████████████████████████████████████████████████| 1/1 [03:42<00:00, 222.28s/it] Convert format: 100%|███████████████████████████████████████████████████████████████████| 1/1 [00:06<00:00, 6.37s/it] <中略>

----------------------------------------------------------------------

0: Paris is the capital city of France. It is the most populous city in France, with an estimated population of 6,848,000 in 2016. It is the second most populous city

----------------------------------------------------------------------

3: Paris is the capital city of France. It is the most populous city in France, with an estimated population of 6,848,000 in 2016. It is the second most populous city

----------------------------------------------------------------------

TorchDevice: mps:0

cur_mem: 2.6426 GB, peak_mem: 0.0000 GB

TorchDevice: cpu

cur_mem: 2.6426 GB, peak_mem: 0.0000 GB

model size: 2.443 GB cache size: 0.398 GB hidden size (p): 0.008 GB

peak gpu mem: 0.000 GB projected: False

prefill latency: 2.763 s prefill throughput: 741.119 token/s

decode latency: 4.132 s decode throughput: 30.009 token/s

total latency: 6.895 s total throughput: 18.563 token/s

コマンドを実行すると、約2.6GBのモデルがダウンロードされます。
MacのGPUなので mps (Metal Performance Shaders) です。

4. 遊んでみる

他の人のブログとかを見るとあるはずの flexgen/apps/chatbot.py がない。
READMEにあるはずのサンプルコードも消えてる。
なんで？

Where is the chatbot? I miss it! · Issue #87 · FMInference/FlexGen https://github.com/FMInference/FlexGen/issues/87

理由はわかりませんが chatbot.py は消されてしまったようです。

過去のファイルより取得して実行できないかな？
試してみましたが、下記はダメでした....

$ cd flexgen/apps/

$ curl -O https://raw.githubusercontent.com/FMInference/FlexGen/9d888e5e3e6d78d6d4e1fdda7c8af508b889aeae/flexgen/apps/chatbot.py

$ python3 chatbot.py --model facebook/opt-1.3b

すると chatbot.py を復活させて且つ進化させてくださった方を発見！！！
deepl のAPIキーがあれば日本語でも動くようになってます！
最初からこっちでやればよかった！

con3office/FlexGen at m1
https://github.com/con3office/FlexGen/tree/m1

というわけで下記コマンドで無事に動作しました。

$ pip3 install deepl <-- chatbot.py でエラーにならないために入れます。

$ cd flexgen/apps/

$ curl -O https://raw.githubusercontent.com/con3office/FlexGen/m1/flexgen/apps/chatbot.py

$ python3 chatbot.py --model facebook/opt-1.3b

CUDA devices not available, mps:0 is used instead Initialize... A chat between a curious human and a knowledgeable artificial intelligence assistant. Human: Hello! What can you do? Assistant: As an AI assistant, I can answer questions and chat with you. Human: What is the name of the tallest mountain in the world? Assistant: Everest. Human: where is Tokyo? <-- Human: のあとに入力 Assistant: Tokyo. [4.08s] <-- これが回答

Human: What country is Tokyo located in? Assistant: Japan. [3.63s] Human: What is the name of the biggest lake in the world? Assistant: Lake Victoria. [4.07s] (カスピ海？)

Human:

※終了させたいときは ctrl ＋ C

ちょっと回答がおかしいときがありますが一番小さいモデルなのでそれはご愛嬌。

改善されたchatbot.pyのオプションの確認や、日本語でのやり取りもしたい場合は下記を参考にしてください。

大規模言語モデルOPTをM1/M2 Mac上のFlexGenで動かしてチャットする｜りん こうたつ｜note 

参考にしたリンク

• FlexGenで遊んだメモ｜ゆぬ｜note
  
  
• FlexGen のインストールと動作確認（大規模言語モデル，チャットボット）（Python，PyTorch を使用）（Windows 上） 
  
  
• 大規模言語モデルOPTをM1/M2 Mac上のFlexGenで動かしてチャットする｜りん こうたつ｜note 
  
  
• PyTorchがApple SiliconのGPUを使えるようになるらしいので試してみた - Qiita
  
  
• PyTorchをM1 MacBook のGPU(MPS)で動かす．実行時間の検証もしたよ