Ubuntu Server 22.04 + GPU 12GB で Rinna 3.6B を動かしたので、その備忘録です。
モチベーションとしては、LlamaIndex で実現しているのと同じようなことをこのモデルなどを活用しながらつくることができないかと思っています。
CUDAは現時点での最新の12.1 を使用しました。
$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Mon_Apr__3_17:16:06_PDT_2023
Cuda compilation tools, release 12.1, V12.1.105
Build cuda_12.1.r12.1/compiler.32688072_0
動かした方の基本は https://huggingface.co/rinna/japanese-gpt-neox-3.6b-instruction-sft のページをご覧ください。
いわゆる Rinna 3.6B と言われているモデルは japanese-gpt-neox-3.6b-instruction-sft と japanese-gpt-neox-3.6b の２つがありますが、 instruction-sft がついている方は、 付いていない方(japanese-gpt-neox-3.6b) をベースにそれを fine-tuning して会話できるようにしたものです。
The model is based on rinna/japanese-gpt-neox-3.6b and has been finetuned to serve as a instruction-following conversational agent.
12GB の GPU では、そのままでは動かないので、モデルのロード部分を以下のようにしました。
#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft")
model = AutoModelForCausalLM.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-sft",
torch_dtype=torch.float16)
モデルロード時のオプションに torch_dtype=torch.float16 を追加。
こちらのページを参考にしました。
【Python】手持ちのGPUがVRAM12Gだけど「Rinna-3.6B」とお話がしたい！！！
コード全体では次のようになります。
main.py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
model_name = "rinna/japanese-gpt-neox-3.6b-instruction-sft"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False, torch_dtype=torch.float16)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
if torch.cuda.is_available():
model = model.to("cuda")
prompt = "ユーザー: 日本のおすすめの観光地を教えてください。<NL>システム: どの地域の観光地が知りたいですか？<NL>ユーザー: 渋谷の観光地を教えてください。<NL>システム:"
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")
print(output)
実行してみると、以下のような結果になりました。
$ time python3 main.py
わかりました。まずは、ハチ公像から始めましょう。ハチ公像は、日本の最も有名な観光スポットの1つです。ハチ公像は、日本の建築家である安藤忠雄によって設計されま
した。ハチ公像は、日本の彫刻家である安藤忠雄によって設計されました。ハチ公像は、日本の彫刻家である安藤忠雄によって設計されました。ハチ公像は、日本の彫刻家である安藤忠雄によって設計されました。ハチ公像は、日本の彫刻家である安藤忠雄によって設計されました。ハチ公像は、日本の彫刻家である安藤忠雄によって設計されました。ハチ公像は
real 1m28.380s
user 1m16.951s
sys 0m9.055s
これで動かすことはできたのですが、一回実行するごとに 1分30秒ほどの時間がかかっています。 ロード時間が遅いだけなので、開発中でなければ問題になることはないかもしれませんが、開発中はちょっと遅すぎてつらいです。
さらに調べてみると int8 まで量子化する例がありましたので、試しました。
話題のrinna-3.6bをColab無料枠で動かしたい！
追加でいくつかのモジュールを入れることで int8 に量子化して少ないメモリで動かすことができるらしいです。 詳しくは https://huggingface.co/docs/accelerate/index を読みましょう。
モデルをロードする部分を以下のように変更します。
model = AutoModelForCausalLM.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-sft",
device_map = "auto",
load_in_8bit=True)
device_map と load_in_8bit オプションを追加します。
device_map の部分は auto ではなく明示的に指定した方がよいかもしれません。 以下のように:
device_map = infer_auto_device_map(
model,
max_memory={0: "10GiB", "cpu": "10GiB"},
no_split_module_classes=["GPTNeoXLayer"],
dtype=torch.int8)
model = AutoModelForCausalLM.from_pretrained(
"rinna/japanese-gpt-neox-3.6b-instruction-sft",
device_map = device_map,
load_in_8bit=True)
コードは全体 main.py
モデルをロードする部分以外のコードに変更はありません。
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
model_name = "rinna/japanese-gpt-neox-3.6b-instruction-sft"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map = "auto",
load_in_8bit=True)
prompt = "ユーザー: 日本のおすすめの観光地を教えてください。<NL>システム: どの地域の観光地が知りたいですか？<NL>ユーザー: 渋谷の観光地を教えてください。<NL>システム:"
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")
print(output)
実行してみます。
$ time python3 main.py
わかりました。まずは、ハチ公像から始めましょう。ハチ公像は、日本の東京にある有名な観光スポットです。ハチ公像は、日本の忠犬ハチ公にちなんで名付けられた像です。ハチ公像は、日本の東京にある有名な観光スポットです。</s>
real 0m28.014s
user 0m19.643s
sys 0m8.601s
約30秒で実行できました。先ほどの float16 と比べて 1/3 の時間で実行できました。これなら耐えられるかも。
さて、一番関心のある Rinna の回答例をいくつかみてみましょう。
例1: 逆に質問された例
わかりました。どのような観光スポットをお探しですか？
質問したのに、質問でかえされるとちょっとイラッとしますね。 しかし、そのあとも続けて会話できるのだから、これはこれで良いのかもしれない。
例2: うまくいかなった例
わかりました。渋谷の観光地は、ハチ公像、スクランブル交差点、109、スクランブルエッグ、ハチ公像、忠犬ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公
像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ公像、ハチ
ハチ公像推しがすごい。
最後に int8 で動かすための環境構築方法をメモしておきます。 OSは Ubuntu Server 22.04 です。
Python:
$ python3 --version
Python 3.10.6
CUDA:
$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Mon_Apr__3_17:16:06_PDT_2023
Cuda compilation tools, release 12.1, V12.1.105
Build cuda_12.1.r12.1/compiler.32688072_0
j-gpt-neox-3.6b 環境を venv でつくります。
$ python3 -m venv ~/.local/j-gpt-neox-3.6b
$ source ~/.local/j-gpt-neox-3.6b/bin/activate
必要なモジュールを入れます。
(j-gpt-neox-3.6b) $ pip install transformers sentencepiece accelerate bitsandbytes scipy
PyTorch も必要ですが、明示的に指定しなくても上記コマンドで入りました。
これでAutoModelForCausualLMの方のモデルのロードはできるのですが、Tokenizer はさらに protobuf モジュールが必要です。
しかも pip install protobuf してからコード(main.py)を動かすとバージョンを下げろ（ダウングレードしろ）とのエラーがでます。 結論としては 3.19.0 を指定してインストールすることで作動しました。
(j-gpt-neox-3.6b) $ pip install protobuf==3.19.0
以上で環境構築完了です。 あとは main.py のコードを実行するだけです。
最後にまとめとして、この段階での pip 環境を確認しておきます。
(j-gpt-neox-3.6b) $ pip list
Package Version
------------------------ ----------
accelerate 0.19.0
bitsandbytes 0.39.0
certifi 2023.5.7
charset-normalizer 3.1.0
cmake 3.26.3
filelock 3.12.0
fsspec 2023.5.0
huggingface-hub 0.14.1
idna 3.4
Jinja2 3.1.2
lit 16.0.5
MarkupSafe 2.1.2
mpmath 1.3.0
networkx 3.1
numpy 1.24.3
nvidia-cublas-cu11 11.10.3.66
nvidia-cuda-cupti-cu11 11.7.101
nvidia-cuda-nvrtc-cu11 11.7.99
nvidia-cuda-runtime-cu11 11.7.99
nvidia-cudnn-cu11 8.5.0.96
nvidia-cufft-cu11 10.9.0.58
nvidia-curand-cu11 10.2.10.91
nvidia-cusolver-cu11 11.4.0.1
nvidia-cusparse-cu11 11.7.4.91
nvidia-nccl-cu11 2.14.3
nvidia-nvtx-cu11 11.7.91
packaging 23.1
pip 22.0.2
protobuf 3.19.0
psutil 5.9.5
PyYAML 6.0
regex 2023.5.5
requests 2.31.0
scipy 1.10.1
sentencepiece 0.1.99
setuptools 59.6.0
sympy 1.12
tokenizers 0.13.3
torch 2.0.1
tqdm 4.65.0
transformers 4.29.2
triton 2.0.0
typing_extensions 4.6.1
urllib3 2.0.2
wheel 0.40.0
ちなみに、CUDA 11.5 で試した場合、これ作動しなかったので注意してください。 CUDA 12.1 にすることで動くようになりました。
CUDA 11.5 で作動しなかった原因は（詳しくはわかりませんが、どうやら） bitsandbytes モジュールが CUDA 11.5 では作動しないからのように思えました。 CUDA 11.8 まで上げればOK的な情報もありました。
以上です。