WSL2でSakana AIを試してみる
「進化的モデルマージにより日本語数学LLMとして構築したEvoLLM-JPは、数学のみならず、日本語の全般的な能力に長けている」らしいEvoLLM-JPを試してみます。
追記(2024/3/22)
10Bのモデルですが、torch_dtypeを"auto"からtorch.bfloat16に変更すると、推論のスピードが改善しました。
モデルEvoLLM-JPは、3種提供されています。今回は3つとも試します。
- 大規模言語モデル
SakanaAI/EvoLLM-JP-v1-10B : MICROSOFT RESEARCH LICENSE TERMS
SakanaAI/EvoLLM-JP-v1-7B : MICROSOFT RESEARCH LICENSE TERMS
SakanaAI/EvoLLM-JP-A-v1-7B : Apache License, Version 2.0
※画像言語モデルは、別の機会に。
SakanaAI/EvoVLM-JP-v1-7B : Apache License, Version 2.0
使用する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. 準備
venv構築
python3 -m venv sakanaai
cd $_
source bin/activate
パッケージのインストール。
pip install torch transformers accelerate flash_attn
2. 流し込むコード
いつものコードです。query.pyという名前で保存します。
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-chat", 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_chat = not args.no_chat
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
) -> List[Dict[str, str]]:
start = time.process_time()
# messages
messages = ""
if is_chat:
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_chat:
prompt = tokenizer.apply_chat_template(
conversation=messages,
add_generation_prompt=True,
tokenize=False
)
else:
prompt = messages
input_ids = tokenizer.encode(
prompt,
add_special_tokens=True,
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_chat:
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
print('history = ""')
print('history = q("ドラえもんとはなにか")')
print('history = q("続きを教えてください", history)')
3. 試してみる
(1) SakanaAI/EvoLLM-JP-v1-10B
pythonコマンドで起動して、
python -i query.py --model-path SakanaAI/EvoLLM-JP-v1-10B
と実行して、聞いてみます。
>>> history = q("ドラえもんとはなにか")
26分かかりました。少し化けてもいます。
追記(2024/3/22)
XにてSakana AIの中の方から「bfloat16にしてみて」とのコメントをいただきました。
で、試したところ、確かに大幅改善しました。
(2) SakanaAI/EvoLLM-JP-v1-7B
>>> history = q("ドラえもんとはなにか")
(3) SakanaAI/EvoLLM-JP-A-v1-7B
>>> history = q("ドラえもんとはなにか")
追記(2024/3/22)
こちらも bfloat16 で試したところ、、、
4. まとめ
ソースとしているモデルが異なるから当然と言えば当然ですが、応答結果と速度がかなり違いますね。
・EvoLLM-JP-v1-10B : 0.087481 [トークン/秒]
・EvoLLM-JP-v1-10B : 8.515899 [トークン/秒]
・EvoLLM-JP-v1-7B : 26.959319 [トークン/秒]・EvoLLM-JP-A-v1-7B : 4.062542 [トークン/秒]
・EvoLLM-JP-A-v1-7B : 24.724166 [トークン/秒]
「うん?推論、遅いなぁ」と思ったら、torch_dtype、"auto"ではなくtorch.bfloat16みたく指定してみるのがよいですね。
この記事が気に入ったらサポートをしてみませんか?