量子化による日本語LLMの実行

モチベーション

この記事において量子化については、否定的な言い方で終わったが、少し調べてみると興味深い領域であると思い直し、量子化について実験したので、その内容をここにまとめる。

参考

今回の量子化は、bitsandbytesを使ったが、他にも手法はあるようだ。関係情報を以下に列挙する。

Hugging Face 先ずはここから。英語なのが少々難儀。
日本語LLMの推論速度検証インターンにより日本語LLM推論速度に関するレポート。よくまとまっており、量子化についても触れられている。
【ELYZA-japanese-Llama-2-13b】東大スタートアップがGPT3.5を超えるLLMを開発！使い方〜実践まで ELYZA-japanese-Llama-2-13bの紹介記事。コードもこれを参考にした。

はじめに

上記・参考の「【ELYZA-japanese-Llama-2-13b】東大スタートアップがGPT3.5を超えるLLMを開発！使い方〜実践まで」の記事に掲載のコードをRTX A4000（GPUメモリ：16GB）で実行した。その際には、「WARNING:root:Some parameters are on the meta device device because they were offloaded to the cpu.」とのメッセージが表示されながらも回答を得た。そのため、TITAN X（GPUメモリ：12GB）では実行を諦め、DeepSpeedによるモデル並列化を調べようと思った。

少し調べる内、量子化によってモデルを圧縮しつつ精度を保つための手法／ライブラリがあることを知ったので、早速試してみることにした。

今回は、bitsandbytesライブラリを使う手法を試したみた。

なお、実行したJupyterLabのdockerコンテナは、この記事のものと同じ。

コードを変更

元のコード

上記「参考」の「【ELYZA-japanese-Llama-2-13b】東大スタートアップがGPT3.5を超えるLLMを開発！使い方〜実践まで」の元のコード部分は次のとおり。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"
text = "仕事の熱意を取り戻すためのアイデアを5つ挙げてください。"

model_name = "elyza/ELYZA-japanese-Llama-2-13b-instruct"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    use_cache=True,
    device_map="auto",
    offload_folder = "/content/ELYZA-japanese-Llama-2-13b-instruct",
    low_cpu_mem_usage=True,
)
model.eval()

上記の部分を次のとおり変更した。実は、各引数の詳細の意味はまだよく分かってないが、試行錯誤の結果、こうすればとりあえず動作したというレベルである。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"
text = "仕事の熱意を取り戻すためのアイデアを5つ挙げてください。"

model_name = "elyza/ELYZA-japanese-Llama-2-13b-instruct"

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    quantization_config=quantization_config,
).eval()

実行結果

この記事を書いているということは、上記の変更の結果、TITAN Xという12GBのメモリで実行できた。

量子化の効果は素晴らしい。しかし、残念ながら動作の詳細についての理解はできていない。今後の課題である。

いくつもの質問ができるように、モデルの構築と質問部分を分けて次のようなコードに変更した。

import torch
from transformers import BitsAndBytesConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "elyza/ELYZA-japanese-Llama-2-13b-instruct"

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    quantization_config=quantization_config,
).eval()

オリジナルのコードに対して、次のとおりmax_new_tokensを1024に増やした。

B_INST, E_INST = "[INST]", "[/INST]"
B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。"
text = "仕事の熱意を取り戻すためのアイデアを5つ挙げてください。"

while text != "":
    text = input("質問を入力してください：")
    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,
    )
    token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")

    with torch.no_grad():
        output_ids = model.generate(
            token_ids.to(model.device),
            max_new_tokens=1024,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    output = tokenizer.decode(output_ids.tolist()[0][token_ids.size(1) :], skip_special_tokens=True)
    print(output)
    print("\n\n\n")

まとめ

量子化については、参考に掲載したHugging Faceにあるとおり、他の手法もあるので、そちらも今後試してみたい。