ELYZA-japanese-Llama-2-13BをEXL2形式で量子化して推論する on Databricks

導入

以下の記事でELYZA-japanese-Llama-2-13B(のInstructチューニングモデル)を試しました。

このときはtransformers上で動作させましたが、やはりVRAM消費量を抑えつつ高速推論したい！ということで、最近よく使っているExLlama V2で使えるようにEXL2形式に変換・量子化して使ってみます。

ちなみに、llama.cpp用のGGUFファイルは以下で公開されておりましたので、llama.cpp使いの方はこちらをどうぞ。

GPTQ/AWQ変換モデルは今のところ見当たらず。
手前みそですが、AWQへの変換は以下のやり方でできるんじゃないかと思います。（未検証）

変換およびその後の推論はDatabricks上で実施しました。
DBRは14.2ML、クラスタタイプはg5.2xlarge(AWS)を使っています。

また、今回使うモデルは以下のinstructモデル(非fast版)です。

以下の操作はこのモデルを事前にダウンロードしてある前提で進めます。

Step1. exllamav2のリポジトリをクローン

EXL2形式への変換はexllamav2のリポジトリ中にconvert.pyという変換用のスクリプトがあり、それを使って行います。
スクリプトを実行できるように、まずはexllamav2のリポジトリをクローンします。

シェルやノートブック上でgit cloneすることでリポジトリをクローンできると思いますが、今回は諸所の都合でDatabricksのRepos機能を使ってクローンしました。

Databricks Reposについては、以下を参照してください。

Repos上から、https://github.com/turboderp/exllamav2.gitを指定してリポジトリを追加します。

無事追加が完了すると、所定の場所に以下のようにリポジトリのクローンが作られます。

Step2. 重みをsafetensors形式に変換

EXL2形式に変換するモデルの重みは、safetensors形式で保存されている必要があるのですが、ELYZA-japanese-Llama-2-13Bはbin形式でしか重みが公開されていません。

そのため、最初にbin形式の重みファイルをsafetensors形式に変換します。

変換処理には、こちらのサイトのスクリプトを利用させていただきました。
Databricks上でsafetensors_conv.pyという名前のファイルを作成し、中身をコピーします。

次にノートブックを作成し、まずは必要（かもしれない）パッケージをインストール。

%pip install -U transformers exllamav2

dbutils.library.restartPython()

上記で作成したスクリプトファイルからconvert_multi関数をインポートし、事前にダウンロードしたモデルのパスを指定して実行。

from safetensors_conv import convert_multi

convert_multi("/Volumes/training/llm/model_snapshots/models--elyza--ELYZA-japanese-Llama-2-13b-instruct/", False)

正常に終了すると、以下のように.safetensorsのファイルが作られます。
(私の環境ではだいたい30分ぐらいかかりました)

Step3. 変換後ファイルの出力先フォルダ作成

EXL2変換処理中の出力ファイルのフォルダを事前に作成しておきます。

dbutils.fs.mkdirs("file:/tmp/exl2/")
dbutils.fs.mkdirs("file:/tmp/models/local--ELYZA-japanese-Llama-2-13b-instruct/4.0bpw/")

Step4. (Option?) exllamav2リポジトリのconvert.pyファイルを一部修正

2023/12/29現在、exllamav2のmainブランチには少し問題があるようで、EXL2変換を正常に行うためには少し細工が必要です。（tag=v0.0.11を指定したリポジトリであればおそらくここの作業は不要。また、将来的には修正されると思います)

Step1でクローンしたリポジトリの直下にあるconvert.pyを開き、ExLlamaV2Configを作成する処理の下、67・68行目に一部の処理を追記します。
(下記画像の水色部分)

※　現在のmasterブランチではnum_expertsのデフォルトがNoneになっているのですが、これが原因で変換処理中にエラーが発生します。そのため、強制的にnum_experts=1を指定しています。Mixtralモデルの変換をする場合はこれを除去してください。

Step5. EXL2への変換処理実行

準備が整いましたので、EXL2の変換スクリプトを実行します。
※　スクリプトに関するDocumentは以下を参照してください。

ノートブックのセル内にexllamav2リポジトリ内のconvert.pyを実行する処理を記載して実行します。
パラメータとして、Step3で作成した出力先フォルダを指定しています。
今回は4.0bitでの量子化を-bオプションで指定しました。

また、キャリブレーションのデータセットは未指定(デフォルト使用)です。
ここは日本語LLM用のキャリブレーションファイルを準備するほうがよい性能になるかもしれません。

!cd /Workspace/Repos/リポジトリをクローンした場所/exllamav2 && \
    python convert.py \
    --no_resume \
    -i /Volumes/training/llm/model_snapshots/models--elyza--ELYZA-japanese-Llama-2-13b-instruct/ \
    -o /tmp/exl2/ \
    -cf /tmp/models/local--ELYZA-japanese-Llama-2-13b-instruct/4.0bpw/ \
    -b 4.0 

私の環境下では約50分かかりました。

最後に、変換されたファイル類を永続化するために、Unity CatalogのVolumes内に保管します。

dbutils.fs.cp(
    "file:/tmp/models/local--ELYZA-japanese-Llama-2-13b-instruct/4.0bpw/", 
    "/Volumes/training/llm/model_snapshots/models--local--ELYZA-japanese-Llama-2-13b-instruct-4.0bpw-h6-exl2/", 
    recurse=True,
)

Step6. 動作テスト

正しく変換できたか確認してみます。準備周りの説明は割愛。

%pip install -U -qq transformers accelerate "exllamav2>=0.0.11" langchain sentencepiece

dbutils.library.restartPython()

from exllamav2_chat import ChatExllamaV2Model
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.prompts.chat import (
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda

model_path = "/Volumes/training/llm/model_snapshots/models--local--ELYZA-japanese-Llama-2-13b-instruct-4.0bpw-h6-exl2"
chat_model = ChatExllamaV2Model.from_model_dir(
    model_path,
    cache_max_seq_len=4096,
    system_message_template="[INST] <<SYS>>\n{}\n<</SYS>>\n",
    human_message_template="{}[/INST]",
    ai_message_template="{}",    
    temperature=0.0001,
    max_new_tokens=1024,
    repetition_penalty = 1.15,
    low_memory=True,
    cache_8bit=True,
)

system_template = "あなたは誠実で優秀な日本人のアシスタントです。"

prompt = ChatPromptTemplate.from_messages(
    [
        SystemMessagePromptTemplate.from_template(system_template),
        HumanMessagePromptTemplate.from_template("{query}"),
        AIMessagePromptTemplate.from_template(" "),
    ]
)

chain = ({"query": RunnablePassthrough()} 
         | prompt 
         | chat_model
         | StrOutputParser())

では、推論させてみましょう。

query = "明日の東京の天気は?"
async for s in chain.astream(query):
    print(s, end="", flush=True)

出力

私はELYZAによって訓練されたAIであり、現状の情報を元に回答します。

2023年4月15日 (金) から4月17日 (日) の間は、曇りや雨の予報となっています。

ただし、天気は変わりやすいため、最新の情報については公式サイトへのリンクからご確認ください。

破綻のない文章が生成できています。どうやらうまく変換・推論実行できているようですね。

いくつか試してみます。

query = "ランダムな10個の要素からなるリストを作成してソートするコードをPythonで書いてください。"
async for s in chain.astream(query):
    print(s, end="", flush=True)

出力

以下がそれに当てはまるものです。
```python
import random

random_list = [3, 6, 8, 2, 5, 4, 7, 1, 9, 0]

# 引き数なしのsorted()でソート
print(sorted(random_list))
```

ちょっと回答がシンプルすぎるかな？
パラメータの問題な気もしますが。。。

query = "まどか☆マギカでは誰が一番かわいい?"
async for s in chain.astream(query):
    print(s, end="", flush=True)

出力

まどか☆マギカに登場するキャラクター全員が魅力的でかわいいと言えます。しかし、以下の3名は特にかわいく見えることが多いようです: 
- 綾瀬 まどか (まどか) 
- 杏子 (あんず) 
- 長門有希 (ながもありさき)

さすがにハルシネーション満載ですね。
量子化前のモデルでも、そこまで正確に答えれてないと思うので、量子化の影響なのかどうなのか悩む。

ちなみに、自作Chatbotアプリで動かすと、これぐらいの推論速度になります。
（GIFファイルなので、速度把握としては不正確ですが。。。）

非fastモデルなのですが、十分早いかなと思います。
（fastモデルだとさらに早くなることが期待できるので、後で試す予定です）

(2023/12/29追記) fast-instructモデルでも試してみました。

生成量が異なるのでわかりづらいですが、明らかにこちらの方が生成速度は速くなっています。

まとめ

ELYZA-japanese-Llama-2-13BをEXL2形式に変換し、exllamav2で推論させてみました。

初めてEXL2への変換をしてみたのですが、最初からsafetensorsのモデルがあれば割と簡単に変換できるのかなと思います。
ただ、CTranslate2と違って変換にもGPUを積んだクラスタが必要なのは注意。

また、キャリブレーションデータを別で用意する方がよいのかどうか、など気を付けないといけないポイントがいろいろありそうです。このあたりはもっと情報を集めたり、自分で探っていく必要があるかなと思いました。

huggingface上の有名モデルはこの方がEXL2変換・公開されていたりするので、あまり自分でやる機会は少ないかもしれませんが、日本語LLMはあまり対象にならなかったりするので、自分で変換できるスキルは身に着けておいても損はないですね。