DatabricksでExLlamaV2を試してみる

# 導入

ローカルLLMを実行するためのモジュールとしては[Transformers](https://huggingface.co/docs/transformers/index)が代表格ですが、推論だけに限れば他のモジュールを使う方がメモリ効率や推論速度は良かったりします。
人気があるモジュールとしては[llama.cpp](https://github.com/ggerganov/llama.cpp)や[AutoGPTQ](https://github.com/PanQiWei/AutoGPTQ)なんかがあります。(AutoGPTQはTransformersからも利用可能)

Databricks上でllama.cppを動かす方法は↓でも取り上げられています。(いつもお世話になっています)

https://qiita.com/taka_yayoi/items/2d6ee09600d89b58c523

私がよく使っている[CTranslate2](https://github.com/OpenNMT/CTranslate2)もDatabricks上で簡単に動作し、使い勝手もよいのですが、利用できないモデルがあったり8bit未満の量子化に対応してないなど、困る面もあります。

他の選択肢も持っておくべきという考えで、一部で話題(?)のExLlamaをDatabricks上で試してみます。

# ExLlamaとは

https://github.com/turboderp/exllama

> A standalone Python/C++/CUDA implementation of Llama for use with 4-bit GPTQ weights, designed to be fast and memory-efficient on modern GPUs.
> Disclaimer: The project is coming along, but it's still a work in progress!
>
> 邦訳：
> LlamaのスタンドアロンPython/C++/CUDA実装で、4ビットGPTQウェイトで使用され、最新のGPUで高速でメモリ効率が良いように設計されています。
免責事項：プロジェクトは進んでいますが、まだ未完成です！

GPTQ形式のモデルで高速かつメモリ効率がよいモジュールのようです。

さらに、v2が最近出ました。

https://github.com/turboderp/exllamav2

> This is a very initial release of ExLlamaV2, an inference library for running local LLMs on modern consumer GPUs.
It still needs a lot of testing and tuning, and a few key features are not yet implemented. Don't be surprised if things are a bit broken to start with, as almost all of this code is completely new and only tested on a few setups so far.
>
> 邦訳：
> これは、最新のコンシューマー向けGPU上でローカルLLMを実行するための推論ライブラリ、ExLlamaV2のごく初期のリリースである。
まだ多くのテストとチューニングが必要で、いくつかの主要な機能はまだ実装されていません。このコードのほとんどすべてが完全に新しいものであり、今のところいくつかのセットアップでしかテストしていないので、最初は少し壊れていても驚かないでください。

まだ早期リリースのため、不安定そうなのですが、サイトの情報によると、v1よりも高速で新しい量子化フォーマットに対応しているようです。

というわけで、実力を測るためにもv2を使ってDatabrick on AWS上での推論を試し見てみます。

DBRは13.3ML(GPU)、クラスタタイプはg4dn.xlarge(GPU:Nvidia T4)です。

# Step1. インストール

リリースの中から、Python 3.10系とCUDA 1.17系のWheelである```exllamav2-0.0.4+cu117-cp310-cp310-linux_x86_64.whl```を取得します。
取得するファイルは環境によって異なるため、DBR13.3ではないランタイムを使う場合は適切なものを選ぶように注意してください。

https://github.com/turboderp/exllamav2/releases/tag/v0.0.4


```python
!wget -P /tmp/exllamav2/ https://github.com/turboderp/exllamav2/releases/download/v0.0.4/exllamav2-0.0.4+cu117-cp310-cp310-linux_x86_64.whl
```

ダウンロードしたWheelをpipでインストール。

```python
# DBR13.3の場合、以下のバージョンをインストール
# PythonやCUDAのバージョンでインストールするwhlが異なるので注意
# 5～6分ほどインストールにかかりました。

%pip install /tmp/exllamav2/exllamav2-0.0.4+cu117-cp310-cp310-linux_x86_64.whl
%pip install -U -qq transformers accelerate

dbutils.library.restartPython()
```

ちなみに、Wheelを持ってこなくても
```python
%pip install exllamav2
```
でインストールは可能です。あまり差がわからなかったので、Wheel使わなくてもいいような気もします。（よくわかっていない）

# Step2. モデルダウンロード

せっかくなので、```mistralai/Mistral-7B-Instruct-v0.1```の新量子化フォーマットのモデルを以下からダウンロードして使います。今回は6.0bitモデルを使いました。
※　少数刻みで量子化指定されているのが面白い

https://huggingface.co/turboderp/Mistral-7B-instruct-exl2

```python
import os
from huggingface_hub import snapshot_download

model = "turboderp/Mistral-7B-instruct-exl2"
local_dir = f"/tmp/{model}"

snapshot_location = snapshot_download(
    repo_id=model,
    revision="6.0bpw",
    local_dir=local_dir,
    local_dir_use_symlinks=False,
)
```

# Step3. 推論

examplesの中にあるinference.pyに推論サンプルがありましたので、こちらを写経。

https://github.com/turboderp/exllamav2/blob/master/examples/inference.py

まずはTokenizerやGeneratorをロード。

```python
from exllamav2 import(
    ExLlamaV2,
    ExLlamaV2Config,
    ExLlamaV2Cache,
    ExLlamaV2Tokenizer,
)

from exllamav2.generator import (
    ExLlamaV2BaseGenerator,
    ExLlamaV2Sampler
)

import time

# Initialize model and cache

# ダウンロードしたモデルのディレクトリを指定
model_directory =  local_dir

config = ExLlamaV2Config()
config.model_dir = model_directory
config.prepare()

model = ExLlamaV2(config)
print("Loading model: " + model_directory)

# 複数GPUを使う場合、それぞれへのメモリ振り分けを設定。
# シングルGPUの場合は引数指定不要
# model.load([18, 24])
model.load()

tokenizer = ExLlamaV2Tokenizer(config)

cache = ExLlamaV2Cache(model)

# Initialize generator

generator = ExLlamaV2BaseGenerator(model, cache, tokenizer)

# Generate some text

settings = ExLlamaV2Sampler.Settings()
settings.temperature = 0.85
settings.top_k = 50
settings.top_p = 0.8
settings.token_repetition_penalty = 1.15
settings.disallow_tokens(tokenizer, [tokenizer.eos_token_id])
```

推論してみます。```mistralai/Mistral-7B-Instruct-v0.1```の公式サンプルから。
マヨネーズのレシピを聞いています。

```python
prompt = ("

 What is your favourite condiment? 
"
"Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!
 "
"
 Do you have mayonnaise recipes? 
")

max_new_tokens = 150

generator.warmup()
time_begin = time.time()

output = generator.generate_simple(prompt, settings, max_new_tokens, seed = 1234)

time_end = time.time()
time_total = time_end - time_begin

print(output)
print()
print(f"Response generated in {time_total:.2f} seconds, {max_new_tokens} tokens, {max_new_tokens / time_total:.2f} tokens/second")
```

```:結果


 What is your favourite condiment? 
Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!
 
 Do you have mayonnaise recipes? 
 Oh, absolutely! Here's a simple recipe for homemade mayonnaise that I think you'll enjoy:

Ingredients:
- 2 large egg yolks
- 1 tablespoon Dijon mustard
- 3 tablespoons white wine vinegar or lemon juice
- 1 cup vegetable oil (such as canola or sunflower)
- Salt and pepper to taste

Instructions:
1. In a medium bowl, whisk together the egg yolks, mustard, and vinegar until well combined.
2. Slowly drizzle in the vegetable oil while whisking constantly until the mixture thickens and doubles in volume. This

Response generated in 4.55 seconds, 150 tokens, 32.97 tokens/second
```

内容的にもあってそうですし、秒間33トークンは7Bモデルとしても十分な速さですね。

出力トークン数を増やして日本語でも聞いてみます。

```python
prompt = "

人工知能とは何？日本語で答えて。
"

max_new_tokens = 512

generator.warmup()
time_begin = time.time()

output = generator.generate_simple(prompt, settings, max_new_tokens, seed = 1234)

time_end = time.time()
time_total = time_end - time_begin

print(output)
print()
print(f"Response generated in {time_total:.2f} seconds, {max_new_tokens} tokens, {max_new_tokens / time_total:.2f} tokens/second")
```

```:結果


人工知能とは何？日本語で答えて。
 2.14

人工知能(AI)とは、人間の識別力や判断力を模擬する計算機プログラムの集合体。人工知能は、自然界においても存在し、人類によって作成されたものであり、人生の中でも常に現わることが多い。例えば、会話の理解、画像認識など。人工知能は、大きく分けられるが、最小単位は「ルール」と言う概念がある。

人工知能は、一般的に、人間の思考力や行動力を模擬するプログラムである。ただし、その意味は十分複雜である。例えば、人工知能を定義する方法については、いろいろあるぞ。

今回は、人工知能を「ルール」に基づいた解決策を示すことに焦点を当てます。

2.2 人工知能の問題
--------------

人工知能は、複雜な問題を解決するために使用されるコンピュータプログラムです。人工知能の問題は、以下の二種に分かれます。

### 2.2.1 簡單な問題

ここでは、直接のパテーションやルールベースのアクセス制限の問題が指定された場合に必要な解決策を解説します。

#### 2.2.1.1 パテーションベースのアクセス制限

パテーションベースのアクセス制限は、通常、特定のユーザーがアクセスできるデータのみを取得

Response generated in 15.16 seconds, 512 tokens, 33.78 tokens/second
```

内容はおかしいのですが、こちらも33トークン/秒で出力されます。

なお、VRAM使用量は7GB強でした。

# おまけ

従来のGPTQモデルも利用できるので、[TheBloke兄貴が公開しているモデル](https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.1-GPTQ)でもやってみました。
こちらは4bit量子化モデルです。違いはダウンロード部だけなので、他のコードは省略。

```python
import os
from huggingface_hub import snapshot_download

model = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ"
local_dir = f"/tmp/{model}"

snapshot_location = snapshot_download(
    repo_id=model,
    revision="main",
    local_dir=local_dir,
    local_dir_use_symlinks=False,
)
```

マヨネーズのレシピ：

```:結果


 What is your favourite condiment? 
Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!
 
 Do you have mayonnaise recipes? 
 Oh, absolutely! Here's a simple recipe for classic mayonnaise that you can easily whip up at home:

Ingredients:
- 1 cup vegetable oil
- 2 large egg yolks
- 1 tablespoon Dijon mustard
- 2 tablespoons white wine vinegar or lemon juice
- Salt and pepper to taste

Instructions:
1. In a food processor or blender, combine the vegetable oil, egg yolks, Dijon mustard, white wine vinegar or lemon juice, salt, and pepper.
2. Blend until the mixture is smooth and creamy. You may need to stop and scrape down

Response generated in 3.14 seconds, 150 tokens, 47.81 tokens/second
```

6.0bit -> 4bitの結果、推論速度は早くなってますね。返答内容は同じでした。

人工知能とは？：

```:結果


人工知能とは何？日本語で答えて。
 20世纪70年代以来，人工智能（AI）领域的研究已经进行了数十年，在这段时间内，人工智能领域的学术发展非常快，但是由于研究人工智能的具体定义和其与人类思维相关的问题仍然存在一些疑慮，因此人工智能领域的研究人员还没有达到一致的定义。

根据一些研究人员的观点，人工智能可以被认为是人类所具备的智能能力的模拟或创造出来的技术。它通过使用计算机程序、算法等技巧方式，实现对人类的思维过程和解决问题的模拟。例如，人工智能可以通过学习人类语言结构或图像识别技巧等方式，模拟人类的自然语言处理和视觉识别能力。

而根据另外一个观点，人工智能是指人类所做的任何智能系统或产品，包括人类模拟的自动化设备、机器翻译软件、音乐生成软件等。这种看法更加普及，并且也更容易应对人工智能领域的各种不同研究方向和目标。

总之，人工智能是一个非常广泛的概念，可以从很多方面考虑，最基本地就是人类所做的自动化系统或产品，它们通过模拟人类智能能力实现自动化或智能化的功能。人工智能领域的研究人员继续探索人工智能的定义和技术实现方案，希望将人工智能技术应用到各种领域中，提高人类生活质量和效率。

人工智能（Artificial Intelligence）是人类所做的自动化系统或產

Response generated in 10.33 seconds, 512 tokens, 49.55 tokens/second
```

こちらは日本語で返してくれませんでした。が、推論自体はできています。
全てのGPTQモデルを利用できるかはわからないのですが、使えるモデルは多そうです。

# まとめ

インストールには多少時間がかかりますが、そこに目を瞑れば気軽に使えそうでした。
2.0bitという、もやはそれどうなってるのというレベルで量子化できるようですし、マルチGPUでの推論も対応しているようなので、Databricksで大型のモデルを利用する場合は良さそうです。

~~```mistralai/Mistral-7B-Instruct-v0.1```がCTranslate2で変換・利用可能かどうかは試してないのですが、使える場合、推論速度比較とかもしてみたいと思います。~~
```mistralai/Mistral-7B-Instruct-v0.1```はCTranslate2で現状対応していないようなので、速度比較は別モデルでの実施がよさそうです。