🕶️
FlexGenをGoogle Colaboratoryで動かす
最近1GPUで大規模な言語モデルを動かせるFlexGenというライブラリが話題になりました。
ただ、2023年2月28日現在強烈なリファクタリングが進んでおり、世間的なブログ記事を読んでもそのまま動かせるものが少ないです。そのためこの記事で手順をまとめてみます。
基本的にはColaboratoryにコードを貼っていけば実行できるはずです。
Google Colaboratoryのセットアップ
Google Driveのマウント
モデルファイルを保存するために、Google Driveをマウントします。
（容量に気をつけてください。けっこう食います。）
# Google Driveのマウント
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にキャッシュされるようにします
# 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'
FlexGenのインストール
FlexGenをインストールします。
最新のものはリファクタリング中のものもあるため、特定のcommitに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 .
実行する
Google Colaboratory Pro+のGPU(プレミアム)+ハイメモリ設定で動作させます。
ベンチマークの実行
軽量なOPT1.3Bでのベンチマークです。
!python -m flexgen.flex_opt --model facebook/opt-1.3b
チャットボットの起動
apps/chatbot.pyを実行することで、実際に対話形式でチャットボットを使えます。
OPT-6.7Bは少し利口ですが、文章が崩れていることもあります。
# OPT6.7Bの実行
!python apps/chatbot.py --model facebook/opt-6.7b
OPT-30Bはだいぶ利口ですが、日本語はまだ不十分です。
以下の設定にすると、モデルの重みを圧縮しつつ、RAMではなくDisk上にモデルをオフローディングして実行します。
# OPT30Bの実行
!python apps/chatbot.py --model facebook/opt-30b --percent 0 100 100 0 100 0 --compress-weight
さらにこれの上位モデルとしてfacebook/opt-66bというモデルも指定できますが、Google Colaboratoryでは動かせないようです。（処理中に止まります）
最後に
手元で大規模言語モデルが動くと、かなり身近になったなという気がしますね。
Discussion