CALM2-7BをAWQで4bit量子化&高速推論
CALM2-7Bをベースに開発中のAIアイネスを4bit量子化して高速化してみました。
1. AWQ
Activation-aware Weight Quantization for LLM Compression and Acceleration
という量子化手法で、活性化の大きさに基づいてパラメータの刈り込みと量子化を行います。
ニューラルネットワークの重みのうち、重要な部分を量子化から除外することで、量子化による性能劣化を大幅に削減できるとされています。
なお、ぱぷりか炒めさんのChatGPTとの会話履歴が分かりやすいです。
2. AutoAWQ
AWQにはGPTQと同様にAutoAWQが開発・公開されています。
今回はこれを使用します。
3. キャリブレーションデータセット
AWQの活性化の大きさの測定には、キャリブレーションセットを用います。
今回は、momongaさんのstockmark-13b-GPTQ-calib-ja-1kを参考に、以下のようなデータセットにしました。
izumi-lab/wikipedia-ja-20230720 (アイネスフウジンの口調にしたもの)
312件elyza/ELYZA-tasks-100 (アイネスフウジンの口調にしたもの)
200件
アイネスフウジンの口調への変換は独自に開発したモデルを使用しました。
なお、繰り返し生成の検知や、パターンマッチングを用いて明らかな変換の失敗は取り除いていますが、それ以上のクリーニングは行っていません。
4. 量子化の実行
インスタンスの準備
CUDA 11.8以降と、Compute Capability 7.5が必要です。
T4かA100を用意しましょう。
V100は Compute Capability 7.0なので使えません。
AutoAWQのインストール
`CUDA 11.8.0 + Python 3.10`なので、以下のようにgithubにアップロードされているwheelからインストールする必要があります。
!pip install https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.6/autoawq-0.1.6+cu118-cp310-cp310-linux_x86_64.whl
量子化のオプション
今回はデフォルトの設定を使います。
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }
Zero point
重みの分布を非対称にするために使用するオフセット値を有効にするかどうか。
zero pointを使用すると、量子化の精度や表現力が向上することがある。
(AFPQ: Asymmetric Floating Point Quantization for LLMs - arxiv)Group size
重みをグループ化する際に使用するグループのサイズ。
グループ化とは、重みを連続的なチャンクに分割することで、量子化の精度を向上させること。
group sizeの値は128が推奨、-1は列ごとの量子化を使用することを意味する。
(Quantize 🤗 Transformers models - Hugging Face)量子化の種類
GEMV: GEMM より 20% 高速ですが、バッチ サイズは 1 のみです (大規模なコンテキストには適していません)。
GEMM: 8 未満のバッチ サイズでは FP16 よりもはるかに高速です (大規模なコンテキストに適しています)。
キャリブレーションデータセットのロード
キャリブレーションデータセットとして指定できるのは、以下の2つです。
str: Hugging Face datasetsのようなデータセットのパス
List[str]: ロード済みのデータセット
また、データセットは合計で512件になるようにロードします。
この数字は、AutoAWQのcalib_data.pyのデフォルト値です。
量子化を行うquantizer.pyで変更できるようにしたかったみたいですが、
現状の実装ではサンプル数512件、シーケンス長128で固定のようです。
wiki_dataset = load_japanese_wiki_ines(wiki_path)
tasks100_dataset = load_japanese_tasks100_ines(ELYZA_tasks100_ines_path)
# 乱数生成器を初期化(シード1990)
rng = np.random.default_rng(1990)
calib_dataset = rng.choice(wiki_dataset, size=512-len(tasks100_dataset), replace=False).tolist() + tasks100_dataset
量子化の実行&結果の保存
# Quantize
model.quantize(tokenizer, quant_config=quant_config, calib_data=calib_dataset)
# Save quantized model
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)
おそらくキャリブレーションデータセットの長さにも依ると思いますが、
T4で40分程度かかります。
5. 推論の実行
モデルのロード
Transformersライクに扱えるのは良いですね。
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer, TextStreamer
import re
import torch
# Load model
model = AutoAWQForCausalLM.from_quantized(quant_path, fuse_layers=False)
tokenizer = AutoTokenizer.from_pretrained(quant_path)
推論の実行
VRAMは7.1GB消費しました。
prompt = "まどか☆マギカでは誰が一番かわいい?その理由も説明して。"
token_ids = tokenizer(prompt, return_tensors="pt").input_ids
with torch.no_grad():
output_ids = model.generate(
token_ids.cuda(),
max_new_tokens=512,
do_sample=True,
temperature=0.6,
top_p=0.95,
top_k=100,
pad_token_id=tokenizer.pad_token_id,
bos_token_id=tokenizer.bos_token_id,
eos_token_id=tokenizer.eos_token_id
)
result = tokenizer.decode(output_ids.tolist()[0], skip_special_tokens=True)
print(result)
……あたし、あんまりアニメ見ないんだよねぇ。でもさ、それならやっぱり一番は巴マミちゃんなんじゃないかなあ。彼女の魅力って、とにかく可愛くて優しくて、それでいて芯が強くて、ちょっと不器用なところもあるんだけどそこがまた愛おしいっていう、そんなところだと思うんだ。だから、そういうところがとっても魅力的なんだよ!
この記事が気に入ったらサポートをしてみませんか?