WSL2でLLM-jp-13B v1.1を試してみる

2024年2月12日 20:40

「インストラクションチューニングの設定の見直しや DPO (Direct Preference Optimization) の追加によって，性能を改善したチューニング済みモデル v1.1 とその際に使用したデータとコードを公開」されたらしいので試してみます。

使用するPCはドスパラさんの「GALLERIA UL9C-R49」。スペックは
・CPU: Intel® Core™ i9-13900HX Processor
・Mem: 64 GB
・GPU: NVIDIA® GeForce RTX™ 4090 Laptop GPU(16GB)
・GPU: NVIDIA® GeForce RTX™ 4090 (24GB)
・OS: Ubuntu22.04 on WSL2（Windows 11）
です。

1. 準備

python3 -m venv llm-jp
cd $_
source bin/activate

パッケージのインストール。

pip install torch transformers accelerate

2. コード

流し込むコード

import sys
import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
from typing import List, Dict
import time

# argv
parser = argparse.ArgumentParser()
parser.add_argument("--model-path", type=str, default=None)
parser.add_argument("--no-instruct", action='store_true')
parser.add_argument("--no-use-system-prompt", action='store_true')
parser.add_argument("--max-tokens", type=int, default=256)

args = parser.parse_args(sys.argv[1:])

model_id = args.model_path
if model_id == None:
    exit

is_instruct = not args.no_instruct
use_system_prompt = not args.no_use_system_prompt
max_new_tokens = args.max_tokens

# トークナイザーとモデルの準備
tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    trust_remote_code=True
)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    device_map="auto",
    #device_map="cuda",
    low_cpu_mem_usage=True,
    trust_remote_code=True
)
#if torch.cuda.is_available():
#    model = model.to("cuda")

streamer = TextStreamer(
    tokenizer,
    skip_prompt=True,
    skip_special_tokens=True
)

DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"

# generation params
generation_params = {
    "do_sample": True,
    "temperature": 0.8,
    "top_p": 0.95,
    "top_k": 40,
    "max_new_tokens": max_new_tokens,
    "repetition_penalty": 1.1,
}


def q(
    user_query: str,
    history: List[Dict[str, str]]=None
):
    start = time.process_time()
    # messages
    messages = ""
    if is_instruct:
        messages = []
        if use_system_prompt:
            messages = [
                {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
            ]
        user_messages = [
            {"role": "user", "content": user_query}
        ]
    else:
        user_messages = user_query
    if history:
        user_messages = history + user_messages
    messages += user_messages
    # generation prompts
    if is_instruct:
        prompt = tokenizer.apply_chat_template(
            conversation=messages,
            add_generation_prompt=True,
            tokenize=False
        )
    else:
        prompt = messages
    input_ids = tokenizer.encode(
        prompt,
        add_special_tokens=False,
        return_tensors="pt"
    )
    print("--- prompt")
    print(prompt)
    print("--- output")
    # 推論
    output_ids = model.generate(
        input_ids.to(model.device),
        streamer=streamer,
        **generation_params
    )
    output = tokenizer.decode(
        output_ids[0][input_ids.size(1) :],
        skip_special_tokens=True
    )
    if is_instruct:
        user_messages.append(
            {"role": "assistant", "content": output}
        )
    else:
        user_messages += output
    end = time.process_time()
    ##
    input_tokens = len(input_ids[0])
    output_tokens = len(output_ids[0][input_ids.size(1) :])
    total_time = end - start
    tps = output_tokens / total_time
    print(f"prompt tokens = {input_tokens:.7g}")
    print(f"output tokens = {output_tokens:.7g} ({tps:f} [tps])")
    print(f"   total time = {total_time:f} [s]")
    return user_messages

チャットテンプレートの準備

LLM-jpはチャットテンプレートの定義がない（このモデルの提供方法ですと仕方がないのですが）ので、以下のように定義したテンプレートを準備します。

{{'以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。\n\n'}}
{% for message in messages %}
	{% if message['role'] == 'user' %}
		{{'### 指示:\n' + message['content'] + '\n\n'}}
	{% endif %}
	{% if message['role'] == 'assistant' %}
		{{'### 応答:\n' + message['content'] + '\n\n'}}
	{% endif %}
{% endfor %}
{% if add_generation_prompt %}
	{{ '### 応答:\n' }}
{% endif %}

（注）見やすいように改行とインデントしています。

3. 試してみる

RTX 4090 (24GB)1枚だけですとVRAMが溢れてしまい、秒あたり0.9トークンとなりますので、2枚使います。

どのモデル？

3つモデルが公開されています。今回は、2つめのモデルを使います。

実行

python -i query.py \
	--model-path llm-jp/llm-jp-13b-instruct-full-dolly_en-dolly_ja-ichikara_003_001-oasst_en-oasst_ja-v1.1

モデルのロード直後のVRAM使用量は、
・RTX 4090 (24GB) : 12.4GB
・RTX 4090 Laptop GPU (16GB) : 12.7GB
なので、25.1GBほどです。

RTX 4090 Laptop GPU(16GB): 13.3 - 0.6 = 12.7GB

チャットテンプレートの設定

pythonのプロンプト（>>>）が表示されたら、先ほど紹介したチャットテンプレートの定義を（改行とインデントを削除して）を流し込み、AutoTokenizerのメンバ変数chat_templateに設定します。

tokenizer.chat_template = "{{'以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。\n\n'}}{% for message in messages %}{% if message['role'] == 'user' %}{{'### 指示:\n' + message['content'] + '\n\n'}}{% endif %}{% if message['role'] == 'assistant' %}{{'### 応答:\n' + message['content'] + '\n\n'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '### 応答:\n' }}{% endif %}"

準備が整いました。それでは聞きましょう。

聞いてみる

>>> history = q("ドラえもんとはなにか")
--- prompt
以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。

### 指示:
ドラえもんとはなにか

### 応答:

ドラえもんは日本のテレビアニメシリーズである。1969年から小学館の雑誌で連載された児童向けSF漫画作品、藤子・F・不二雄によって書かれた『ドラえもん』が原作。主人公の小学生、野比のび太のもとへ、勉強もスポーツも苦手な、けれど勉強もスポーツも苦手だった彼のもとへタイムマシンに乗ってやってきたネコ型ロボット「ドラえもん」。ドラえもんには数々のひみつ道具がついています。ドラえもんの持ち主はのび太ですが、のび太以外にも持ち主がいて、持ち主に代わってのび太の世話役をしてくれたりします。のび太の世話役として登場したキャラクターたちは、ドラえもんと共にのび太を助けたり、のび太を叱咤激励してのび太を成長させたりしました。のび太やその仲間たちとともに、のび太とドラえもんは数々の冒険を繰り広げました。22世紀からきたのび太とドラえもんが繰り広げる物語は、当時の子供たちだけではなく、今大人になった人たちの心にも響き続けているようです。
prompt tokens = 43
output tokens = 223 (16.954024 [tps])
total time = 13.153219 [s]

llm-jp-13b-instruct-full-dolly_en-dolly_ja-ichikara_003_001-oasst_en-oasst_ja-v1.1

素晴らしい。
続きを聞いてみます。

>>> history = q("続きを教えてください。", history)
--- prompt
以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。

### 指示:
ドラえもんとはなにか

### 応答:
ドラえもんは日本のテレビアニメシリーズである。1969年から小学館の雑誌で連載された児童向けSF漫画作品、藤子・F・不二雄によって書かれた『ドラえもん』が原作。主人公の小学生、野比のび太のもとへ、勉強もスポーツも苦手な、けれど勉強もスポーツも苦手だった 彼のもとへタイムマシンに乗ってやってきたネコ型ロボット「ドラえもん」。ドラえもんには数々のひみつ道具がついています。ドラえ もんの持ち主はのび太ですが、のび太以外にも持ち主がいて、持ち主に代わってのび太の世話役をしてくれたりします。のび太の世話役 として登場したキャラクターたちは、ドラえもんと共にのび太を助けたり、のび太を叱咤激励してのび太を成長させたりしました。のび 太やその仲間たちとともに、のび太とドラえもんは数々の冒険を繰り広げました。22世紀からきたのび太とドラえもんが繰り広げる物語 は、当時の子供たちだけではなく、今大人になった人たちの心にも響き続けているようです。

### 指示:
続きを教えてください。

### 応答:

ドラえもんは主人公ののび太の元へやってきたネコ型ロボットです。のび太の世話役として、のび太を助けたり、のび太を励ましたりしました。ドラえもんはのび太の世話役になる前は、のび太より少し高いところに住んでいました。のび太と同じ学校に通っているわけではありませんが、のび太を助けるためにのび太の家へよく訪れています。のび太はドラえもんに助けられながら、勉強もスポーツも苦手なのび太が努力し、のび太自身の力で成長していきます。のび太はドラえもんと共にのび太にとって大きな存在になりました。
prompt tokens = 289
output tokens = 129 (16.563327 [tps])
total time = 7.788290 [s]