【解消】rinna のモデルを使おうとして OutOfMemoryError にはまった話

はじめに

前回の記事の続き。

pancokeiba.hatenablog.com

Windows での環境構築に挫折したため WSL2 で構築することとした。今回はその作業内容と rinna のサンプルコードの実行結果について記載する。

実行環境

• PC
  • Windows 11 Home
  • CPU intel Core i5-14500
  • メモリ 16GB
  • GPU NVIDIA GEFORCE RTX 4060
    • VRAM 8GB
    • Compute Capability 8.9
    • Micro-architecture Ada Lovelace
    • Driver version 551.76
• WSL 関連
  • WSL 2.2.4.0
  • Ubuntu 22.04.3 LTS
  • VSCode 1.92.2
  • conda 環境の諸々
    • conda 24.5.0
    • python 3.11.3
    • pytorch 2.4.0
    • pytorch-cuda 12.4
    • sentencepiece 0.2.0
    • transformers 4.41.2
• 汎用言語モデル rinna/bilingual-gpt-neox-4b

環境構築

実施した作業を時系列に沿って記載する。

WSL2 の導入

過去に導入済みだったが、なぜかネットワークが遅すぎてコマンド実行に数十秒かかる状態だったため、アンインストール・インストールを実施した。
以下のサイトを参考にさせていただいた。

WSL2 のインストールとアンインストール #初心者 - Qiita

CUDA toolkit のインストール

公式サイトのインストールコマンドを実行した。

CUDA Toolkit 12.6 Update 1 Downloads | NVIDIA Developer

cuDNN のインストール

公式サイトのインストールコマンドを実行した。

cuDNN 9.3.0 Downloads | NVIDIA Developer

動作確認

nvidia-smi

上記コマンドを実行すると結果が返ってきた。たぶんこれで OK なはず。

nvcc -V

上記コマンドを実行すると Command not found のエラーが出たため、~/.bashrcに環境変数の設定コマンドを追記した。

~/.bashrc

export PATH=/usr/local/cuda:/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/lib:/usr/local/cuda/lib64:$LD_LIBRARY_PATH

~/.bashrcの保存後、WSL を再起動して環境変数の設定は完了。
再度 nvcc コマンドを実行すると結果が返ってきたため OK。

Anaconda のインストール

WSL でも conda 環境で実行したかったため、Anaconda を導入した。
導入手順はこちらを参考にさせていただいた。

【備忘録】WSL2にanaconda3をインストールしてcondaを使えるようにするまで #Python - Qiita

conda 環境の作成

conda の仮想環境を作成した。手順は割愛する。

Pytorch のインストール

公式サイトからインストールコマンドをコピペして実行した。

Start Locally | PyTorch

VSCode の導入

ソースを書く・ターミナルを操作する・GitHub に上げるなどの作業を VSCode に集約したかったため導入した。と言っても、Ubuntu にはインストールされているようで、codeコマンドを実行するだけで VSCode が起動された。

VSCode 起動後は、画面左下のアイコンから Connect to WSL を選択すれば OK。

資材格納

以下のディレクトリ構成で資材を格納した。

• bilingual-gpt-neox-4b 配下
  • 手動でダウンロードした rinna モデルのファイル一式を格納した
• rinna_test.py
  • サンプルコードが書かれた Python ファイル

.
├── bilingual-gpt-neox-4b
│   ├── README.md
│   ├── config.json
│   ├── gitattributes
│   ├── model.safetensors
│   ├── pytorch_model.bin
│   ├── rinna.png
│   ├── spiece.model
│   ├── spiece.vocab
│   ├── tokenizer_config.json
└── rinna_test.py

サンプルコードを実行

以下で公開されているサンプルコードを一部変更して実行した。

rinna/bilingual-gpt-neox-4b · Hugging Face

モデルのパスの記載部分を変更し、bilingual-gpt-neox-4b フォルダに格納されたモデルを参照するようにした。

変更後のコード

tokenizer = AutoTokenizer.from_pretrained("./bilingual-gpt-neox-4b", use_fast=False)

サンプルコード実行エラー

Killed

初めてのエラーに遭遇した。Killed と表示されてそれ以外の情報がない。

調べた限り、メモリ不足の際に Killed が表示されるようだが、詳細を確認するために、dmesgコマンドを実行した。

実行したコマンド

dmesg -T

コマンド実行結果

[Sun Sep  1 14:48:46 2024] Out of memory: Killed process 8833 (pt_main_thread) total-vm:52131812kB, anon-rss:6358652kB, file-rss:0kB, shmem-rss:0kB, UID:1000 pgtables:23936kB oom_score_adj:0

たしかにサンプルコードを実行した時間帯に、Out of Memory によってプロセスがキルされていた。今のサンプルコードだと、自分の実行環境ではメモリ不足でキルされてしまうということが分かった。

OOM Killer への対処と結果

キルされないようにサンプルコードを修正した。具体的には、torch_dtypeとdevice_mapの設定値を追加した。

修正後の rinna_test.py

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# メモリ使用状態を取得する
# torch.cuda.memory._record_memory_history()

# rinnaモデルをローカルパス指定
tokenizer = AutoTokenizer.from_pretrained("./bilingual-gpt-neox-4b", use_fast=False)

# 変更前 → Killed
# model = AutoModelForCausalLM.from_pretrained("./bilingual-gpt-neox-4b")
# 変更後 → OK
model = AutoModelForCausalLM.from_pretrained("./bilingual-gpt-neox-4b", torch_dtype=torch.float16, device_map='auto')

# if torch.cuda.is_available():
#    model = model.to("cuda")

text = "西田幾多郎は、"
token_ids = tokenizer.encode(text, add_special_tokens=False, return_tensors="pt")

with torch.no_grad():
    output_ids = model.generate(
        token_ids.to(model.device),
        max_new_tokens=100,
        min_new_tokens=100,
        do_sample=True,
        temperature=1.0,
        top_p=0.95,
        pad_token_id=tokenizer.pad_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id
    )

output = tokenizer.decode(output_ids.tolist()[0])
print(output)

# メモリのスナップショットを取得
# torch.cuda.memory._dump_snapshot("my_snapshot.pickle")

rinna_test.py 実行結果

T5Tokenizer に対するメッセージが出ているが、プログラムは正常終了し、期待する結果が出力された。

You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
Some parameters are on the meta device device because they were offloaded to the cpu.
西田幾多郎は、明治初期の物理学の第一人者であり、理神論者としても知られています。
彼は、万物は全て霊であり、その本質は、永遠に変わらないものと信じてやまなかったのです。 (以上「万物の霊の本質と、人心の科学」より)
「神がかり」が、なぜ、霊能者に、必要なのか、といえば、それこそ、神秘の世界に神を思うのが、真の宗教だからです。 霊の存在は、宗教そのものです。
宗教には、いろいろな

rinna_test.py の実行時間について

Out of Memory が発生したことから、自分の実行環境はサンプルコードの実行に耐えうるスペックではなさそうということが分かった。参考程度に実行時間について記載しておく。
測定には、time.perf_counter()を使用した。5 回実行したが、どれも 30 秒近くかかっていることが分かる。

1 回目: 29.19523430600384
2 回目: 27.53580800799682
3 回目: 28.29200663699885
4 回目: 28.88365111400344
5 回目: 28.652095464000013

プログラム実行中は、メモリも GPU もかなり使っていた。16GBのGPUが必要なのか・・・

参考資料

WSL2 のインストールとアンインストール #初心者 - Qiita
WSL2 上の PyTorch に GPU を認識させて深層学習環境をつくる
anaconda - CUDA Toolkitインストール時に発生するnvcc missingエラーについて - pytorch
【備忘録】WSL2にanaconda3をインストールしてcondaを使えるようにするまで #Python - Qiita
linux - PyTorch code stops with message "Killed". What killed it? - Stack Overflow