FlexGenをGoogle Colaboratoryで動かして、ChatGPTライクな大規模言語モデルで遊ぶ

2023.03.07Update

中村

N1! Machine Learning Product Engineerの中村です。
最近は自然言語処理がかなりのブームになっており、弊社でもNotion AIなどでブログ記事を書かせるような不届者が多数出ているとの噂です。自分も最近はOpenAIのAPIなどを検証しているのですが、大規模言語モデルは手元で動かせないモデルが多く存在します。（大量のGPUが必要など）

しかし、最近登場したFlexGenを使うことで、1GPUで大規模言語モデルを動かせるとの噂です。
今回はこのFlexGenを使ってMeta社のOPTという大規模言語モデルを動かしてみます。 https://github.com/FMInference/FlexGen
https://arxiv.org/abs/2205.01068

FlexGenとは

FlexGenは、大規模言語モデル（Large Language Model）をシングルGPUで動かすことを可能にする、高スループット生成エンジンです。

近年の大規模言語モデルはモデルの重みが公開されてもそれを動かすことが困難な場合が多いです。以下は公開されている大規模言語モデルですが、最大のモデルを使おうとする場合には複数枚の高性能GPUが不可欠です。

Meta社のOPT（175億パラメータ）
https://ai.facebook.com/blog/democratizing-access-to-large-scale-language-models-with-opt-175b/
BigScienceのBLOOM（176億パラメータ）
https://huggingface.co/bigscience/bloom
直近だとMeta社がLLaMa（65億パラメータ）を公開しています。
https://research.facebook.com/publications/llama-open-and-efficient-foundation-language-models/

このようなモデルを動かそうとした場合には、一般的にはGPUが複数必要であり、家庭で動かすのは多くの場合困難でした。そうなると大企業で潤沢な計算資源がない限り、ビジネス転用も難しいでしょう。

しかし、そこで登場したのがFlexGenです。大規模言語モデルをロードする場合には、ディスクなどにオフローディングし、家庭用のGPUでも動かせるように自動的にチューニングしてくれます。
一般にDiskよりもRAMの方が高速に動作するため、このような仕組みはレイテンシーが急激に低下しやすいですが、できる限り高速に動作した状態を維持するように設計されているようです。

以下では、実際にFlexGenをGoogle Colaboratory上で使用してみます。
この記事では以下のcommitを使用します。
（記事執筆2022/2/28現在、リファクタリングが急速に進んでおり、READMEの記述でさえもそのままCloneしても動かないことが多いため、記事の途中でgit reset –hardすることで対象のcommitを使用します）
https://github.com/FMInference/FlexGen/commit/f79b8950487d0af96eaf5ef5f75c1be15e24e9bf

Google Colaboratoryの環境を整備する

Google ColaboratoryはPro+版をGPU(プレミアム)+ハイメモリで利用します。
（いろいろな環境を使ってきましたが、試す場合にはColabが一番使いやすいと個人的には思います。特に昨今の技術の乱発状態では、Pro+に5000円払った方がいろいろ楽だなと思います。）

最近だとGPUもA100を引けますが、コンピューティングユニット数に注意してください。

Google Driveのマウント

使用するごとにモデルをダウンロードするのは大変なので、Google Driveにモデルをキャッシュするようにします。（容量に注意してください）

# Googleドライブのマウント

from google.colab import drive

drive.mount("/content/drive")

作業フォルダを作成

# 作業フォルダの作成

import os

os.makedirs("/content/drive/MyDrive/Colab Notebooks/FlexGen/work", exist_ok=True)

HuggingFaceのキャッシュパスを設定

途中でダウンロードするモデルをGoogle Driveにキャッシュするようにします。
（最新版のFlexGenだとこの環境変数が効かなかったため、何か変更があるのかもしれません）

# HuggingFaceのキャッシュパス

import os

os.environ['TRANSFORMERS_CACHE'] = '/content/drive/MyDrive/Colab Notebooks/FlexGen/work'

os.environ['HF_DATASETS_CACHE'] = '/content/drive/MyDrive/Colab Notebooks/FlexGen/work'

pipのアップグレード

一応pipをアップグレードします

1	!python -m pip install --upgrade pip

FlexGenのインストール

FlexGenをソースからインストールします。
最新版ではColabでうまく動かなかったため、特定のコミットにresetして実行します。

# FlexGenのインストール

%cd "/content/drive/MyDrive/Colab Notebooks/FlexGen/work"

!git clone https://github.com/Ying1123/FlexGen.git

%cd "/content/drive/MyDrive/Colab Notebooks/FlexGen/work/FlexGen"

!git reset --hard f79b895

!pip install -e .

ここまででColabの準備は完了です。

実際に動かしてみる

ベンチマークの実行

まず、OPT1.3Bという比較的小型の言語モデルでベンチマークを動かしてみます。

1 2	# デモの実行 !python -m flexgen.flex_opt --model facebook/opt-1.3b

prefill latency: 0.064 s prefill throughput: 31823.288 token/s

decode latency: 0.815 s decode throughput: 152.217 token/s

total latency: 0.879 s total throughput: 145.623 token/s

良い感じですね。145トークン/秒は良い感じです。（AIでしか文章を書けない人の顔）

モデルのキャッシュ場所を指定したので、Google Drive内にフォルダが作成されてそこにモデルがダウンロードされるはずです。

チャットボットの実行

では次はチャットボットを動かしてみます。 OPT-6.7Bという先ほどよりも少し大きめのモデルを動かしてみます。

1 2	# チャットbotの起動 !python apps/chatbot.py --model facebook/opt-6.7b

1-2分ほどでモデルのダウンロードが終わり、実際に入力できるようになります。

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: Hello?

Assistant: This is a bot. I cannot understand you.

Human: Very Sad

Assistant: I'm sorry. I can't understand you.

日本語訳

Human: こんにちは！あなたは何ができますか？

Assistant： AIアシスタントとして、質問に答えたり、チャットをすることができます。

Human：世界で一番高い山の名前は何ですか？

Assistant：エベレストです。

Human：もしもし？

Assistant：こちらはボットです。あなたのことを理解することはできません。

Human: とても悲しい

Assistant: 申し訳ありません。あなたを理解することはできません。

突然「あなたのことを理解することができない」というトーンになって悲しい気持ちになりました。

もっと強いチャットボットを使う

所詮小型のモデルに私の気持ちなんて理解できないのです。もっと大きなモデルを使ってみましょう。

1	!python apps/chatbot.py --model facebook/opt-30b --percent 0 100 100 0 100 0 --compress-weight

percentはオフローディングをどのように行うかを指定します。この場合はRAM上に全くモデルをロードせずに実行します。
compress-weightを指定することで重みが圧縮されます。（精度はやや悪化します）

だいたい1回の応答に10秒程度かかります

Human: Hello?

Assistant: What's up?

Human: I ate curry for lunch today.

Assistant: That sounds interesting.

Human: What are your hobbies?

Assistant: I love reading and watching movies.

Human: That's good. What kind of movies do you watch?

Assistant: I like romantic movies.

Human: Do you know the specific title name of the movies?

Assistant: I don't know all the movie titles, but I know the names of some of my favorite movies.

Human：もしもし？

Assistant：どうしたんですか？

Human: 今日のお昼にカレーを食べました。

Assistant：それは良さそうですね。

Human：趣味は何ですか？

Assistant：読書と映画を見るのが好きです。

Human：それはいいですね。どんな映画を見ますか？

Assistant：恋愛映画が好きです。

Human：なるほど。具体的な映画の題名はご存知ですか？

Assistant：すべての映画のタイトルを知っているわけではありませんが、好きな映画のいくつかは名前を知っています。

無難な回答を返してきたり、最後の回答は的を得ない表現ですが、比較的正しい回答をしているように思います。もう少しチューニングすれば十分使えるかもしれないですね。

日本語でチャットしてみる

Stable Diffusionが話題になったように、大規模自然言語処理モデルが家庭で動かせるとなるともっと話題になっても良さそうですが、日本だとあまり話題になっていないのはおそらく日本語のモデルがまだ少ないからでしょう。
Meta社のOPTも日本語を学習しているモデルではあるので、日本語でチャットをしてみましょう。

Colab上での日本語入力はEnterキーで送信されてしまうので、変換時は注意してください（n敗）

Human: こんにちは

Assistant: まんまよ。

Human: あなたが好きな食べ物はなんですか？

Assistant: あなたのお茶を飲みます。

Human: 日本語は話せますか？

Assistant: 日本語は話せます。

Human: 本当ですか？

Assistant: 本当です。

Human: あなたの趣味はなんですか？

Assistant: あなたの趣味は、日本語を話すことです。

Human: 日本語を話すことが趣味ってどういうことですか？

Assistant: 日本語を話すことが趣味です。

Human: 本当は日本語話せないですよね？

Assistant: 好きな日本語話せることが趣味です。

お茶を強奪されたり、日本語を喋るのが趣味だったり、おかしな文法を話すわけではないですが回答は変ですね・・・。このあたりの日本語が強い大規模言語モデルが登場すると一気に日本語の環境も変わってきそうな気配です。

もっと強いモデルを動かしたい

ちなみにさらに上位モデルのOPT-66Bもあるのですが、これはGoogle Colabでは動かないようです。（Loading checkpoint shardsという処理の途中で止まります・・・）

1	!python apps/chatbot.py --model facebook/opt-66b --percent 0 100 100 0 100 0 --compress-weight

OPT-30Bぐらいのサイズで日本語に特化したモデルが出れば大きなゲームチェンジャーになり得そうですね。

終わりに

今回はFlexGenをGoogle Colaboratory上で動作させてみました。
想像以上に手軽に動いてしまったので、ここからさらに大規模言語モデルが登場すると考えると、家庭用GPUや安価なインスタンスで動作させて、実際にプロダクション環境などに適用することも可能になるかもしれません。

自分も最近はGithub CopilotやNotion AIに頼りっぱなしですが、今後の自然言語処理界隈の動向にも注目ですね。

We are hiring!

ニフティでは、さまざまなプロダクトへ挑戦するエンジニアを絶賛募集中です！
ご興味のある方は以下の採用サイトよりお気軽にご連絡ください！

ニフティ株式会社採用情報

Tech TalkやMeetUpも開催しております！
こちらもお気軽にご応募ください！

Event – NIFTY engineering