概要

本稿はモデルの 8bit 量子化について、です
せっかくなので、昨日(2024年3月21日)、アナウンスされた Rakuten AI 7B Chat でお試ししてみます。

実行環境

NVIDIA RTX A5000
Python 3.11.4

モデルの8ビット量子化

16bit(このモデルは BF16テンソル) と 8bit への量子化の前後でメモリ使用量の変化をみてみます。

load_in_8bit は BitsAndBytesConfig オブジェクトを経由して使う

AutoModelForCausalLM.from_pretrained で直接指定すると以下のように「deprecatedですよ」と指摘されるので quantization_config にくるんでつかいました

The `load_in_4bit` and `load_in_8bit` arguments are deprecated and will be removed
in the future versions. 
Please, pass a `BitsAndBytesConfig` object in `quantization_config` argument instead.

以前は load_in_8bit 引数を AutoModelForCausalLM.from_pretrained に直接指定していましたが、冬眠から明けたらやりかたが少し変化していました。

今後は quantization_config を経由して指定します。

必要なパッケージの追加

pip install accelerate
pip install -i https://pypi.org/simple/ bitsandbytes

モデル 8bit 量子化して実行

以下のようなコードを書きました。
ミソは↓です。

quantization_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(model_path, quantization_config=quantization_config, device_map="auto")

コード全体

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


def cout_memory_availability(label):
    """
    GPUのメモリ利用状況と現在時刻を表示する
    :param label: メモリ状況のラベル
    :return: なし
    """

    # 現在時刻を取得して指定されたフォーマットでフォーマットする
    now = datetime.now()
    formatted_now = now.strftime("%Y年%m月%d日%H時%M分%S秒")

    total_memory = torch.cuda.get_device_properties(0).total_memory / (1024 ** 2)  # 総メモリをMB単位で取得
    reserved_memory = torch.cuda.memory_reserved(0) / (1024 ** 2)  # 予約済みメモリをMB単位で取得
    available_memory = total_memory - reserved_memory  # 利用可能なメモリを計算

    print(f"{formatted_now} - {label}: GPUの総メモリ: {total_memory:.2f} MB")
    print(f"{formatted_now} - {label}: 予約済みメモリ: {reserved_memory:.2f} MB")
    print(f"{formatted_now} - {label}: 利用可能なメモリ: {available_memory:.2f} MB")


set_seed(42) # シード固定

cout_memory_availability("モデル読み込み前")

# 8ビット量子化の設定を作成
quantization_config = BitsAndBytesConfig(load_in_8bit=True)

model_path = "Rakuten/RakutenAI-7B-chat"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, quantization_config=quantization_config, device_map="auto")
model.eval()

cout_memory_availability("モデル読み込み後")

system_message = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: {user_input} ASSISTANT:"

user_input_text = "「馬が合う」はどう言う意味ですか"
input_req = system_message.format(user_input=user_input_text)
input_ids = tokenizer.encode(input_req, return_tensors="pt").to(device=model.device)
tokens = model.generate(
    input_ids,
    max_new_tokens=1024,
    do_sample=True,
    pad_token_id=tokenizer.eos_token_id,
)
assistant_out_text = tokenizer.decode(tokens[0][len(input_ids[0]):], skip_special_tokens=True)
print(f"USER:\n{user_input_text}")
print(f"ASSISTANT:\n{assistant_out_text} ({len(tokens[0])}トークン)")
print()
print()

cout_memory_availability("推論後")

8bit 量子化の実行結果

モデル読み込み後のGPUメモリ消費量は8GB程度でした(推論後は 12GB 程度※)
（※推論後に torch.cuda.empty_cache() を呼び出し未使用のメモリを解放してあげると再び 8GB 程度に戻った）
トークン生成速度: 6.6トークン/秒

2024年03月22日19時35分52秒 - モデル読み込み前: GPUの総メモリ: 24563.50 MB
2024年03月22日19時35分52秒 - モデル読み込み前: 予約済みメモリ: 0.00 MB
2024年03月22日19時35分52秒 - モデル読み込み前: 利用可能なメモリ: 24563.50 MB
2024年03月22日19時36分03秒 - モデル読み込み後: GPUの総メモリ: 24563.50 MB
2024年03月22日19時36分03秒 - モデル読み込み後: 予約済みメモリ: 8058.00 MB
2024年03月22日19時36分03秒 - モデル読み込み後: 利用可能なメモリ: 16505.50 MB
USER:
「馬が合う」はどう言う意味ですか
ASSISTANT:
「馬が合う」とは、人との間には何か相性が合わないものを感じているが、何となく気が合う、という意味です。
「馬が合う」と言う言葉は、一般的に、人同士の間には、相性が合う合わないがあることを指しています。友情や恋愛において、二人の間に相性の良し悪しがある場合に使われることも多く、「馬が合う」と言う言葉で、何となく「気の置けない間柄」であることを表現することができるのです。
「馬」は、古くからのことわざで使われる言葉で、「馬があう・あわない」と言う言葉で、相性が良いか悪いかを表すために使われます。自然発生的に仲良くなる関係、というような意味があります。
「馬が合う」と言う言葉を使うと、相性が合っている間柄でないと、分かりにくく伝わらないことがありますが、相性が良いときに使うと、「何となく気が合う」という意味を表現することができます。 (258トークン)


2024年03月22日19時36分42秒 - 推論後: GPUの総メモリ: 24563.50 MB
2024年03月22日19時36分42秒 - 推論後: 予約済みメモリ: 12172.00 MB
2024年03月22日19時36分42秒 - 推論後: 利用可能なメモリ: 12391.50 MB

4ビット量子化で実行してみる

4ビットの場合も同様に以下のようにする

quantization_config = BitsAndBytesConfig(load_in_4bit=True)

モデル読み込み後のGPUメモリ消費量は5GB程度でした(推論後も 5GB 程度)
トークン生成速度: 24.3トークン/秒

2024年03月22日19時40分21秒 - モデル読み込み前: GPUの総メモリ: 24563.50 MB
2024年03月22日19時40分21秒 - モデル読み込み前: 予約済みメモリ: 0.00 MB
2024年03月22日19時40分21秒 - モデル読み込み前: 利用可能なメモリ: 24563.50 MB

2024年03月22日19時40分32秒 - モデル読み込み後: GPUの総メモリ: 24563.50 MB
2024年03月22日19時40分32秒 - モデル読み込み後: 予約済みメモリ: 5202.00 MB
2024年03月22日19時40分32秒 - モデル読み込み後: 利用可能なメモリ: 19361.50 MB

USER:
「馬が合う」はどう言う意味ですか
ASSISTANT:
「馬が合う」とは、人同士が気が合う、気が似ていると感じ合わることです。
または、人の気質や性格が似ていることを指します。
例文としては、「私と彼は馬が合うようで、いつもよく気が合う」などという使えます。
「馬が合う」は、旧来は、馬同士を比較して類似性を表して使ったものですが、そこから人が人に対して似たような気質や性格を比べて、類似性や合うと感じることを表す言葉として使われるようになりました。
「馬が合う」はもちろんいい意味でも悪い意味でも使える言葉です。例えば、「私はいつもあなたと馬が合うのでいつもいい物を分けてあげたのにいい物をなくした」などと言うとしたら、自分といつもあ気が合うこと、いつも気に入るものをこっそり分けてあげたのにそれを盗んだといった悪口の意味に受け取られます。
使い方によっていいこと悪気のどちらの意味にも取れる曖昧な言葉です。
「合わせる」は「調べる」「比較する」という意味があります。
「馬が合うという」は、「気が合う」と言う意味になります。意味と同じ使い方ができる言葉です。 (316トークン)


2024年03月22日19時40分45秒 - 推論後: GPUの総メモリ: 24563.50 MB
2024年03月22日19時40分45秒 - 推論後: 予約済みメモリ: 5724.00 MB
2024年03月22日19時40分45秒 - 推論後: 利用可能なメモリ: 18839.50 MB

16ビットで実行してみる

上記のコードのモデル読み込み部分を以下のように torch.bfloag16 を指定してやります

model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map="auto")

実行結果

モデル読み込み後のGPUメモリ消費量は15GB程度
トークン生成速度: 33トークン/秒

2024年03月22日16時36分02秒 - モデル読み込み前: GPUの総メモリ: 24563.50 MB
2024年03月22日16時36分02秒 - モデル読み込み前: 予約済みメモリ: 0.00 MB
2024年03月22日16時36分02秒 - モデル読み込み前: 利用可能なメモリ: 24563.50 MB

2024年03月22日16時36分11秒 - モデル読み込み後: GPUの総メモリ: 24563.50 MB
2024年03月22日16時36分11秒 - モデル読み込み後: 予約済みメモリ: 14834.00 MB
2024年03月22日16時36分11秒 - モデル読み込み後: 利用可能なメモリ: 9729.50 MB
USER:
「馬が合う」はどう言う意味ですか
ASSISTANT:
「馬が合う」とは、人との間には何か得体に欠つけた不健全なものがあるようなことを言う時によく使われる表現です。
「馬が合う」の由来は、人間にとっては大事な協力者で友情のある仲間の「馬」を、女らしい、男らしいといった性別で表すことからです。
馬は人と同じように、性格もよく似た仲間が、互いを必要とし、互いの分をかためていく関係になるそうです。
それが「馬が合う」と表現されてきた由来ですが、現代では「馬が合う」とよく言われるのは、上記の由来とは別に、仲良くなった人同士が、自然に触れ合い、気楽に過ごせる友人関係のことを指していわれているそうです。
「馬が合う」と表現されている例をあげると、「男性同士で馬が合う友情ができる」「私はこの人と馬が合うような気がする」などがあります。
「馬が合う」という表現は、仲良くなった人との間に、何かしら得体に欠けた不健全な何かがあることを表していますが、上記の由来や例に挙げられたような、男女間や人との交流、協力し合うことの健全性を意味することが多く、大事な友情を温かく表現してくれる表現でもあります。 (327トークン)


2024年03月22日16時36分21秒 - 推論後: GPUの総メモリ: 24563.50 MB
2024年03月22日16時36分21秒 - 推論後: 予約済みメモリ: 14964.00 MB
2024年03月22日16時36分21秒 - 推論後: 利用可能なメモリ: 9599.50 MB

まとめ

7Bのモデル(RakutenAI-7B-chat)を使って、 8bit 量子化をやってみました
16bit と 8bit （と4bit) の結果比較
- 8bit
  - GPUメモリ消費量: 8GB (推論後に12GB)
  - トークン生成速度: 6.6トークン/秒
- 16bit
  - GPUメモリ消費量: 15GB
  - トークン生成速度: 33トークン/秒
- 4bit
  - GPUメモリ消費量: 5GB (推論後も5GB)
  - トークン生成速度: 24.3トークン/秒

8bit 量子化すると、GPUメモリは理論値(7GB強い)に近いサイズになりましたが推論後は大きく増えました。（torch.cuda.empty_cache()により解放され使用メモリは8GB程度まで戻る)。トークン生成速度はBF16に比較して80%も遅くなりました。生成された文章の質ですが、大きく劣化したようには見えませんでした。
　
4bit 量子化すると GPUメモリは大幅に抑えられました。トークン生成速度は 25%程度おそくなりましたが、8bit量子化に比べるとかなり速かったです。

今後

8bit 量子化時のトークン生成速度の遅さは何に起因するのか、条件を変えて確かめてみようとおもいます。

シード値を変える
BitsAndBytesConfig量子化パラメータを変える（↓あたり）

        llm_int8_threshold=6.0,
        llm_int8_skip_modules=None,
        llm_int8_enable_fp32_cpu_offload=False,
        llm_int8_has_fp16_weight=False,

モデルを変える

LLMモデル "RakutenAI-7B-chat" を 8bit 量子化して実行してみた

概要