Paperspace Gradientでtext-generation-webuiを起動させてRWKVを動かす

2023年4月26日 14:17

クラウド上でGPUを使った開発（検証）環境を利用したい場合に一番メジャーなのはGoogle Colaboratryですが、課金プランが従量課金制になったのを受けての移住先として注目されていたのがPaperspace Gradientです。

Paperspace Gradientの使い方については以下の記事に詳しいのでご参考ください。

私自身はGrowthプランを契約しています。月額39ドルでA100 80GBが使い放題という破格の価格設定に惹かれてのことだったのですが、ハイスペックGPUは用意されているリソース自体が少なく取り合い状態になっているため、実際に利用できているのはA6000 48GBあたりです。

とはいえ実物を買うよりはだいぶ安く収まっていますし、検証用途であれば今のところ十分なスペックなのでPaperspace Gradientを愛用しております。

Notebookのセットアップ

そんなPaperspace Gradientで様々な言語モデルを試すためにtext-generation-webuiを起動させてみたので、そのセットアップについてメモしておきたいと思います。

text-generation-webuiはAUTOMATIC1111/stable-diffusion-webuiの言語モデル版のようなツールで、Gradioを用いたUIで言語モデルを操作することが可能です。

Gradientへの会員登録は完了している前提で、Notebookの作成方法から記載していきます。

まずはランタイムを選びます。Start from ScratchでOKです。

GPUは無料で使える範囲内で強めのものを選んでおけば良いと思います。料金が記載されているものについては、時間あたりいくらで従量料金が発生します。

自動シャットダウン時間も選択できますが、最大の6時間を選択しておけば良いかと思います。有料プランに加入していても、従量料金の発生しないGPUを利用する場合は連続6時間までの起動になります。

text-generation-webuiのセットアップ

Notebookが起動したら直下に.ipynbファイルを作成します。

Paperspace Gradientは/tmpディレクトリ以外の場所にデータを保存すると課金対象になります。Growthプランでは50GBまでは自由に使える容量となっていますが、言語モデルを扱う上で50GBはかなり小さい容量です。

/tmpディレクトリに保存すると、再度マシンを立ち上げたときに/tmpに保存したデータが削除されてしまいますが、毎回Notebookを全実行すればtext-generation-webuiが起動できるようにコードを書いて行こうと思います。

まずは /tmp ディレクトリ以下にtext-generation-webuiのリポジトリをクローンします。

# /tmpディレクトリで作業する
%cd /tmp

# 1. text-generation-webuiのダウンロード
repository_url = "https://github.com/oobabooga/text-generation-webui.git"
git_branch = "main"

!git clone {repository_url}
%cd text-generation-webui
!git checkout {git_branch}
!pip install -r requirements.txt

上記のコードを実行すると、/tmpディレクトリ以下にtext-generation-webuiディレクトリが作成され、必要なライブラリがインストールされます。

（立ち上げ環境によってはPyTorch 2.0が事前にインストールされていてエラー表示が出る可能性がありますが、無視してもWebUIは立ち上がります）

次に言語モデルをダウンロードします。今回はRWKV-4-Raven-14Bを動かそうと思うので、以下のコードでリポジトリの /models ディレクトリ内にデータをダウンロードします。

%cd /tmp/text-generation-webui

## RWKV

# 使用前にTokenizerをダウンロード
!curl -o models/20B_tokenizer.json -O https://raw.githubusercontent.com/BlinkDL/ChatRWKV/main/v2/20B_tokenizer.json

# RWKV-4-Raven-14B
!curl --location -o models/RWKV-4-Raven-14B-v8-EngAndMore-20230408-ctx4096.pth -O https://huggingface.co/BlinkDL/rwkv-4-raven/resolve/main/RWKV-4-Raven-14B-v8-EngAndMore-20230408-ctx4096.pth

容量がでかいので気長に待ちます・・・。

ダウンロードが終わったら以下のコードでWebUIを起動します。

%cd /tmp/text-generation-webui

!python server.py --listen --share --no-stream --model RWKV-4-Raven-14B-v8-EngAndMore-20230408-ctx4096.pth

--model引数には先ほどダウンロードしたモデルのファイル名を指定します。

--shareをつけないとPaperspace Gradient上で起動したサーバーにアクセスできないので注意です。

ロードが完了すると、以下のようにPublic URLが割り振られます。

起動時の画面はこんな感じ。

プロンプトを投げたところ、15秒ほどで応答が返ってきました。

Output generated in 14.89 seconds (13.44 tokens/s, 200 tokens, context 32, seed 1770166076)

CPUをゴリゴリに使っている感じです。

WebUI内の画面

パラメータタブでは非常に多くのパラメータをチューニングすることができます。

モデルを引数につけずにサーバ起動すると、modelsディレクトリ内にあるモデルが全て読み込まれ、この画面でモデルの切り替えを行うことができます。ただ、起動が重くなるだけなので、引数に指定して起動した方が良いかと思います。

トレーニングのUIもありますが、こちらはまだ使ったことがないです。

起動時のパラメータを設定して再起動することができるタブもありますが、Notebookから起動しているのでこれは上手くいきません。

RWKV参考記事

所感

ここまで紹介しておいてなんですが、画像モデルと違って言語モデルはUIがあってもあまり便利じゃないかも・・・。パラメータやプロンプトに対する応答などはNotebookに履歴として残っていた方が後から参照しやすいですし。

とはいえ、言語モデルで遊んでみるとっかかりには良いかなーと思いました。

現場からは以上です。