GPT-NeoX-20BをRTX 3090で動かして日本語の問題を解かせてみた。

# はじめに

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になります。

https://qiita.com/kekekekenta/items/d0d64688c889b860beba

これで、Windows 11 + RTX 3090(24GB) + CUDA 12.1 + Docker Desktop の環境を用意できました。

## GPT-NeoX-20B をダウンロードするために Git LFS をインストール
Git for Windows を Windows11 にインストールします。既にインストールしている方はインストールする必要はありません。

https://git-scm.com/downloads

次に、Gitを拡張する、Git Large File Storageをダウンロードします。こちらも既にインストールしている方はインストールする必要はありません。

https://git-lfs.com/

ダウンロードしたファイルを実行した後は、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コンテナにアタッチします。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/32580/210c8691-14c1-4682-6947-e64a3322f604.png)

アタッチしたら、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を作成します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/32580/fc86a1c5-507d-3982-65ff-140e98892752.png)

Notebook を開いたら、Ctrl + Shift + P で「Notebook: Select Notebook Kernel」を選択して「/usr/bin/python」を選択します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/32580/dbca1afe-ecf9-cb69-ef0e-5074001ea942.png)

## Notebook で GPT-NeoX-20B モデルを GPU に読み込む
作成したNotebookのファイルと同じディレクトリに、git cloneでダウンロードしたgpt-neox-20bのディレクトリがあることを確認して以下を実行します。

```python
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 ぐらい使っている状態）

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/32580/b728a8fa-1839-6453-91c6-04958ec19a42.png)

## 係り受けの問題を解いてみる
以下のプログラムとプロンプトを実行します。「ブラックホール」と答えてくれるといいのですが…。

```python
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. 太陽

答え: 正解はブラックホールです！
```

## 照応の問題を解いてみる
次に、以下のプログラムとプロンプトを実行します。答えは「火星」です。

```python
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. 生命

答え: 火星
```

## 同義文判定の問題を解いてみる
最後に、以下のプログラムとプロンプトを実行します。答えは「同じである」です。

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

https://note.com/npaka/n/n09aeb00bf72b