Llama 3.1 8B Instruct を量子化して試せるようにした

昨晩リリースされた Llama 3.1 を AWQ 4Bit 量子化して使えるようにしました。

公式ページの説明だけだと、量子化がうまくいかなかったので、そこを補足的に記述しました

これから、色々試して順次追記していこうとおもいます

【デモ】チャット公開しました

本記事でつくった量子化モデル＋ ChatStream をつかって以下URLにて Llama3.1 8B Instruct をチャットできるようにしています

推論エンジンは vLLM をつかっていますので、そこそこ高速です。

↓がチャットデモです

環境

WSL + NVIDIA A5000

量子化の方法

量子化アルゴリズムは AWQ を使用しました。
以下の AutoAWQをつかいます

AutoAWQと関連パッケージインストール方法

pip install -q --upgrade transformers autoawq accelerate

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "meta-llama/Meta-Llama-3.1-8B-Instruct"
quant_path = "/mnt/d/Meta-Llama-3.1-8B-Instruct-AWQ"
quant_config = {
  "zero_point": True,
  "q_group_size": 128,
  "w_bit": 4,
  "version": "GEMM",
}

# Load model
model = AutoAWQForCausalLM.from_pretrained(
  model_path, low_cpu_mem_usage=True,
  use_cache=False,
  device_map="auto"  # わすれずにこれを追加
)

tokenizer = AutoTokenizer.from_pretrained(model_path)

# 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}"')

量子化チェックポイントの作成は約20分でおわりました。

ハマったポイント

ハマったというほどではないですが、２点ほど注意点があります

その１. transformers パッケージは最新に！

transformers が 4.43.1 未満だと、以下のエラーがでるので注意です。

    raise ValueError(
ValueError: `rope_scaling` must be a dictionary with two fields, `type` and `factor`, got {'factor': 8.0, 'low_freq_factor': 1.0, 'high_freq_factor': 4.0, 'original_max_position_embeddings': 8192, 'rope_type': 'llama3'}

transformers を最新にしましょう

いつも同じ環境をつかってると transformers が古くってなりがちですね。

pip install --upgrade transformers

その２.モデルをロードするときは、デバイス指定をいれる

AWQ実行で以下のようなエラーがでました。

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat2 in method wrapper_CUDA_bmm)

モデルロードに device_map="auto" を追加しましょう。

model = AutoAWQForCausalLM.from_pretrained(
  model_path, low_cpu_mem_usage=True,
  use_cache=False,
  device_map="auto"  # わすれずにこれを追加
)

これで無事4Bit量子化ができます。