ELYZA-japanese-Llama-2-7bをAWQ化して利用する

ローカル環境（RTX3090）にて量子化と推論が両方問題なく動いたので共有です。やってることは、autoawqのREADMEをそのままELYZA-japanese-Llama-2-7bに適用しただけです。

量子化

!pip install autoawq
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = 'elyza/ELYZA-japanese-Llama-2-7b-instruct'
quant_path = 'ELYZA-japanese-Llama-2-7b-instruct-awq'
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4 }

# Load model
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize
model.quantize(tokenizer, quant_config=quant_config)

# Save quantized model
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

これだけでOKです。15分程度で完了しました。

量子化中のリソースモニターです。VRAMにも余裕がありました。

推論

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer, TextStreamer
import torch

quant_path = './ELYZA-japanese-Llama-2-7b-instruct-awq'
quant_file = "pytorch_model.bin"

# Load model
model = AutoAWQForCausalLM.from_quantized(quant_path, quant_file, fuse_layers=True)
tokenizer = AutoTokenizer.from_pretrained(quant_path, trust_remote_code=True)
streamer = TextStreamer(tokenizer, skip_special_tokens=True)

# Convert prompt to tokens
B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"


def chat(text, max_new_tokens=512):
    prompt = "{bos_token}{b_inst} {system}{prompt} {e_inst} ".format(
        bos_token=tokenizer.bos_token,
        b_inst=B_INST,
        system=f"{B_SYS}{DEFAULT_SYSTEM_PROMPT}{E_SYS}",
        prompt=text,
        e_inst=E_INST,
    )

    tokens = tokenizer(
        prompt,
        return_tensors='pt'
    ).input_ids.cuda()

    # Generate output
    generation_output = model.generate(
        tokens,
        do_sample=True,
        temperature=0.7,
        top_p=0.95,
        top_k=40,
        max_new_tokens=max_new_tokens
    )

    return tokenizer.decode(generation_output[0])
    

雑にchatという関数を生やします。それでは利用してみます。

%%time
print(chat("クマが海辺に行ってアザラシと友達になり、最終的には家に帰るというプロットの短編小説を書いてください。"))

他の例です。

%%time
print(chat("小学生向けの四則演算の問題を5問出してください。最初は問題だけ、次に解答のみをください。"))

これはELYZA自身の性能ですが、日本語の指示に対して正確に回答をくれます。ただ計算結果はたまたまかもしれませんが、全然合ってなかったです。

推論時のリソース消費は、GPU使用率96%張り付き、VRAM 9GB程度の利用となりました。

素のTransformersと比較する

ソースコードは公式そのままです。

VRAM消費は90%張り付き、VRAM14GB強でした。

同じ推論のPromptに対してAWQは2秒弱、素のTransformersだと4.5秒だったので、推論時間が倍以上速くなっていると思います。

おわりに

AWQを使うとVRAMの消費が2/3くらいになって推論速度が倍強になるということが確認できました。

またこれは地味ですが、modelをメモリに載せるまでの時間は素だと2分程度、AWQだと10秒でした。これによって試すカジュアルさも代わりますね。

ただし量子化したことでバカになってる？みたいな気もするので、少なくとも定量的な評価はした方が良さそうです。

以上。