MPT-30Bの4bit量子版をローカルRTX4090で動作させる

WSL2 Ubuntu 20.04
GPU:RTX 4090 24GB
CUDA: 11.8

mosaicMLが6/22に発表した約300億パラメータのLLMです。ベースモデル、Instructionモデル、Chatモデルの３種類が提供されています。

• mpt-30b: ベースモデル。入力後に続くと推測される文章を生成する。
• mpt-30b-instruct: Instructionモデル。入力の指示に従ってタスクを実行する文章を生成する。
• mpt-30b-chat: Chatモデル。ChatGPTのように会話形式で返答を生成する。

自分用に作っているものですが、再現できるようにpoetryで書いているので環境構築が簡単に出来ると思います。
簡単なテスト等に使っていただけると嬉しいです。(機能追加等のPRも募集中です！)
japanese-llm-experiments

MPT-30Bのbaseモデルを動作させます。

1import torch
2from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
3
4tokenizer = AutoTokenizer.from_pretrained("mosaicml/mpt-30b")
5model = AutoModelForCausalLM.from_pretrained(
6            "mosaicml/mpt-30b",
7            trust_remote_code=True,
8            load_in_4bit=True,
9            device_map="auto",
10        )
11
12prompt = "Q:Vtuberってなんですか？\nA:"
13
14token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
15with torch.no_grad():
16    output_ids = model.generate(
17        token_ids.to(model.device),
18        max_new_tokens=256,
19        do_sample=True,
20        temperature=0.7,
21    )
22
23    output = self.tokenizer.decode(output_ids.tolist()[0])
24
25print(output)

まずは英語入力の結果から。入力は「Q: Whats Vtuber?\nA: 」というQA方式にしました。

1Q:What's Vtuber? 
2A: A Vtuber is a virtual YouTuber. They use a virtual character to create videos (usually on YouTube) and interact with their fans. These characters are often designed to look like the YouTuber themselves, but can also be completely different.
3
4Q:What is Nijisanji?
5A: Nijisanji (ニジサンジ) is a Vtuber management company based in Japan that was founded in April 2016. It is the largest Vtuber management company in the world and has over 300 Vtubers under its management.
6
7Q:What is the difference between VTuber and Vtuber?
8A: There is no difference between VTuber and Vtuber. They are both terms used to refer to virtual YouTubers.

次に日本語入力の結果。入力は「Q: Vtuberって何?\nA: 」です。

1Q:Vtuberってなんですか? 
2A:Vtuberとは、YouTuberとして活動する主に女性が多い視聴者を集めるYouTuberのことです。
3このカテゴリには、YouTuberとして主に女性が多い視聴者を集めているYouTuberを含むことに注意してください。
4
5Q:Vtuberがあるのはなぜですか?
6A:YouTuberをしたいという人が多いのですが、実際に目に見えたところのYouTuberは多くないため、Vtuberを通してYouTuberを始めることができます。
7
8Q:Vtuberは、どうして初めてYouTubeに参加しましたか?
9A:Vtuberは、多くの場合、YouTuberを始めたいと思っていたとき以前からVtuberを見ていたため、最初からYouTubeに参加したいと思っていました。

日本語もいけるみたいですが、あまり出力結果は良くない感じでした。
まだ簡単なテストしかしてないので、今後の活用方法に期待ですね。

環境構築した時に下記のようなエラーが出る可能性があります。このときはbitsandbytesの_cpu.soファイルを_cuda118.soで上書きすればOKです。(別のCUDAバージョンの場合はそれに対応するsoファイルをコピーしてください)

AttributeError: japanese-llm-experiments/.venv/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cpu.so: undefined symbol: cget_col_row_stats

コマンド例

1cd japanese-llm-experiments/.venv/lib/python3.10/site-packages/bitsandbytes
2cp libbitsandbytes_cuda118.so libbitsandbytes_cpu.so