Zenn
🎃

sarashina2.2-1b-instruct-v0.1をAWQで量子化する

に公開
量子化
LLM
awq
tech

３行で

  • AutoAWQを用いて、sbintuitions/sarashina2.2-1b-instruct-v0.1を量子化しました。
  • 量子化したモデルのhuggingfaceのrepoは以下です。

https://huggingface.co/sinchir0/sarashina2.2-1b-instruct-v0.1-awq

  • 量子化した結果、CPUメモリは約55%削減、GPUメモリは約53%削減できました。

詳細

量子化の理解を深めたく、AutoAWQを用いた量子化をsarashina2.2-1b-instruct-v0.1に対して行いました。

https://huggingface.co/sbintuitions/sarashina2.2-1b-instruct-v0.1

環境構築

vast.aiの1xA40(GPUメモリ45GB)環境を用いました。テンプレートは「NVIDIA CUDA (Ubuntu)」を利用しました。
GPUメモリについて、大きめの環境を使いましたが、推論含め小さくても問題なく行えるかと思います。

以下は環境構築に用いたコードです。参考として。

python3 -m pip install --upgrade pip # pipのバージョンアップ
pip install autoawq

量子化コード

ライブラリはAutoAWQを利用します。

https://github.com/casper-hansen/AutoAWQ

AutoAWQのREADMEを、model_pathとquant_pathだけ変更しました。

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = 'sbintuitions/sarashina2.2-1b-instruct-v0.1'
quant_path = 'sarashina2.2-1b-instruct-v0.1-awq'
quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }

# Load model
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# Quantize
model.quantize(tokenizer, quant_config=quant_config)

# Save quantized model
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

print(f'Model is quantized and saved at "{quant_path}"')

推論 & 評価

コード

以下コードを使いました。公式に載っているコードについて、与えるプロンプトを、技術ブログにて紹介された例に変更し、メモリ使用量を計測できるようにしました。
https://huggingface.co/sbintuitions/sarashina2.2-1b-instruct-v0.1#how-to-use

import torch
import os, psutil
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, set_seed

def print_gpu_memory(tag=""):
    allocated = torch.cuda.memory_allocated() / 1024**2
    reserved = torch.cuda.memory_reserved() / 1024**2
    peak = torch.cuda.max_memory_allocated() / 1024**2
    print(f"[{tag}] Allocated: {allocated:.2f} MB | Reserved: {reserved:.2f} MB | Peak: {peak:.2f} MB")

def print_cpu_memory(tag=""):
    process = psutil.Process(os.getpid())
    mem = process.memory_info().rss / 1024**2
    print(f"[{tag}] CPU RSS Memory: {mem:.2f} MB")

# 初期化 & メモリリセット
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()

print_cpu_memory("Before loading")
print_gpu_memory("Before loading")

# モデルのロード
model_name = "sbintuitions/sarashina2.2-1b-instruct-v0.1"
# model_name = "sinchir0/sarashina2.2-1b-instruct-v0.1-awq"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)
chat_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
set_seed(123)

print_cpu_memory("After loading")
print_gpu_memory("After loading")

# ユーザーの入力
text = """
1. 彼は比類のない陸上選手だ。
2. 彼は比較的に良い陸上選手だ。

1の文が難しいので2の文に直そうと思っているのですが、これってあってますか？
"""
user_input = [{"role": "user", "content": text}]

# 応答生成前にメモリ統計をリセット
torch.cuda.reset_peak_memory_stats()

# 応答生成
responses = chat_pipeline(
    user_input,
    max_length=128,
    do_sample=True,
    num_return_sequences=3,
)

print_cpu_memory("After generation")
print_gpu_memory("After generation")

# 応答を表示
for i, response in enumerate(responses, 1):
    print(f"Response {i}: {response['generated_text']}")

評価

まず、元のモデルの出力例を1つ貼ります。

はい、その変換は適切です。「比類のない」という言葉は、他のどんな選手とも比較できないほど優れていることや、群を抜いていることを指します。一方、「比較的に良い」という表現は、平均的または標準的に優れていることを意味します。したがって、最初の文を「彼は比較的に良い陸上選手だ」と言い換えることで、彼が平均的なレベルにあることをより明確に伝えることができます。この変更により、あなたの意図する意味を正確に伝えることができます。

次に、量子化後の出力例を1つ貼ります。

はい、その変換は適切です。最初の文「彼は比類のない陸上選手だ」は、他のどの選手と比較しても群を抜いているという意味なので、 quite exceptional（非常に優れた）や inimitable（比較するものがないほど優れている）といったニュアンスが含まれます。一方、二番目の文「彼は比較的に良い陸上選手だ」は、特に群を抜いているわけではなく、標準的なレベルよりも高い能力を持っているが、特に exceptional（例外的に良い）というわけではないことを示しています。したがって、二番目の文は最初の文の意味を正確に伝えるために適切な言い換えとなります。

量子化前後でほぼ同様の回答となりました。

メモリ

元のモデルのメモリ使用量は以下となりました。

[After loading] CPU RSS Memory: 2138.09 MB
[After loading] Allocated: 2725.18 MB | Reserved: 2766.00 MB | Peak: 2725.18 MB
[After generation] CPU RSS Memory: 2221.10 MB
[After generation] Allocated: 2733.30 MB | Reserved: 2836.00 MB | Peak: 2795.17 MB

次に、量子化後のモデルのメモリ使用量は以下となりました。

[After loading] CPU RSS Memory: 808.30 MB
[After loading] Allocated: 1236.35 MB | Reserved: 1238.00 MB | Peak: 1236.35 MB
[After generation] CPU RSS Memory: 1567.81 MB
[After generation] Allocated: 1244.47 MB | Reserved: 1392.00 MB | Peak: 1328.09 MB

Peak同士で比較すると、CPUメモリは約55%削減、GPUメモリは約53%削減できました。

終わりに

環境構築から量子化まで、AutoAWQライブラリを使うことで、30分程度でできました。今後は、量子化ライブラリにてどのようなことが具体的に行われているかの理解を深めていきたいと思います。

Discussion