はじめに

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 さんの記事を参考にしました。