SGLangとDatabricksを使ってLlama-3.1を高スループットで推論する

# 導入

LMSYSから[SGLang](https://github.com/sgl-project/sglang)のバージョン0.2.x系がリリースされました。
公式Blogでは、[TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM)や[vLLM](https://github.com/vllm-project/vllm)と匹敵/超えるスループットを実現したと主張しています。

https://lmsys.org/blog/2024-07-25-sglang-llama3/

TensorRT-LLMはDatabricksの[プロビジョニングされたスループットの基盤モデルAPI](https://docs.databricks.com/ja/machine-learning/foundation-models/deploy-prov-throughput-foundation-model-apis.html)でも利用されている(はずの)パッケージであり、同水準のスループットが出るというのはすごい。

SGLangは以前[こちら](https://qiita.com/isanakamishiro2/items/aeb73d09bd81349fa023)でも試用していたのですが、非常に使い勝手のよいLLM/VLM推論のパッケージです。

Databricks東京リージョンでは、2024/7月末現在まだLlama-3.1のプロビジョニングされたスループットの基盤モデルAPIは対応していないようなので、今回は久しぶりにSGlangを使ってLlama3.1の推論を高効率に行ってみます。

実施環境はDatabricks on AWS、DBRは15.4ML LTS(Beta)、インスタンスタイプはg5.xlarge(GPU)です。

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

SGLangに必要なパッケージをインストール。
vllmやflashinferが必須となるパッケージなので合わせてインストールします。
(この時点でのSGLang最新バージョンは```0.2.5```でした)

```python
# 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からLlama3.1 8B Instructをダウンロードしておき、保管したモデルファイルをロードします。

https://huggingface.co/meta-llama/Meta-Llama-3.1-8B-Instruct

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

# Llama-3.1 8B Instructのモデルを保管したパス
model_path = "/Volumes/training/llm/model_snapshots/models--meta-llama--Meta-Llama-3-8B-Instruct/"

runtime = sgl.Runtime(
    model_path=model_path,
    mem_fraction_static=0.80,
)
# Chat用テンプレートの適用
runtime.endpoint.chat_template = get_chat_template("llama-3-instruct")

sgl.set_default_backend(runtime)
```

# Step3. バッチ推論

[ここ](https://qiita.com/isanakamishiro2/items/2169219ae537dab79a34)と同じプロンプト内容で推論を実行してみます。

```python
@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=1,
        top_p=0.6,
    )

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

batch()
```

````md:出力
 ------------------ 
user : Hello, what is your name?
assistant : Konnichiwa! My name is JARVIS, and I'm a friendly AI assistant. I'm here to help you with any questions or tasks you may have. How can I assist you today? 😊

 ------------------ 
user : Databricksとは何ですか？詳細に教えてください。
assistant : Databricksは、データ分析と機械学習のためのクラウドベースのプラットフォームです。Databricksは、Apache Sparkを基盤としており、データ分析、機械学習、データサイエンスのためのワークフローを提供しています。

Databricksは、以下のような機能を提供しています。

1. **データ分析**: Databricksは、Apache Sparkを基盤としており、データ分析のためのワークフローを提供しています。ユーザーは、データを読み込み、分析、可視化することができます。
2. **機械学習**: Databricksは、機械学習のためのワークフローを提供しています。ユーザーは、機械学習のためのコードを書き、実行することができます。
3. **データサイエンス**: Databricksは、データサイエンスのためのワークフローを提供しています。ユーザーは、データを読み込み、分析、可視化することができます。
4. **Collaboration**: Databricksは、チームでの協力的な開発を可能にします。ユーザーは、チームメンバーと共にワークフローを共有、協力的に開発することができます。
5. **Security**: Databricksは、セキュリティーのための機能を提供しています。ユーザーは、データをセキュアに保存、保護することができます。

Databricksは、以下のような利点があります。

1. **スケーラビリティ**: Databricksは、スケーラビリティーの高いプラットフォームです。ユーザーは、大量のデータを分析、機械学習することができます。
2. **Flexibility**: Databricksは、柔軟なプラットフォームです。ユーザーは、異なる言語、フレームワークを使用して開発することができます。
3. **Cost-effective**: Databricksは、コスト効果的なプラットフォームです。ユーザーは、必要なリソースを使用して開発することができます。

Databricksは、以下のような用途があります。

1. **データ分析**: Databricksは、データ分析のためのプラットフォームです。ユーザーは、データを読み込み、分析、可視化することができます

 ------------------ 
user : まどか☆マギカでは誰が一番かわいい?
assistant : まどか☆マギカは、かわいいキャラクターが多いですが、誰が一番かわいいかは、個人の好みによって異なります。

しかし、人気のあるキャラクターとして、モモエは挙げられます。彼女の可愛らしい容姿と、可愛らしい性格が、ファンの人気を高めているのです。

また、カナメも人気のあるキャラクターです。彼女の可愛らしい容姿と、可愛らしい性格が、ファンの人気を高めているのです。

以上、まどか☆マギカのキャラクターの中で、かわいいと感じるキャラクターは、モモエとカナメが挙げられます。

 ------------------ 
user : ランダムな10個の要素からなるリストを作成してソートするコードをPythonで書いてください。
assistant : Pythonでランダムな10個の要素からなるリストを作成してソートするコードは以下のようになります。
```python
import random

# ランダムな10個の要素を作成
elements = [random.randint(1, 100) for _ in range(10)]

# ソート
elements.sort()

print(elements)
```
このコードでは、`random` モジュールを使用してランダムな整数を10個生成し、リスト`elements`に格納します。次に、`sort()` メソッドを使用してリストをソートします。最後に、ソートされたリストを出力します。

実行すると、ランダムな10個の要素からなるリストが生成され、ソートされた状態で出力されます。例えば、以下のような出力結果が得られます。
```
[3, 5, 7, 11, 14, 17, 23, 29, 31, 42]
```
このリストはランダムに生成された10個の整数であり、ソートされた状態で出力されています。

 ------------------ 
user : 現在の日本の首相は誰？
assistant : 日本の現在の首相は、岸田文雄（Kishida Fumio）です。2021年10月4日から首相を務めています。

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

現在、3位の人を抜きました。ということは、3位の人より前にいますね！

前回の順位は、4位でした。3位の人を抜きましたということは、4位の人より前にいますね！

4位の人より前にいるということは、3位の人より前にいるということになります！

つまり、現在の順位は2位です！🏃‍♂️
````

実行時間は18秒くらいでした。
ちなみに別途4bit量子化したGPTQフォーマットでも試しましたが、そちらは7秒弱でした。
生成トークン数が異なるので単純比較になりませんが、速い。
自分でベンチを回しているわけではありませんが、確かにvLLMよりも高いスループットを達成しているかもしれません。

# まとめ

Llama 3.1自体は日本語非対応なのですが、Gemmaと同じ程度に自然な日本語となっているように見えます。（内容の正しさは別として）

SGLang自体は結構久しぶりに使ったのですが、対応モデルも増えていて使い勝手も良さそうです。もう少し使い倒してみようと思います。