はじめに
Azure OpenAIだと当たり前のように動くGPTモデルですが、ローカルの Windows11 環境の GPU でも動かしたいですよね。今回は、GPT-NeoX-20B を Windows11 の RTX 3090 で動かしてみましたので忘れないようにメモしておきます。
GPT-NeoX-20B を 8bit でロードすると、だいたい 22GB 使用するので RTX 3090(24GB) だとギリ動きます。
準備
Windows11 で GPT-NeoX-20B を動かす環境を整えます。
Windows11 に CUDA 12.x と Docker Desktop をインストール
Windows11 の DeepLearning 環境は、以前書いたこちらの記事を参考に環境を整えます。ただし、CUDA は最新の12.1をインストールしてください。11.xではなく、12.xになります。
これで、Windows 11 + RTX 3090(24GB) + CUDA 12.1 + Docker Desktop の環境を用意できました。
GPT-NeoX-20B をダウンロードするために Git LFS をインストール
Git for Windows を Windows11 にインストールします。既にインストールしている方はインストールする必要はありません。
次に、Gitを拡張する、Git Large File Storageをダウンロードします。こちらも既にインストールしている方はインストールする必要はありません。
ダウンロードしたファイルを実行した後は、Windows Terminalもしくは、コマンド プロンプトから以下のコマンドで有効にします。
git lfs install
GPT-NeoX-20B のダウンロード
Windows Terminalもしくは、コマンド プロンプトから以下のように、GPT-NeoX-20Bのモデルをダウンロードします。（私の Windows11 環境では、Dドライブの work ディレクトリにダウンロードしています。）
d:
cd \work
git clone https://huggingface.co/EleutherAI/gpt-neox-20b
NVIDIA の pytorch コンテナを実行
Windows Terminalもしくは、コマンド プロンプトから以下のように docker コマンドを使用して pytorch 23.05-py3 のコンテナを実行します。（Windows11 の Dドライブの work ディレクトリをコンテナの /work にマウントしています）
docker run --gpus all -it -v D:\work:/work nvcr.io/nvidia/pytorch:23.05-py3 bash
起動したら以下のpipコマンドで、Pythonのモジュールをインストールしておきます。
pip install transformers accelerate bitsandbytes
GPT-NeoX-20B を使ってみる
今回は、Windows11 上の Visual Studio Code の Notebook から GPT-NeoX-20B を動かしてみます。
Windows11 上の Visual Studio Code でコンテナをアタッチ
Windows11 上の Visual Studio Code を起動して、左下の「><」アイコンをクリックするか、Ctrl + Shift + P で「attach to Running Container」を選択して、実行中の pytorchコンテナにアタッチします。
アタッチしたら、Visual Studio Code の「File」メニューから「Open Folder」を選択して、GPT-NeoX-20Bをダウンロードした「/work/」ディレクトリを開きます。（Windows11 の Dドライブの work ディレクトリをコンテナの /work にマウント済みであることを確認）
Visual Studio Code で Notebookを作成
Ctrl + Shift + P で「Create:New jupyter Notebook」を選択して Notebookを作成します。
Notebook を開いたら、Ctrl + Shift + P で「Notebook: Select Notebook Kernel」を選択して「/usr/bin/python」を選択します。
Notebook で GPT-NeoX-20B モデルを GPU に読み込む
作成したNotebookのファイルと同じディレクトリに、git cloneでダウンロードしたgpt-neox-20bのディレクトリがあることを確認して以下を実行します。
from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
model_path = "./gpt-neox-20b"
tokenizer = AutoTokenizer.from_pretrained(model_path)
config = AutoConfig.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
model_path,
config=config,
device_map="auto",
load_in_8bit=True,
)
読み込むとタスクマネージャーの専用GPUメモリが 22.0GB/24.0GBぐらいになっていると思います。（専用GPUメモリは、Windows11 で動くアプリケーションも 0.8GB ぐらい使っている状態）
係り受けの問題を解いてみる
以下のプログラムとプロンプトを実行します。「ブラックホール」と答えてくれるといいのですが…。
prompt = """天の川銀河の中心には、太陽の400万倍程度の質量をもつブラックホールがあると推定されている。
この文脈において、天の川銀河の中心にあると推定されているのは何か、次の1から4から選択してください。
1. 天の川
2. 銀河
3. ブラックホール
4. 太陽
答え: """
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
input_ids = input_ids.to('cuda')
output_tokens = model.generate(
input_ids,
do_sample=True,
max_length=500,
min_length=100,
temperature=0.9,
num_return_sequences=1,
)
output = tokenizer.batch_decode(output_tokens)[0]
print(output)
少しすると以下のように文章を生成してくれます。正解ですね！自信満々に答えましたね！
天の川銀河の中心には、太陽の400万倍程度の質量をもつブラックホールがあると推定されている。
この文脈において、天の川銀河の中心にあると推定されているのは何か、次の1から4から選択してください。
1. 天の川
2. 銀河
3. ブラックホール
4. 太陽
答え: 正解はブラックホールです！
照応の問題を解いてみる
次に、以下のプログラムとプロンプトを実行します。答えは「火星」です。
prompt = """火星には、生命が存在する可能性がある。かつて大量の水があった証拠が見つかっており、現在も地下には水がある可能性がある。
この文脈において、かつて大量の水があった証拠が見つかっているのはどれか、次の1から4から選択してください。
1. 火星
2. 可能性
3. 地下
4. 生命
答え: """
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
input_ids = input_ids.to('cuda')
output_tokens = model.generate(
input_ids,
do_sample=True,
max_length=500,
min_length=100,
temperature=0.9,
num_return_sequences=1,
)
output = tokenizer.batch_decode(output_tokens)[0]
print(output)
おお！こちらも正解ですね！
火星には、生命が存在する可能性がある。かつて大量の水があった証拠が見つかっており、現在も地下には水がある可能性がある。
この文脈において、かつて大量の水があった証拠が見つかっているのは何か、次の1から4から選択してください。
1. 火星
2. 可能性
3. 地下
4. 生命
答え: 火星
同義文判定の問題を解いてみる
最後に、以下のプログラムとプロンプトを実行します。答えは「同じである」です。
prompt = """義経は平氏を追いつめ、ついに壇ノ浦でほろぼした。
この文が表す内容は、平氏は義経に追いつめられ、ついに壇ノ浦でほろぼされたということか。次の1から2の中から選択してください。
1. 同じである
2. 異なる
答え: """
input_ids = tokenizer(prompt, return_tensors="pt").input_ids
input_ids = input_ids.to('cuda')
output_tokens = model.generate(
input_ids,
do_sample=True,
max_length=500,
min_length=100,
temperature=0.9,
num_return_sequences=1,
)
output = tokenizer.batch_decode(output_tokens)[0]
print(output)
問題の出し方が悪いのか、不正解でした。残念。
義経は平氏を追いつめ、ついに壇ノ浦でほろぼした。
この文が表す内容は、平氏は義経に追いつめられ、ついに壇ノ浦でほろぼされたということか。次の1から2の中から選択してください。
1. 同じである
2. 異なる
答え: 異なる
さいごに
8bitにしてロードしているので少しは精度に影響があると思いますが、200億パラメーターのGPTが家の Windows11 で動くのは凄いですね！
今回試して分かったのは、Azure OpenAIだと工夫しなくても正解を出してくれますが、GPT-NeoX-20B だと正解を出せるようにプロンプトを工夫する必要がることです。安心して使う場合はAzure OpenAIのほうが良いと思います。
ではー。
参考記事
npaka さんの記事を参考にしました。