Google Colab で FlexGen を試す

「Google Colab」で「FlexGen」を試したのでまとめました。

【注意】「OPT-30B」のチャットデモを実行するには、「Google Colab Pro/Pro+」の「プレミア」の「ハイメモリ」と、300GBほどのGoogleドライブが必要です。

1. FlexGen

「FlexGen」は、限られたGPUメモリ (16GB T4 または 24GB RTX3090など) で大規模言語モデルを実行するための高スループット生成エンジンです。

大規模言語モデル (LLM) は、ChatGPTやCopilotなどのアプリケーションのコアになりますが、LLM推論の計算とメモリの要件が高いため、従来では複数のハイエンドアクセラレータを使用した場合にのみ実行可能でした。「FlexGen」は、LLM推論のリソース要件を1つのコモディティGPU (T4、3090など) にまで下げ、さまざまなハードウェアセットアップの柔軟な展開を可能にすることを目的としています。

「FlexGen」の主な機能は次のとおりです。

・ライトニング高速オフロード : 単一GPUで175Bモデルを実行する場合、他のオフロードベースのシステムよりも最大100 倍高速です。
・エクストリームコンプレッション : OPT-175Bなどのモデルのパラメータとアテンションキャッシュの両方を4ビットに圧縮し、精度の損失はほとんどありません。
・スケーラビリティ : より多くのGPUが与えられた場合にスケーリングできるように、分散パイプライン並列処理ランタイムが付属しています。

2. OPT-6.7Bのチャットデモの実行

Google ColabでのOPT-6.7Bのチャットデモの実行手順は、次のとおりです。
少なくとも15GB以上のGPUメモリが必要です。

(1) Googleドライブのマウント。

# Googleドライブのマウント
from google.colab import drive
drive.mount("/content/drive")

(2) 作業フォルダの作成と移動。
今回はマイフォルダ直下の「work」を作業フォルダとしました。

# 作業フォルダの作成と移動
import os
os.makedirs("/content/drive/My Drive/work", exist_ok=True)
%cd "/content/drive/My Drive/work"

(3) HuggingFaceのキャッシュパスの設定。
モデルを毎回ダウンロードするのは大変なので、Googleドライブに保存してます。

# HuggingFaceのキャッシュパスの設定
import os
os.environ['TRANSFORMERS_CACHE'] = '/content/drive/MyDrive/My Drive/work'
os.environ['HF_DATASETS_CACHE'] = '/content/drive/MyDrive/My Drive/work/Datasets'

(4) FlexGenのインストール。

# FlexGenのインストール
!git clone https://github.com/Ying1123/FlexGen.git
%cd FlexGen
!pip install -e .

(5) 「OPT-6.7B」のチャットデモの実行。

# チャットデモの実行
!python apps/chatbot.py --model facebook/opt-6.7b

パラメータは、次のとおりです。

・--model : モデル名 (facebook/opt-6.7b)
・--path : モデルの重みのパス (~/opt_weights)
・--offload-dir : テンソルのオフロード先ディレクトリ (~/flexgen_offload_dir)
・--percent : オフロード戦略の割合 ([100, 0, 100, 0, 100, 0])
・--compress-weight : 重みを圧縮するかどうか
・--compress-cache : キャッシュを圧縮するかどうか

(6) 以下のように表示されたらセリフを入力。

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:

「日本の有名なゲームキャラクターは？」に対して、「アシスタント：マリオ、リンク、ピカチュウ、カービィ。」と答えてくれました。

Human: What are some famous Japanese game characters?
Assistant: Mario, Link, Pikachu, and Kirby.

3. OPT-30Bのチャットデモの実行

Google ColabでのOPT-30Bのチャットデモの実行手順は、次のとおりです。
少なくとも64GBのCPU メモリが必要です。

(1) 「OPT-30B」のチャットデモの実行。
モデルのダウンロードに2時間ほどかかりました。

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

(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:

「日本の有名なゲームキャラクターは？」に対して、「アシスタント：マリオ、リンク、ピカチュウ、カービィ。」と答えてくれました。

Human: アメリカの大統領の名前は？
Assistant : トランプ。

4. オフロード戦略と--percentの設定

メモリ不足を回避するために、「--percent」の設定で、より多くのテンソルを CPU とディスクにオフロードする割合を調整できます。後日、自動ポリシーオプティマイザをリリースする予定ですが、今は手動で設定する必要があります。

「--percent」には、以下の6つの値を指定します。

・GPUの重みの割合
・CPUの重みの割合
・GPUのアテンションキャッシュの割合
・CPUのアテンションキャッシュの割合
・GPUのアクティベーションの割合
・CPUのアクティベーションの割合

5. メモリ不足時の対応方法

十分な GPU/CPU メモリがない場合は、以下を試してみてください。 より多くのメモリを節約できますが、実行速度は遅くなります。

・--compress-weight : 重みの圧縮の有効化
・--percent 0 0 100 0 100 0 : ウェイトをディスクにオフロード

Unityではじめる ROS・人工知能 ロボットプログラミング実践入門