ローカルでLLM-jp-13Bを動かす

NII主宰のLLM勉強会（LLM-jp）が公開したモデルをローカルのGPUで動かします。利用したモデルは、instructの全部のせ？のモデルです。

llm-jp/llm-jp-13b-instruct-full-jaster-dolly-oasst-v1.0 · Hugging Face

環境

GPU：NVIDIA GeForce RTX3060 12GB
RAM：32GB
torch：2.0.1
transformers：4.30.2
Docker上にJupyterLabを構築しています。

モデル読み込み

デモコードのままロードするとクラッシュするので、8bit量子化でロードします。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

tokenizer = AutoTokenizer.from_pretrained("llm-jp/llm-jp-13b-instruct-full-jaster-dolly-oasst-v1.0")
model = AutoModelForCausalLM.from_pretrained("llm-jp/llm-jp-13b-instruct-full-jaster-dolly-oasst-v1.0", load_in_8bit=True)

ロード後のサイズは約26GBで、共有GPUメモリ(RAM)まで使っていますが、クラッシュはしないのでギリギリ許容範囲だと思って進めます。

GPT2アーキテクチャのモデルは、Linear層ではなくConv1D層のため、下記のエラーがでます。

You are loading your model in 8bit or 4bit but no linear modules were found in your model. this can happen for some architectures such as gpt2 that uses Conv1D instead of Linear layers. Please double check your model architecture, or submit an issue on github if you think this is a bug.

Conv1Dの層が量子化できておらずサイズがあまり減ってないだけで、推論に影響はないと思うので放置します。（詳しい方教えてください）

推論の実行

%%time
text = """好きな映画はあったりしますか？"""
text = text + "### 回答："
tokenized_input = tokenizer.encode(text, add_special_tokens=False, return_tensors="pt").to(model.device)
with torch.no_grad():
    output = model.generate(
        tokenized_input,
        max_new_tokens=512,
        do_sample=True,
        top_p=0.95,
        temperature=0.7,
    )[0]

input_length = len(tokenized_input[0])
print(tokenizer.decode(output[input_length:]))

そうだね、『ショーシャンクの空に』は好きだ。<EOD|LLM-jp>
CPU times: user 52.8 s, sys: 71.6 ms, total: 52.9 s
Wall time: 52.8 s

第1印象としては、他の日本語LLMより圧倒的に回答がそれっぽくて自然です。ただ、回答がかなりあっさりとした短文や単語で帰ってくることも多いです。

プロンプトで長めに回答するように指示した場合

text = """長めの文で以下の問いかけに回答してください。
好きな映画はなんですか？"""

私は映画が大好きです。私が好きな映画は、例えば『ミッションインポッシブル』や『アベンジャーズ』、『ハリー・ポッター』、『スター・ウォーズ』などです。<EOD|LLM-jp>
CPU times: user 2min 6s, sys: 70.2 ms, total: 2min 6s
Wall time: 2min 6s

「アベンジャーズ」、「アバター」、「タイタニック」、「フォレスト・ガンプ」、「ショーシャンクの空に」、「インセプション」<EOD|LLM-jp>
CPU times: user 1min 57s, sys: 116 ms, total: 1min 57s
Wall time: 1min 57s

数十回程度やってみた雰囲気では、プロンプトによる制御も効いてそうです。

今後の日本語LLMは、このモデルをベースに追加学習などで調整するのが良さそうです。