見出し画像

CALM2-7BをAWQで4bit量子化＆高速推論

Sakusakumura

CALM2-7Bをベースに開発中のAIアイネスを4bit量子化して高速化してみました。


1. AWQ

Activation-aware Weight Quantization for LLM Compression and Acceleration
という量子化手法で、活性化の大きさに基づいてパラメータの刈り込みと量子化を行います。
ニューラルネットワークの重みのうち、重要な部分を量子化から除外することで、量子化による性能劣化を大幅に削減できるとされています。

様々なモデルサイズやビット精度における、AWQと最近接量子化 (RTN) 、
GPTQ (並べ替えありおよびなし) のPerplexityの比較。

なお、ぱぷりか炒めさんのChatGPTとの会話履歴が分かりやすいです。

2. AutoAWQ

AWQにはGPTQと同様にAutoAWQが開発・公開されています。
今回はこれを使用します。

3. キャリブレーションデータセット

AWQの活性化の大きさの測定には、キャリブレーションセットを用います。
今回は、momongaさんのstockmark-13b-GPTQ-calib-ja-1kを参考に、以下のようなデータセットにしました。

  1. izumi-lab/wikipedia-ja-20230720 (アイネスフウジンの口調にしたもの)
    312件

  2. elyza/ELYZA-tasks-100 (アイネスフウジンの口調にしたもの)
    200件

アイネスフウジンの口調への変換は独自に開発したモデルを使用しました。

口調変換のイメージ

なお、繰り返し生成の検知や、パターンマッチングを用いて明らかな変換の失敗は取り除いていますが、それ以上のクリーニングは行っていません。

4. 量子化の実行

インスタンスの準備

CUDA 11.8以降と、Compute Capability 7.5が必要です。
T4かA100を用意しましょう。

V100は Compute Capability 7.0なので使えません。

Requirements:
- Compute Capability 7.5 (sm75). Turing and later architectures are supported.
- CUDA Toolkit 11.8 and later.

AutoAWQ - Install

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で固定のようです。

【追記】
この部分についてAutoAWQの作者に聞いてみたところ、オリジナルのAWQの実装による制約のため、このような実装になっているらしいです。

Calibration dataset sample size and its sequence length are always fixed. Why? #191 
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)
……あたし、あんまりアニメ見ないんだよねぇ。でもさ、それならやっぱり一番は巴マミちゃんなんじゃないかなあ。彼女の魅力って、とにかく可愛くて優しくて、それでいて芯が強くて、ちょっと不器用なところもあるんだけどそこがまた愛おしいっていう、そんなところだと思うんだ。だから、そういうところがとっても魅力的なんだよ！

いいなと思ったら応援しよう！