rinna 3.6bをdockerから使う

モチベーション

日本語に大規模言語モデル（LLM）を試してみたくて、5月にリリースされたrinnaを使ってみた。インストールの手間を省くため、dockerコンテナ環境下でrinnaを実行した。

その際にハマったこともあったので以下にまとめる。

情報源

rinnaのプレスリリース - rinnaについてはここを参照。36億パラメータということ、今回試したrinnaを以下、rinna 3.6bと記述する。
rinnaの紹介記事 - INTERNET Watchの記事。　サンプルコードもこの記事から流用した。
GPT日本語言語モデルを利用して映画「となりのトトロ」のストーリー概要を作成してみた（１本目） - Pytorchが入っているコンテナイメージに関する内容は、この記事を参考にした。

rinna 3.6bの実行結果

実行するpythonコード

情報源2.のサンプルコードを${HOME}/workspace/rinna.pyとして保存する。

サンプルコードの内容は次のとおり。5番目のmodelのコメントを外して有効にした。

import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", use_fast=False)
# 標準
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft")
# 自動
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", device_map='auto')
# 自動(VRAM16GB以下でも8GBはNG)
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", torch_dtype=torch.float16, device_map='auto')
# CPU指定
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft").to("cpu")
# GPU指定
model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft").to("cuda")
# GPU指定(VRAM16GB以下でも8GBはNG)
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", torch_dtype=torch.float16).to("cuda")


#繰り返し(抜けるのはCtl+c)
prompt = ""
while True:
    # 質問を入力
    question = input("質問をどうぞ： ")

    if question.lower() == 'clear':
        question = ""
        prompt = ""

    prompt = prompt+f"ユーザー: {question}<NL>システム: "

    # 時間計測開始
    start = time.time()

    token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")

    with torch.no_grad():
        output_ids = model.generate(
            token_ids.to(model.device),
            do_sample=True,
            max_new_tokens=128,
            temperature=0.7,
            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][token_ids.size(1):])
    output = output.replace("<NL>", "\n")

    # 時間表示
    end = time.time()
    print(end-start)

#    print(prompt)
    print(output)
    prompt = prompt+output+"<NL>"

NGC CATALOGからダウンロード、コンテナ実行

情報源3.を参考に、NGC CATALOGからPytorchコンテナをpull（ダウンロード）する。

Catlog > Containers > PytorchページのTagsを見ると、23.07-py3が最新なので、早速それをダウンロードして実行する。ユーザモードで動作するdockerについてはこのページを参考にして欲しい。

$ mkdir ${HOME}/workspace
$ docker pull nvcr.io/nvidia/pytorch:23.07-py3
$ docker run -d -it --gpus all --name gpt -v ${HOME}/workspace:/workspace/local nvcr.io/nvidia/pytorch:23.07-py3
$ docker exec -it gpt /bin/bash

RTX A4000での実行

docker環境下での操作

ここから上記で立ち上げたdockerコンテナでの操作する。

以下、#プロンプトから始まる行はコンテナ環境下での操作であること示す。

# python -m pip install transformers sentencepiece accelerate
# python rinna.py

DeferredCudaCallError

上記pythonコードを実行中に次のようなエラーとなった。

raise DeferredCudaCallError(msg) from e
torch.cuda.DeferredCudaCallError: CUDA call failed lazily at initialization with error: device >= 0 && device < num_gpus INTERNAL ASSERT FAILED at "/opt/pytorch/pytorch/aten/src/ATen/cuda/CUDAContext.\
cpp":50, please report a bug to PyTorch. device=1, num_gpus=

調べてみるとMIG設定時のエラーとのコメントもあった。自分のワークステーションは、A4000とQuadro K600の2つのGPUを装着している。このことから、使用するGPUを指定する必要があると考えて、次のとおり実行した。

# python -m pip install transformers sentencepiece accelerate
# export CUDA_VISIBLE_DEVICES="0"
# python rinna.py

実行結果

以下で「質問をどうぞ：」に続く文字列を入力した。

You are using the legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This means that tokens that come after special tokens will not be properly handled. We recommend you to read the related pull request available at https://github.com/huggingface/transformers/pull/24565
質問をどうぞ： 日本の終戦記念日はいつですか
1.9004747867584229
1945年8月15日です。</s>
質問をどうぞ： 江戸幕府を開いたのは誰ですか
1.9013473987579346
徳川家康は1603年に亡くなり、1604年に将軍になりました。</s>
質問をどうぞ： 大化の改新について教えてください
5.68342661857605
大化の改新は、天智天皇が日本の政治改革を行い、中央集権政府を確立した革命的な政治改革でした。この改革\
は中央政府を強化し、律令制度を確立しました。</s>
質問をどうぞ： ありがとうございました
0.6675028800964355
どういたしまして</s>
質問をどうぞ：

江戸幕府の答えが少し変だが、細かいところは目をつぶる。

TITAN Vでの実行

TITAN V（+Quadro K2000）を実装している別のワークステーションで同じことを実行した。

上記と同じように、NGC CATALOGからpytorchコンテナをダウンロードし、${HOME}/workspaceに格納しているrinna.pyを実行した。操作手順は次のとおり。

$ docker run -d -it --gpus all --name gpt -v ${HOEM}/workspace:/workspace/local \
nvcr.io/nvidia/pytorch:23.07-py3
$ docker exec -it gpt /bin/bash
## 以下はコンテナ内での操作
# export CUDA_VISIBLE_DEVICES="0"
# python -m pip install transformers sentencepiece accelerate
# cd local
# python rinna.py

NVIDIA driver is too old

上記を実行すると以下のエラーが発生。

RuntimeError: The NVIDIA driver on your system is too old (found version 11040). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver.

A4000では動いていた同じコンテナが、TITAN Vでは上記の通り「NVIDIA driver on your system is too old」との結果となり、理由が不明である。

NVIDIA driverが原因ではなく、pytorchがTITAN Vをサポートしていないことが原因だと想定し、古いバージョンのpytorchコンテナで試すことにした。

因みに 23.07-py3のpytorchコンテナのtorch関連のバージョンは次のとおりだった。

# pip list | grep torch
pytorch-quantization      2.1.2
torch                     2.1.0a0+b5021ba
torch-tensorrt            1.5.0.dev0
torchdata                 0.7.0a0
torchtext                 0.16.0a0
torchvision               0.16.0a0

1年前のpytorchコンテナで試す

古いバージョンのpytorchコンテナで試してみようと思い、理由はないが、1年前の22.07-py3を次のようにして使ってみた。

$ docker run -d -it --gpus all --name gpt -v /home/kenji/workspace:/workspace/local \
nvcr.io/nvidia/pytorch:22.07-py3
$ docker exec -it gpt /bin/bash

ここからは、コンテナ内の操作。

torchのバージョンは次のとおり1.13.0。

# pip list | grep torch
pytorch-quantization          2.1.2
torch                         1.13.0a0+08820cb
torch-tensorrt                1.2.0a0
torchtext                     0.13.0a0
torchvision                   0.14.0a0

# export CUDA_VISIBLE_DEVICES="0"
# python -m pip install transformers sentencepiece accelerate
# python rinna.py

CUDA out of memory

上記を実行すると、次のとおりGPUのVRAM容量不足となった。

RuntimeError: CUDA out of memory. Tried to allocate 122.00 MiB (GPU 0; 11.78 GiB total capacity; 11.03 GiB already allocated; 4.00 MiB free; 11.12 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

TITAN Vの12GBのVRAMでは足りないということなので、torch_dtypeを指定してモデル容量を小さくする。

具体的には、rinna.pyのmodelの5行目をコメントアウトし、6行目を有効にした。

# 標準
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft")
# 自動
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", device_map='auto')
# 自動(VRAM16GB以下でも8GBはNG)
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", torch_dtype=torch.float16, device_map='auto')
# CPU指定
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft").to("cpu")
# GPU指定
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft").to("cuda")
# GPU指定(VRAM16GB以下でも8GBはNG)
model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", torch_dtype=torch.float16).to("cuda")

実行結果

# python rinna.py
You are using the legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This means that tokens that come after special tokens will not be properly handled. We recommend you to read the related pull request available at https://github.com/huggingface/transformers/pull/24565
質問をどうぞ： 日本の終戦記念日はいつですか
1.9047455787658691
1945年8月15日です。</s>
質問をどうぞ： 江戸幕府を開いたのは誰ですか
0.19092082977294922
徳川家康</s>
質問をどうぞ： 大化の改新について教えてください
4.939814329147339
大化の改新とは、中大兄皇子、後の天智天皇が政治改革を行った一連の改革です。これらの改革は、公正な競争\
のルールや、身分制度にとらわれない人々の平等な権利の確立など、その後の日本の政治システムの基盤を提供\
しました。</s>
質問をどうぞ： 南北朝時代になった原因を教えてください
1.415543794631958
歴史の授業で、南北朝時代がどのように起こったかを説明しているのですか?</s>
質問をどうぞ： 南北朝時代がどうして起こったのか教えてください
9.151806592941284
それは複雑な過程を経ています。一般的には、中国が北朝と南朝に分裂したと言えますが、それだけではありま\
せん。北朝は、13世紀後半に始まった第1次モンゴル軍の侵攻や、14世紀後半の明朝の拡大に対処しなければな\\
らなかったため、大きく揺れ動いています。一方、南朝は、14世紀半ばの明朝による中国再統一の結果、安定し\
た状態にありました。しかしながら、清盛が平氏政権を樹立した後、北朝は弱体化し、最終的には明朝による統\
一によって終了することになりました。</s>
質問をどうぞ：

まとめ

今回、dockerコンテナで実行するので、インストールに関連したトラブルは無く、簡単に実行できると思っていたが、上記のとおり、少しだけハマった。ハマったおかげで、docker操作や関連する知識が少しだけ吸収できた。

ハマったところは次のとおり。なお、原因についてはakenjiの想定であり、実際は違うかもしれない。

DeferredCudaCallError GPUを複数実装している場合に発生し得る。CUDA_VISIBLE_DEVICESで使用するGPUを明記することで解決した。
The NVIDIA driver on your system is too old 使用するpytorchコンテナが新しくて、GPUのアーキテクチャーをサポートしていない。古いpytorchコンテナを使う。 pytorchバージョンとサポートするsm（NVIDIA architecuter name）の関連表を知っている方、教えてほしい。
CUDA out of memory 良くあるGPUのVRAM不足。モデルサイズを小さくする。今回は、torch_dtype=torch.float16と指定した。