🐡

Tanuki-8x8Bをダウンロードした後に推論させる方法について

2024/08/30に公開

GENIAC 松尾研LLM開発プロジェクトメンバーのArataです。
本記事では、独自アーキテクチャを採用したTanuki-8x8Bモデルの推論方法についてまとめます。

はじめに

Tanuki-8x8Bは独自アーキテクチャである関係上、推論を行う際に一部対応が必要となります。
この記事では、Tanuki-8x8Bおよびその量子化モデルの推論方法についてまとめます。

CUDA 12.1およびPython 3.10、Transformers 4.44.2の環境で動作を確認しています。

モデルの種類と推論方法

Tanuki-8x8Bには以下の5種類のモデルがあります。

これらのモデルについて、推論方法を以下の5通り紹介します。

  1. vLLMによる推論(最推奨)
  2. Huggingface Transformersによる推論(推奨)
  3. AutoGPTQによる推論
  4. AutoAWQによる推論
  5. llama.cppによる推論(非推奨)

1. vLLMによる推論(最推奨)

対応するモデル

Tanuki-8x8Bが推論できるようチームで独自改変したvLLMを利用して推論する方法です。オリジナルだけでなく、AWQやGPTQ版も対応しています。高速で推論が可能であるため、最も推奨いたします。

推論方法

以下のようにソースからvLLMをビルドします。また、flash attentionもインストールします。

git clone https://github.com/team-hatakeyama-phase2/vllm.git
cd vllm
LD_LIBRARY_PATH="" MAX_JOBS=16 pip install -e .
pip install --no-build-isolation flash_attn

その後、モデルを読み込んで推論します。サンプルコードを以下に示します。

# test_vllm.py

from time import time
from vllm import LLM, SamplingParams

model_name = "weblab-GENIAC/Tanuki-8x8B-dpo-v1.0"
# model_name = "team-hatakeyama-phase2/Tanuki-8x8B-dpo-v1.0-AWQ"
# model_name = "team-hatakeyama-phase2/Tanuki-8x8B-dpo-v1.0-GPTQ-4bit"
# model_name = "team-hatakeyama-phase2/Tanuki-8x8B-dpo-v1.0-GPTQ-8bit"

# vllm = LLM(model_name, trust_remote_code=True, tensor_parallel_size=1)  # 1GPUの場合
vllm = LLM(model_name, trust_remote_code=True, tensor_parallel_size=2)  # 2GPUを使う場合
tokenizer = vllm.get_tokenizer()

messages = [
    {"role": "system", "content": "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。"},
    {"role": "user", "content": "AIによって私たちの暮らしはどのように変わりますか?"}
]

inputs_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
print(f"inputs_text: {inputs_text}")

sampling_params = SamplingParams(temperature=0.0, max_tokens=1024, seed=1, repetition_penalty=1.1)
start = time()
outputs = vllm.generate(inputs_text, sampling_params=sampling_params, use_tqdm=False)
end = time()
outputs_text = outputs[0].outputs[0].text
print(f"outputs_text: {outputs_text}")
print(f"Elapsed time: {(end - start):.4f} sec.")

2. Huggingface Transformersによる推論(推奨)

対応するモデル

Huggingface Transformersを使って推論する方法です。独自改変したライブラリのビルド等も必要なく最もシンプルに実装出来ますが、推論速度はvLLMと比較すると遅いです。

なお、本来はGPTQ版もHuggingface Transformersで推論できるはずですが、このissueと同じ問題が発生したため動作未確認です。

推論方法

必要なライブラリをインストールします。

pip install transformers accelerate bitsandbytes
pip install --no-build-isolation flash_attn
pip install autoawq # AWQ版も推論したい場合

その後、モデルを読み込んで推論します。サンプルコードを以下に示します。

# test_transformers.py

from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

model_name = "weblab-GENIAC/Tanuki-8x8B-dpo-v1.0"
# model_name = "team-hatakeyama-phase2/Tanuki-8x8B-dpo-v1.0-AWQ"

model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_name)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

messages = [
    {"role": "system", "content": "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。"},
    {"role": "user", "content": "AIによって私たちの暮らしはどのように変わりますか?"}
]

input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
output_ids = model.generate(input_ids, max_new_tokens=1024, temperature=0.5, streamer=streamer)

3. AutoGPTQによる推論

対応するモデル

*本来は推論可能なはずですが、原因不明の問題で上手く行っていません。詳細は後述します。

Tanuki-8x8Bに対応するよう独自改変したAutoGPTQを使って推論する方法です。Huggingface Transformersを使った推論と同じようなコードで推論できます。

推論方法

以下のように必要なライブラリをインストールし、ソースからAutoGPTQをビルドします。

pip install accelerate transformers
pip install numpy gekko pandas
git clone https://github.com/team-hatakeyama-phase2/AutoGPTQ.git
cd AutoGPTQ
pip install --no-build-isolation -e .
pip install --no-build-isolation flash_attn

その後、モデルを読み込んで推論します。サンプルコードを以下に示します。

# test_autogptq.py

from transformers import AutoTokenizer, TextStreamer
from auto_gptq import AutoGPTQForCausalLM

model_name = "team-hatakeyama-phase2/Tanuki-8x8B-dpo-v1.0-GPTQ-4bit"
# model_name = "team-hatakeyama-phase2/Tanuki-8x8B-dpo-v1.0-GPTQ-8bit" # 8bitはエラーが出る

model = AutoGPTQForCausalLM.from_quantized(model_name, device_map="auto", trust_remote_code=True, use_marlin=True)
tokenizer = AutoTokenizer.from_pretrained(model_name)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

messages = [
    {"role": "system", "content": "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。"},
    {"role": "user", "content": "AIによって私たちの暮らしはどのように変わりますか?"}
]

input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
inputs = tokenizer(input_text, return_tensors="pt", return_token_type_ids=False).to(model.device)

output_ids = model.generate(**inputs, max_new_tokens=1024, temperature=0.5, streamer=streamer)

なお、GPTQ 8bitモデルを上記コードで読み込もうとするとこのissueと同じ問題が発生するため、動作未確認です。

4. AutoAWQによる推論

対応するモデル

Tanuki-8x8Bに対応するよう独自改変したAutoAWQを使って推論する方法です。AWQはHuggingface Transformersでそのまま読み込めるので、あまりこれを使う意味はないと思います。

推論方法

以下のようにソースからAutoAWQをビルドします。また、flash attentionもインストールします。

git clone https://github.com/team-hatakeyama-phase2/AutoAWQ.git
cd AutoAWQ
pip install -e .
pip install --no-build-isolation flash_attn

その後、モデルを読み込んで推論します。サンプルコードを以下に示します。

# test_autoawq.py

from transformers import AutoTokenizer, TextStreamer
from awq import AutoAWQForCausalLM

model_name = "team-hatakeyama-phase2/Tanuki-8x8B-dpo-v1.0-AWQ"

model = AutoAWQForCausalLM.from_quantized(model_name, device_map="auto", trust_remote_code=True, fuse_layers=False) # fuse_layers=Trueだとエラーが出る
tokenizer = AutoTokenizer.from_pretrained(model_name)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

messages = [
    {"role": "system", "content": "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。"},
    {"role": "user", "content": "AIによって私たちの暮らしはどのように変わりますか?"}
]

input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
inputs = tokenizer(input_text, return_tensors="pt", return_token_type_ids=False).to("cuda")

output_ids = model.generate(**inputs, max_new_tokens=1024, temperature=0.5, streamer=streamer)

5. llama.cppによる推論(非推奨)

対応するモデル

llama.cppを使ってGGUFモデルを元に推論する方法です。ただし、この方法には以下の問題があります。

  • tokenizerが正しくGGUF変換出来ているのか未確認
  • 独自アーキテクチャであるTanukiForCausalLM として推論されず、MixtralForCausalLM として推論されてしまう

上記の問題から、この方法による推論では性能低下が発生します。具体的には、内部のテストでJMT-Benchにおいて全体スコア-0.5点程度、特に数学やコーディング、リーズニングの難易度が高いタスクで大きな性能低下が起こることを確認しています。また、人手評価では更に大きな性能低下が確認される可能性もあります。
そのためこの推論方法は非推奨としています。利用する際は性能が低下していることをご理解の上お使いください。

推論方法

以下のようにllama-cpp-pythonをインストールします。(CUDAのバージョンによってコマンドは一部調整)

pip install llama-cpp-python --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu121

その後、モデルを読み込んで推論します。サンプルコードを以下に示します。

# test_llamacpp.py

from llama_cpp import Llama

model_name = "team-hatakeyama-phase2/Tanuki-8x8B-dpo-v1.0-GGUF"
filename = "Tanuki-8x8B-dpo-v1.0-IQ4_NL.gguf"

llm = Llama.from_pretrained(repo_id=model_name, filename=filename, n_gpu_layers=99)

prompt = """<s>以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。

### 指示:
AIによって私たちの暮らしはどのように変わりますか?

### 応答:
"""

output = llm.create_completion(prompt, max_tokens=1024, temperature=0.5, echo=True)

print(output['choices'][0]['text'])

まとめ

この記事では、Tanuki-8x8Bの推論について以下の5通りの方法を紹介しました。

  1. vLLMによる推論(最推奨)
  2. Huggingface Transformersによる推論(推奨)
  3. AutoGPTQによる推論
  4. AutoAWQによる推論
  5. llama.cppによる推論(非推奨)

特にvLLMによる推論は正確かつ高速であるため、基本的にこれを最も推奨いたします。モデル利用の際の参考になれば幸いです。

東大松尾・岩澤研究室 | LLM開発 プロジェクト[GENIAC]

Discussion