A5000にサイバーエージェント製220億パラLLMを量子化して入れる

ちょうどいい感じのサイズのモデルがリリースされたので、これをNvidia RTX A5000に入れてみました。22b！！

import transformers
import torch

from transformers import \
    LlamaForCausalLM, GPTNeoXTokenizerFast,\
    BitsAndBytesConfig,\
    TextStreamer
from torch import \
    inference_mode, no_grad

name = "D:/TextGenerationModels/cyberagent_calm3-22b-chat/"
tokenizer = GPTNeoXTokenizerFast.from_pretrained(name)
streamer = TextStreamer(tokenizer)

量子化して読み込む

8bit

quantization_config = BitsAndBytesConfig(load_in_8bit=True)
model = LlamaForCausalLM.from_pretrained(name, quantization_config=quantization_config)

かなりギリギリでしたが入った！

4bitその1

quantization_config = BitsAndBytesConfig(load_in_4bit=True)
model = LlamaForCausalLM.from_pretrained(name, quantization_config=quantization_config)

余裕があります！が、このモデルの推論時に次のような警告が出ました。

警告の内容

Input type into Linear4bit is torch.float16, but bnb_4bit_compute_dtype=torch.float32 (default). This will lead to slow inference or training speed.
warnings.warn(

4bitその2

警告を受けて、量子化計算時のデータ型を指定するようにしました。

quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)
model = LlamaForCausalLM.from_pretrained(name, quantization_config=quantization_config)

4bitその3

多重量子化とかいう設定をONにしてみました。

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True, 
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True)

生成実験

以下のようにテキスト生成を行いました。

コンテクスト：無し

def create_user_message(content: str) -> dict:
    return dict(role="user", content=content)
input_ids = tokenizer.apply_chat_template(
    [create_user_message("Windows11を擬人化した女の子Windows11-tanに関する面白い話を聞かせて")],
    add_generation_prompt=True,return_tensors="pt"
).to(model.device)
output = model.generate(input_ids, max_new_tokens=128, temperature=0.5, streamer=streamer)

コンテクスト:inference_mode

def create_user_message(content: str) -> dict:
    return dict(role="user", content=content)
input_ids = tokenizer.apply_chat_template(
    [create_user_message("Windows11を擬人化した女の子Windows11-tanに関する面白い話を聞かせて")],
    add_generation_prompt=True,return_tensors="pt"
).to(model.device)
with inference_mode():
    output = model.generate(input_ids, max_new_tokens=128, temperature=0.5, streamer=streamer)

生成速度

量子化モード	コンテクスト	速度 $\mathrm{[tokens/s]}$
8bit	無し	$2.23$
8bit	inference_mode	$2.33$
4bitその1	無し	$9.55$
4bitその1	inference_mode	$9.85$
4bitその2	無し	未計測
4bitその2	inference_mode	$11.2$
4bitその3	inference_mode	$7.80$

8bitに比べて4bitの方が格段に生成速度が速いです。容量に余裕があるからでしょうか？また、コンテクスト無しとinference_modeを付けた場合の違いはあまり（わずかにはありそうですけど）見られませんでした。

最も生成速度が速かったのは、量子化モード：4bitその2、コンテクスト：inference_modeの時でした。

生成テキスト

テキスト中において/は改行を表します。

量子化モード：8bit、コンテクスト：無し

Windows 11-tanは、Microsoftが開発した最新のオペレーティングシステム、Windows 11をキャラクター化したもので、そのデザインは親しみやすさと機能性を兼ね備えています。彼女は未来的なデザインの制服を身にまとい、頭にはWindowsのロゴをあしらったキャップをかぶっています。このキャップは、彼女の知識と技術の象徴であり、どんな問題にもスマートに対応する自信を表しています。//Windows 11-tanの性格は明るく前向きで、新しい機能や改善点をユーザーに楽しく紹介する役割を担っています。彼女は特にマルチタスクやエンターテインメントが得意で、ゲームからビデオ編集、そしてクリエイティブな作業まで、あらゆるニーズに応えることが

量子化モード：8bit、コンテクスト：inference_mode

Windows11-tanは、Windows11のオペレーティングシステムをキャラクター化した存在です。彼女は、新しいインターフェースや機能を持つWindows11の特徴を反映しており、そのデザインにはモダンな要素が取り入れられています。//Windows11-tanの外見は、シンプルながらも洗練されたスタイルが特徴です。彼女は、ショートカットヘアで、髪の色は通常、Windowsのロゴカラーである青や緑を基調にしています。また、彼女の瞳には、スタートメニューやタスクバーのアイコンがデザインされており、彼女自身がWindows11の機能を象徴しています。//彼女の性格は、非常にフレンドリーでエネルギッシュです。彼女は

量子化モード：4bitその1、コンテクスト：無し

Windows 11を擬人化した女の子、Windows 11-tanは、Microsoftの最新OSであるWindows 11を親しみやすく、そして楽しく紹介するキャラクターです。Windows 11-tanは、特に若者やアニメファンにアピールするためにデザインされました。//まず、彼女の外見についてですが、Windows 11-tanは現代的でスタイリッシュな服装をしています。たとえば、白を基調としたミニスカートに、肩や腕にフィットするデザインが特徴的なジャケットを羽織っています。このジャケットには、Windows 11のロゴやアイコンがさりげなくデザインされており、一目で彼女のキャラクター

量子化モード：4bitその1、コンテクスト：inference_mode

Windows 11を擬人化した女の子、Windows 11-tanは、その新しいインターフェースや機能を強調する存在として多くのファンに愛されています。彼女は明るく、エネルギッシュで、ユーザーエクスペリエンスの向上を目指す姿が特徴的です。//Windows 11-tanの外見は、最新のWindows 11のデザイン要素を取り入れており、例えばスタートメニューやタスクバーのアイコンが彼女の衣装やアクセサリーに反映されています。彼女はしばしば、モダンなグラフィックデザインとアニメ風のキャラクターを融合させたスタイルで描かれています。//彼女の性格は非常にフレンドリーで、新しい機能やカスタマイズオプション

量子化モード：4bitその2、コンテクスト：無し
まだ実験してません

量子化モード：4bitその2、コンテクスト：inference_mode

Windows 11を擬人化した女の子「Windows 11-tan」は、新しい機能やデザインを取り入れた未来志向のOSとして、多くのファンに愛されています。彼女の魅力は、その洗練されたデザインやユーザーインターフェースだけでなく、親しみやすい性格にもあります。//まず、Windows 11-tanは、明るく元気な性格で、いつも最新のテクノロジーやトレンドに敏感です。彼女は、最新のガジェットやアプリケーションについての情報収集を欠かさず、その知識をユーザーと共有するのが大好きです。//彼女の見た目にも注目が集まっています。Windows 11-tanは、シンプルかつ

生成されたテキストのクオリティもあまり変わらない印象でした。

備考

Windows11-tanとは以下のWindows公式Instagramアカウントで紹介されていたキャラクタです。