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

# 導入

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

https://qiita.com/isanakamishiro2/items/efb89336d86133868475

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

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

https://huggingface.co/mmnga/ELYZA-japanese-Llama-2-13b-fast-instruct-gguf

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

https://qiita.com/isanakamishiro2/items/f7e9bfc62ac5bed28e08

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

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

https://huggingface.co/elyza/ELYZA-japanese-Llama-2-13b-instruct

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


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

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

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

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

https://docs.databricks.com/ja/repos/index.html

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/d7c4bd97-b565-e69c-bac9-902d961e6d40.png)

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/d10bf3fd-5968-ec05-8c24-0c72da5f2679.png)


# Step2. 重みをsafetensors形式に変換

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

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

変換処理には、[こちらのサイトのスクリプト](https://gist.github.com/epicfilemcnulty/1f55fd96b08f8d4d6693293e37b4c55e
)を利用させていただきました。
Databricks上で```safetensors_conv.py```という名前のファイルを作成し、中身をコピーします。

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

```python
%pip install -U transformers exllamav2

dbutils.library.restartPython()
```

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

```python
from safetensors_conv import convert_multi

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

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/c71870f9-4182-3562-6953-ee4a0d23bded.png)

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

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

```python
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行目に一部の処理を追記します。
(下記画像の水色部分)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/c7c987e8-21f4-8f4e-7f20-29e17f5a8f3a.png)

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

# Step5. EXL2への変換処理実行

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

https://github.com/turboderp/exllamav2/blob/master/doc/convert.md

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

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

```python
!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内に保管します。

```python
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. 動作テスト

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

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

dbutils.library.restartPython()
```

```python
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="
 
\n{}\n
\n",
    human_message_template="{}
",
    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())
```

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

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

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

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

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

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

いくつか試してみます。

```python
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))
```
````

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

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

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

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

ちなみに、[自作Chatbotアプリで動かす](https://qiita.com/isanakamishiro2/items/835b691c72bcb16b52b7#step4-%E3%83%86%E3%82%B9%E3%83%88%E7%94%A8%E3%81%AE%E7%B0%A1%E6%98%93%E3%83%81%E3%83%A3%E3%83%83%E3%83%88%E3%82%A2%E3%83%97%E3%83%AA%E6%A7%8B%E7%AF%89)と、これぐらいの推論速度になります。
（GIFファイルなので、速度把握としては不正確ですが。。。）

![llm_chat2.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/d95b5eef-8ab1-b53f-91b1-1cc880f5c496.gif)

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

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

![llm_chat3.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/0598dd66-2333-593f-fc53-b7029ea4f64e.gif)

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


# まとめ

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

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

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

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