CALM2をAWQフォーマットに変換する in Databricks

思ったよりハマってしまい、時間がかかりました。。。

導入

以下の記事でCyberAgent社のCALM2を動かしました。
TransformersへのAWQフォーマット統合も来ましたので、CALM2をAWQフォーマットに量子化して使ってみます。
（いつもならTheBloke兄貴がAWQフォーマットのリポジトリを提供してくれるのですが、今回なぜかキャンセルされてしまったので）

2023/11/06追記: TheBloke兄貴も改めて上げられていますね。こっち使えばよさそうです。以降の内容は自前でAWQ量子化したい場合に参考にしてください。

Step1. 量子化時のキャリブレーションデータセットの準備

量子化する際にキャリブレーションするための適切なデータセットが必要になりますので、準備します。

CALM2のGPTQモデルを提供されている方がいらっしゃるので、こちらと同じ以下のキャリブレーションデータセットを使わせていただきます。

なお、GPTQモデルはこちら。

適当にノートブックを作成し、huggingfaceからデータを取得する処理を実行します。
必ずCPUクラスタ+MLランタイムで以下の処理を実行してください。

GPUクラスタ・ランタイムだとファイルシステムの違いのせいなのか、load_datasetでエラーを出すため。
（Databricks利用における固有の問題だと思いますが、ここが最大のハマりポイントでした）

from datasets import load_dataset

def load_ja_calib():
    data = load_dataset("mmnga/wikipedia-ja-20230720-1k", split="train")    
    return spark.createDataFrame(data)

load_ja_calib().write.saveAsTable("training.llm.mmnga_wikipedia_ja_20230720_1k")

今回はtraining.llmスキーマにmmnga_wikipedia_ja_20230720_1kというテーブル名で保管します。スキーマは事前に用意をしておいてください。

しかし、Sparkとhuggingface datasetの相互変換、楽ですね。

Step2. モデルのロードと量子化

ここからGPUクラスタで処理を実行します。
私はg5.4xlargeのAWSインスタンスタイプを利用しました。DBRは14.1MLです。

必要なモジュールをインストール。
autoawqは0.1.5と0.1.6で内部処理が結構変わっているようなのですが、最新版(0.1.6)で実行します。
※pytorch >= 2.1.0が求められるため、インストールは結構時間がかかります。

%pip install -U transformers accelerate
%pip install autoawq=="0.1.6"

dbutils.library.restartPython()

キャリブレーションデータを読み込むための関数を定義。

def load_ja_calib():
    df = spark.table("training.llm.mmnga_wikipedia_ja_20230720_1k")
    pdf = df.select("text").toPandas()

    return list(pdf["text"])

AutoAWQでモデルをロード。
CALM2のモデルは事前に保管しておいたものを利用しています。

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

UC_VOLUME = "/Volumes/training/llm/model_snapshots"
uc_dir = "/models--cyberagent--calm2-7b-chat"
model_path = f"{UC_VOLUME}{uc_dir}"

quant_config = { "zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM" }

# Load model
model = AutoAWQForCausalLM.from_pretrained(model_path, device_map="auto", torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

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

問題がなければ、ロード＆量子化処理が実行され、20分程度で完了します。

Step3. 量子化モデルの保管

transformersで読み込むための設定を追加し、モデルを一時フォルダに保管します。

from transformers import AwqConfig, AutoConfig

quant_path = "/tmp/local-calm2-7b-chat-AWQ"

# modify the config file so that it is compatible with transformers integration
quantization_config = AwqConfig(
    bits=quant_config["w_bit"],
    group_size=quant_config["q_group_size"],
    zero_point=quant_config["zero_point"],
    version=quant_config["version"].lower(),
).to_dict()

model.model.config.quantization_config = quantization_config

# Save quantized model
model.save_quantized(quant_path, safetensors=True)
tokenizer.save_pretrained(quant_path)

このままだとクラスタ終了時にファイルが消えるため、Unity Catalog Volumeに保管します。

dbutils.fs.cp("file:/tmp/calm2-7b-chat-AWQ", "/Volumes/training/llm/model_snapshots/models--local--cyberagent-calm2-7b-chat-AWQ-calib-ja-1k", recurse=True)

Step4. 推論する

正常に変換できたかどうかを確認するために、変換後モデルを読み込んで推論してみます。

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import TextStreamer

model_path = "/Volumes/training/llm/model_snapshots/models--local--cyberagent-calm2-7b-chat-AWQ-calib-ja-1k"

model = AutoModelForCausalLM.from_pretrained(model_path, device_map="cuda")
tokenizer = AutoTokenizer.from_pretrained(model_path)
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

def generate_stream_text(prompt: str, max_new_tokens: int = 512) -> str:

    tokens = tokenizer(prompt, return_tensors="pt").input_ids.cuda()

    generation_config = GenerationConfig(
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_k=40,
        top_p=0.95,
        temperature=0.7,
        eos_token_id=model.config.eos_token_id,
        pad_token_id=model.config.pad_token_id,
    )

    generation_output = model.generate(
        tokens,
        streamer=streamer,
        generation_config=generation_config,
    )

    return tokenizer.decode(generation_output[0], skip_special_tokens=True)

npaka先生リスペクトの質問で検証してみます。

prompt = """USER: まどか☆マギカで誰が一番かわいい？
ASSISTANT: """

max_new_tokens = 128
_ = generate_stream_text(prompt, max_new_tokens)

出力

まどか☆マギカのキャラクターは、それぞれ個性的な性格やビジュアルを持っています。

その中でも、一番かわいいと人気が高いのは鹿目まどかでしょう。彼女の優しさや純粋さ、そして何より魔法少女としての強さを持っていることが、多くのファンを魅了している理由でしょう。

また、巴マミや暁美ほむらなども、女性ファンからの人気が高いです。彼女たちの強さや優しさ、そして魔法少女としての責任感などが、多くのファンを引きつける魅力となっています。

他にも、美樹さやかや佐倉杏子など、個性的なキャラクターがたくさんいるのがまどか☆マギ

大丈夫そうです。

英語でも聞いてみます。

prompt = """USER: What is Databricks?
ASSISTANT: """

max_new_tokens = 128
_ = generate_stream_text(prompt, max_new_tokens)

出力

1. Databricks, Inc. is an American cloud computing company that provides a platform for building, managing, and securing data pipelines and workloads.
2. Databricks provides a cloud-based platform that allows users to build and manage data pipelines and workloads for their data lakes, data warehouses, and analytics workflows. It offers a range of tools and services, including Apache Spark, Apache Kafka, Apache Flink, and Apache Beam, to help users manage their data pipelines and workflows.
3. Databricks provides a comprehensive cloud platform for building and managing

内容の正誤はともかく、破綻していない文章ですね。

まとめ

自前でAWQの量子化をしてみました。
日本語LLMの場合、キャリブレーションデータをきちんと準備するのが大事になりそうですが、正直このあたりはド素人なのでベストプラクティスが知りたいですね。。。

主要なLLMはTheBloke兄貴がAWQ変換済みモデルを公開するため、あまり自前変換する機会は多くないかもしれませんが、Transformersと統合されたこともあってニーズは増えそうな気がしています。

個人的にDatabricks上でTransformersを利用する際は、GPTQよりAWQの方が扱いやすいと思っています。
（私の技術的な問題なのですが、Databricks上でGPTQのモデルを正直うまく扱えておらず。。。）

ただ、下のBlogだとGPTQ(というかExllama)が一番性能がよい結果が出ていたり、このあたりうまく使いこなせればなと思います。(vLLMのチューニングが進んだ後の比較とかも見たいところ）

あと、誰かhuggingface hubに上げてください。(自分でやれ)