N番煎じでMistral AIのMistral-Large2(量子化版)をDatabricks上で試す

導入

先日、Mistral AI社が大規模言語モデルMistral Large2を公開しました。

123Bパラメータというかなり大型のモデルです。
日本語含めた多くの言語に対応しており、コード生成も多種のプログラミング言語をサポートしています。

これまでのモデルと異なり、研究用および非商用利用のための使用と変更を許可するMistral Research Licenseの下でリリースされています。商用利用の場合は連絡を取って許可を得る必要があるようです。

モデルは以下のHuggingfaceリポジトリで公開されています。

というわけで、いつものように軽く試してみます。

検証はDatabricks on AWS上で実施しました。
DBRは15.4ML(Beta)、クラスタタイプはg5.12xlargeです。
推論エンジンにはSGLangを利用します。

検証には量子化モデルを利用します。
そのため、本来のモデルとは出力結果が異なることに注意ください。

Step1. パッケージインストール

SGLangと関係するパッケージをインストール。

# transformers 最新化
%pip install -U transformers

# vLLM
%pip install vllm

# flashinfer
%pip install https://github.com/flashinfer-ai/flashinfer/releases/download/v0.1.1/flashinfer-0.1.1+cu121torch2.3-cp311-cp311-linux_x86_64.whl

# sglang
%pip install "sglang[srt]"

dbutils.library.restartPython()

Step2. ランタイムの実行

HuggingfaceからGPTQで4bit量子化したモデルダウンロードしておき、保管したモデルファイルをロードします。

VRAMの関係上なのか、モデルのロードにはdisable_cuda_graph指定が必要でした。

import sglang as sgl
from sglang.lang.chat_template import get_chat_template

model_path = "/Volumes/training/llm/model_snapshots/models--ModelCloud--Mistral-Large-Instruct-2407-gptq-4bit/"

runtime = sgl.Runtime(
    model_path=model_path,
    tp_size=4,
    mem_fraction_static=0.9,
    disable_cuda_graph=True,
)
runtime.endpoint.chat_template = get_chat_template("llama-2-chat")

sgl.set_default_backend(runtime)

Step3. バッチ推論

ここと同じプロンプト内容で推論を実行してみます。

@sgl.function
def single_turn_question(s, sp, question):
    s += sgl.system(sp)
    s += sgl.user(question)
    s += sgl.assistant(sgl.gen("answer", max_tokens=512))

def batch():
    prompts = [
        "Hello, what is your name?",
        "Databricksとは何ですか？詳細に教えてください。",
        "まどか☆マギカでは誰が一番かわいい?",
        "ランダムな10個の要素からなるリストを作成してソートするコードをPythonで書いてください。",
        "現在の日本の首相は誰？",
        "あなたはマラソンをしています。今3位の人を抜きました。あなたの今の順位は何位ですか?ステップバイステップで考えてください。",
    ]

    system_prompt = """あなたは親切なAIアシスタントです。日本語で回答してください。"""

    prompts_dict = [{"sp": system_prompt, "question": p} for p in prompts]

    states = single_turn_question.run_batch(
        prompts_dict,
        temperature=0.0,
        top_k=5,
        top_p=0.5,
    )

    for s in states:
        print(" ------------------ ")
        for m in s.messages():
            if m["role"] != "system":
                print(m["role"], ":", m["content"])
        print()

batch()

出力

 ------------------ 
user : Hello, what is your name?
assistant : こんにちは！私の名前はAIアシスタントです。何かお手伝いできることがあれば教えてください。

 ------------------ 
user : Databricksとは何ですか？詳細に教えてください。
assistant : Databricksは、データエンジニアリング、データサイエンス、機械学習のための統合プラットフォームです。このプラットフォームは、Apache Sparkの創設者によって設立され、Sparkの機能を最大限に活用するために設計されています。Databricksは、クラウドベースの環境でデータ処理、分析、機械学習のワークフローを効率的に管理するためのツールを提供します。

以下に、Databricksの主要な機能と特徴を詳しく説明します：

1. **統合環境**:
   - Databricksは、データエンジニアリング、データサイエンス、機械学習のための統合環境を提供します。これにより、異なる役割のチームが一つのプラットフォームで協力して作業を進めることができます。

2. **Apache Sparkの最適化**:
   - Databricksは、Apache Sparkのパフォーマンスを最適化するための独自の技術を採用しています。これにより、データ処理の速度と効率が向上します。

3. **クラウドネイティブ**:
   - Databricksは、AWS、Azure、Google Cloudなどの主要なクラウドプロバイダーと統合されており、クラウド上でのデータ処理と分析が容易に行えます。

4. **ノートブックインターフ

 ------------------ 
user : まどか☆マギカでは誰が一番かわいい?
assistant : 「魔法少女まどか☆マギカ」に登場するキャラクターはそれぞれ魅力的で、ファンの好みによって「一番かわいい」と感じるキャラクターが異なることが多いです。例えば、主人公の鹿目まどかはその優しさと勇気で多くのファンに愛されています。他にも、暁美ほむらのクールで知的な魅力や、美樹さやかの元気で明るい性格、巴マミの大人びた雰囲気、佐倉杏子の強さと芯のある性格など、それぞれのキャラクターに独自の魅力があります。

どのキャラクターが一番かわいいかは、個々のファンの好みや感じ方によるので、皆さんそれぞれのお気に入りがあると思います。あなたがどのキャラクターが一番好きか、ぜひ教えてください！

 ------------------ 
user : ランダムな10個の要素からなるリストを作成してソートするコードをPythonで書いてください。
assistant : もちろんです。以下は、ランダムな10個の要素からなるリストを作成し、それをソートするPythonコードです。

```python
import random

# ランダムな10個の要素からなるリストを作成
random_list = [random.randint(1, 100) for _ in range(10)]

# リストをソート
sorted_list = sorted(random_list)

# 結果を表示
print("ランダムなリスト:", random_list)
print("ソートされたリスト:", sorted_list)
```

このコードは、まず`random`モジュールを使用してランダムな10個の整数（1から100の範囲）を生成し、それをリストに格納します。その後、`sorted`関数を使用してリストをソートし、結果を表示します。

 ------------------ 
user : 現在の日本の首相は誰？
assistant : 現在の日本の首相は岸田文雄です。彼は2021年10月に首相に就任しました。

 ------------------ 
user : あなたはマラソンをしています。今3位の人を抜きました。あなたの今の順位は何位ですか?ステップバイステップで考えてください。
assistant : もちろんです。ステップバイステップで考えてみましょう。

1. あなたはマラソンをしています。
2. あなたは今3位の人を抜きました。
3. 3位の人を抜いたということは、あなたがその人よりも前に出たということです。
4. したがって、あなたの順位は3位になります。

ですので、あなたの今の順位は3位です。

実行時間は30秒弱でした。
全体的にそつのない回答です。日本語知識もかなり保持しているようです。
パラメータサイズに見合った性能のように思います。

まとめ

リリースされた時期から何かとLlama 3.1と比較されやすいですが、非常に高い性能を持っているモデルだと思います。
惜しいのはライセンス体系でしょうか。
これがApache 2.0での公開だと、かなり話題になったと思うのですが。