高速なテキスト生成用推論サーバ「text-generation-inference」を試す

ここで知った
https://x.com/rohanpaul_ai/status/1866474376738320405
Embedding用は以前に試してたけど、テキスト生成用もあるのを知らなかった。
https://zenn.dev/kun432/scraps/7f04e5f8e4b3cb
GitHubレポジトリ
https://github.com/huggingface/text-generation-inference

 Text Generation InferenceRust、Python、gRPCサーバーを用いたテキスト生成推論ツール。Hugging Faceで実際に使用されており、Hugging Chat、Inference API、Inference Endpointの動力源となっています。
**Text Generation Inference（TGI）**は、大規模言語モデル（LLM）をデプロイおよび提供するためのツールキットです。TGIは、Llama、Falcon、StarCoder、BLOOM、GPT-NeoX、その他 の最も人気のあるオープンソースのLLM向けに、高性能なテキスト生成を可能にします。TGIには以下の機能が含まれています：
人気のあるLLMを簡単に提供できるランチャー
実運用対応（Open Telemetryを用いた分散トレーシング、Prometheusによるメトリクス収集）
マルチGPUでの推論高速化のためのテンソル並列処理
Server-Sent Events（SSE）を利用したトークンストリーミング
リクエストの継続的バッチ処理による総スループット向上
OpenAI Chat Completion API互換のメッセージAPI

Flash AttentionやPaged Attentionを活用した最適化されたTransformers推論コード
量子化サポート：
bitsandbytes
GPT-Q
EETQ
AWQ
Marlin
fp8


Safetensorsを用いたモデルウェイトの読み込み

大規模言語モデル用ウォーターマーク生成
ロジットワーパー（温度スケーリング、top-p、top-k、繰り返しペナルティなど、詳細はtransformers.LogitsProcessor参照）
ストップシーケンス
ログ確率の取得

スペキュレーション（約2倍のレイテンシ改善）

Guidance/JSON出力：推論の高速化と仕様に準拠した出力の保証
カスタムプロンプト生成：モデルの出力を誘導するカスタムプロンプトの提供が可能
微調整対応：特定のタスクに最適化されたモデルで高精度・高性能を実現

 ハードウェアサポート：Nvidia

AMD（-rocm）
Inferentia
Intel GPU
Gaudi
Google TPU

kun432

ドキュメントはこちら

Quick Tourに従って進めてみる

前提

Ubuntu 22.04
RTX4090(VRAM 24GB)
NVIDIA Container Toolkitはインストール済み
CUDA-12.6

作業ディレクトリ作成

mkdir tgi-work && cd tgi-work

ではモデルを指定して起動。使用可能なモデルはここにあるが、書いていないものでも動かせるものはある模様。今回はサポートされている"google/gemma-2-2b-it"を使う。　なお、モデル利用に制限がある・プライベートモデルの場合にはHuggingFaceのトークンが必要になる（ここ）

model="google/gemma-2-2b-it"
volume=$PWD/data
token=XXXXXXXXX

docker run --gpus all \
    --shm-size 1g \
    -p 8080:80 \
    -v $volume:/data \
    -e HF_TOKEN=$token \
    ghcr.io/huggingface/text-generation-inference:3.0.1 \
    --model-id $model

起動後にモデルがダウンロードされる。以下のような出力がされればOKっぽい。

出力

2024-12-12T06:45:34.414706Z  INFO text_generation_router::server: router/src/server.rs:2440: Serving revision 299a8560bedf22ed1c72a8a11e7dce4a7f9f51f8 of model google/gemma-2-2b-it
2024-12-12T06:45:42.405963Z  INFO text_generation_router::server: router/src/server.rs:1873: Using config Some(Gemma2)
(snip)
2024-12-12T06:45:42.492242Z  INFO text_generation_router::server: router/src/server.rs:2402: Connected

使い方はこちら

APIエンドポイントは、generateエンドポイントを使う。

curl 127.0.0.1:8080/generate \
    -X POST \
    -d '{"inputs":"競馬の楽しみ方を5つ、簡潔にリストアップして。","parameters":{"max_new_tokens":4096}}' \
    -H 'Content-Type: application/json' | jq -r .

出力

{
  "generated_text": "\n\n1. **予想の喜び:**  馬券の組み合わせを考え、予想する喜び。\n2. **レースの展開:**  馬の動きや競争状況を分析し、予想を修正する楽しみ。\n3. **馬の個性:**  それぞれの馬の性格や得意なコースなどを理解し、魅力を感じること。\n4. **競馬場の雰囲気:**  興奮する観客、緊張する騎手、そして馬の力強い走り。\n5. **歴史と伝統:**  長い歴史の中で培われた伝統的なレースや、その背景を学ぶこと。 \n\n\n"
}

またはOpenAI Chat Completion APIと互換性がある以下のエンドポイントでもOK。

curl "http://127.0.0.1:8080/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -d '{
        "messages": [
            {
                "role": "user",
                "content": "競馬の楽しみ方を5つ、簡潔にリストアップして。"
            }
        ]
    }' | jq -r .

出力

{
  "object": "chat.completion",
  "id": "",
  "created": 1733986290,
  "model": "google/gemma-2-2b-it",
  "system_fingerprint": "3.0.1-sha-bb9095a",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": "競馬の楽しみ方は人それぞれ！ですが、以下に 5 つの一般的な楽しみ方を提供します。\n\n1. **予想を楽しむ:** オメーに、馬券の組み合わせを立てたり、馬の得意条件や競走条件を分析して予想してみる。成功すれば気持ちよいですが、予想は楽しむことが重要です。\n2. **競走の様子を楽しもう:** 馬の力強さや競走中の戦略、 jockey の技術性など、レース全体の緊張感や興奮感を味わう。\n3. **美しいコースを眺めながら:**eventの雰囲気を満喫したり、馬場状態を分析したり、競馬場での環境を楽しんだり。\n4. **馬が走る過程を体感:** 馬と jockey の連携する過程、競走の瞬間に集中し、ならではの感動を味わう。 \n5. **交流を楽しむ:** 競馬場では人と人との交流も楽しむことができます。 \n\n競馬の魅力はそれ以外にたくさんあります。 ぜひあなた自身の楽しみ方を見つけてください! 🐎 \n"
      },
      "logprobs": null,
      "finish_reason": "stop"
    }
  ],
  "usage": {
    "prompt_tokens": 23,
    "completion_tokens": 215,
    "total_tokens": 238
  }
}

Pythonで。huggingface_hubを使う場合。

pip install huggingface_hub

sample_hfhub.py

from huggingface_hub import InferenceClient

client = InferenceClient(
    base_url="http://localhost:8080/v1/",
)

output = client.chat.completions.create(
    model="tgi",
    messages=[
        {"role": "user", "content": "競馬の楽しみ方を5つ、簡潔にリストアップして。"},
    ],
    stream=True,
    max_tokens=4096,
)

for chunk in output:
    print(chunk.choices[0].delta.content, end="", flush=True)

python sample_hfhub.py

出力

競馬の楽しみ方5選：

1. **馬の力強さと競争力を楽しむ:**  レースの展開を予測し、動作や体格などを分析して、勝敗を左右する能力や戦略性を楽しむ。
2. **予想を立てて楽しむ:**  過去のデータと分析結果に基づき、予想を立て、その結果ワクワクを味わう。
3. **馬場状態や天候の影響を意識する:**  馬場状態や天候がレース結果に影響を与えることを理解し、その変化を意識して楽しむ。
4. **地域のファンと交流する:**  競馬場での人と人との交流を通して、地域の文化や伝統を感じて楽しむ。
5. **それぞれの馬の個性や能力を楽しむ:**  違う馬タイプが競い合う中で、それぞれの個性や能力を理解し、その魅力を感じて楽しむ。

OpenAI SDKでもそのまま使える。

pip install openai

sample_openai.py

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8080/v1/",
    api_key="-"
)

stream = client.chat.completions.create(
    model="tgi",
    messages=[
        {"role": "user", "content": "競馬の楽しみ方を5つ、簡潔にリストアップして"}
    ],
    stream=True
)

for chunk in stream:
    if chunk.choices[0].delta.content is not None:
        print(chunk.choices[0].delta.content, end="", flush=True)

python sample_openai.py

出力

## 競馬の楽しみ方５選

1. **レースの結果にドキドキ！** 予想通りにレースが展開される時、興奮が絶頂！勝利する馬と、予想外の結果が待ち受けるドキドキ感。
2. **馬を深く知る！** 各馬の血統、 jockey、戦略、過去のレースなど、馬の歴史と現在の姿を知ることで、より楽しみ方が広がる。
3. **競馬場は活気あふれる場所！** スリル満点のレース観戦だけでなく、競馬場の雰囲気、人々の熱気も味わえる。
4. **仲間との会話を楽しむ！**  予想に合っているときは、お互いの意見交換、喜びを分かち合う。 逆に、予想が外れたときも、笑い合える。
5. **競馬は新しい発見！** 運だめしも、競馬場での出会い、新しい出会う機会を増やす。

ドキュメントにはGradioの例も載っている。

kun432

量子化

量子化は--quantizeオプションを指定すればいいみたい。対応している量子化方式は以下。

bits-and-bytes（8ビット:bitsandbytes、4ビット: bitsandbytes-fp4 / bitsandbytes-nf4）
GPT-Q (gptq)
AWQ (awq)
Marlin (marlin)
EETQ (eetq)
EXL2 (exl2)
fp8 (fp8)

bitsandbytes-nf4の場合

docker run --gpus all \
    --shm-size 1g \
    -p 8080:80 \
    -v $volume:/data \
    -e HF_TOKEN=$token \
    ghcr.io/huggingface/text-generation-inference:3.0.1 \
    --model-id "google/gemma-2-2b-it" \
    --quantize bitsandbytes-nf4   # 量子化オプションでbitsandbytes-nf4を指定

ただ、VRAM使用量見てる限り、オプションつけても違いは見られなかった。

GPTQやAWQの場合にはそれ用のモデル（こことかここ）を指定する必要がある。例えばGPTQで"Qwen/Qwen2.5-7B-Instruct-GPTQ-Int8"だとこんな感じ。

docker run --gpus all \
    --shm-size 1g \
    -p 8080:80 \
    -v $volume:/data \
    -e HF_TOKEN=$token \
    ghcr.io/huggingface/text-generation-inference:3.0.1 \
    --model-id "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int8" \  # GPTQなモデルを指定
    --quantize gptq  # 量子化オプションでGPTQを指定

環境やモデルに依存する部分もあるのかな？自分の環境だとAWQは動かなかったりした。

AWQ（"LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct-AWQ"）の場合

docker run --gpus all \
    --shm-size 1g \
    -p 8080:80 \
    -v $volume:/data \
    -e HF_TOKEN=$token \
    ghcr.io/huggingface/text-generation-inference:3.0.1 \
    --model-id "LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct-AWQ" \
    --quantize awq

出力

NotImplementedError: awq quantization is not supported for AutoModel rank=0
2024-12-12T08:12:13.856321Z ERROR text_generation_launcher: Shard 0 failed to start
2024-12-12T08:12:13.856365Z  INFO text_generation_launcher: Shutting down shards
Error: ShardCannotStart

とりあえず量子化については、自分の知見が足りないのもあってちょっとわかっていない。

あとドキュメントにはRoPE Scalingについても記載がある。

kun432

構造化出力・ツール

最初見たときGuidanceってMSのやつかなと思ったけど、関係ないみたい。

で、TGIは以下をサポートしている

JSONと正規表現の文法を使用した構造化出力
Tool Calls (Function Calling)

TGI独自の/generateエンドポイント、そしてOpenAI Chat Completion互換のAPIエンドポイントの両方で対応している

"Qwen/Qwen2.5-7B-Instruct"がツールに対応しているようなのでそれで試してみる。

docker run --gpus all \
    --shm-size 1g \
    -p 8080:80 \
    -v $volume:/data \
    -e HF_TOKEN=$token \
    ghcr.io/huggingface/text-generation-inference:3.0.1 \
    --model-id "Qwen/Qwen2.5-7B-Instruct"

/generateエンドポイントにcurlで。

curl localhost:8080/generate \
    -X POST \
    -H 'Content-Type: application/json' \
    -d '{
    "inputs": "自転車で公園を走っていたら、子犬と猫とアライグマを見かけました。",
    "parameters": {
        "repetition_penalty": 1.3,
        "grammar": {
            "type": "json",
            "value": {
                "properties": {
                    "場所": {
                        "type": "string"
                    },
                    "行動": {
                        "type": "string"
                    },
                    "視覚できる動物の数": {
                        "type": "integer",
                        "minimum": 1,
                        "maximum": 5
                    },
                    "動物": {
                        "type": "array",
                        "items": {
                            "type": "string"
                        }
                    }
                },
                "required": ["場所", "行動", "視覚できる動物の数", "動物"]
            }
        }
    }
}' | jq -r .

出力

{
  "generated_text": "{ \"場所\": \"公園\", \"行動\": \"走る\", \"視覚できる動物の数\":3, \"動物\": [\"子犬\",\"猫\",\"アライグマ\"] }"
}

一般的なTool Callsとは異なるスキーマになるけれども、同じような感覚で使えるのがわかる。

huggingface_hubを使ったPythonでの書き方。

sample_hfhub_grammer.py

from huggingface_hub import InferenceClient

client = InferenceClient(
    base_url="http://localhost:8080/",
)

schema = {
    "properties": {
        "場所": {"title": "場所", "type": "string"},
        "行動": {"title": "行動", "type": "string"},
        "視覚できる動物の数": {
            "maximum": 5,
            "minimum": 1,
            "title": "視覚できる動物の数",
            "type": "integer",
        },
        "動物": {"items": {"type": "string"}, "title": "動物", "type": "array"},
    },
    "required": ["場所", "行動", "視覚できる動物の数", "動物"],
    "title": "動物",
    "type": "object",
}

user_input = "自転車で公園を走っていたら、子犬と猫とアライグマを見かけました。"
resp = client.text_generation(
    f"JSONに変換してください: '{user_input}'. 次のスキーマを使用してください: {schema}",
    grammar={"type": "json", "value": schema},
)

print(resp)

python sample_hfhub_grammer.py | jq -r .

出力

{
  "場所": "公園",
  "行動": "走っていた",
  "視覚できる動物の数": 3,
  "動物": [
    "子犬",
    "猫",
    "アライグマ"
  ]
}

スキーマはPydanticモデルで定義することもできる。

sample_hfhub_grammer_pydantic.py

from huggingface_hub import InferenceClient
from pydantic import BaseModel, conint
from typing import List

client = InferenceClient(
    base_url="http://localhost:8080/",
)


class Animals(BaseModel):
    場所: str
    行動: str
    視覚できる動物の数: conint(ge=1, le=5)
    動物: List[str]


user_input = "自転車で公園を走っていたら、子犬と猫とアライグマを見かけました。"
resp = client.text_generation(
    f"JSONに変換してください: '{user_input}'. 次のスキーマを使用してください: {Animals.model_json_schema()}",
    grammar={"type": "json", "value": Animals.model_json_schema()},
)

print(resp)

python sample_hfhub_grammer_pydantic.py | jq -r .

出力

{
  "場所": "公園",
  "行動": "走っていた",
  "視覚できる動物の数": 3,
  "動物": [
    "子犬",
    "猫",
    "アライグマ"
  ]
}

あとちょっと独自な気がするけど、正規表現でスキーマを指定できる。

sample_hfhub_grammer_regex.py

from huggingface_hub import InferenceClient

client = InferenceClient(
    base_url="http://localhost:8080/",
)

section_regex = "(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)"
regexp = f"{section_regex}\.{section_regex}\.{section_regex}\.{section_regex}"

user_input = f"GoogleのパブリックDNSサーバのIPアドレスは何ですか？次の正規表現を使用してください: {regexp}"
resp = client.text_generation(
    user_input,
    grammar={"type": "regex", "value": regexp},
)

print(resp)

python sample_hfhub_grammer_regex.py

出力

1.1.1.123

ハルシネーションしてるのは御愛嬌。

OpenAI互換APIではTool Callにも対応している様子。

curl localhost:8080/v1/chat/completions \
    -X POST \
    -H 'Content-Type: application/json' \
    -d '{
    "model": "tgi",
    "messages": [
        {
            "role": "user",
            "content": "神戸の天気を教えて。"
        }
    ],
    "tools": [
        {
            "type": "function",
            "function": {
                "name": "get_current_weather",
                "description": "現在の天気を取得する",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "都市名や都道府県を指定。例: 京都府京都市、東京都、大阪市"
                        },
                        "format": {
                            "type": "string",
                            "enum": ["摂氏", "華氏"],
                            "description": "気温の単位。ユーザが指定した地域から推論すること。"
                        }
                    },
                    "required": ["location", "format"]
                }
            }
        }
    ],
    "tool_choice": "get_current_weather"
}'  | jq -r .

出力

{
  "object": "chat.completion",
  "id": "",
  "created": 1734129503,
  "model": "Qwen/Qwen2.5-7B-Instruct",
  "system_fingerprint": "3.0.1-sha-bb9095a",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "tool_calls": [
          {
            "id": "0",
            "type": "function",
            "function": {
              "description": null,
              "name": "get_current_weather",
              "arguments": {
                "format": "摂氏",
                "location": "神戸"
              }
            }
          }
        ]
      },
      "logprobs": null,
      "finish_reason": "stop"
    }
  ],
  "usage": {
    "prompt_tokens": 271,
    "completion_tokens": 28,
    "total_tokens": 299
  }
}

ただし、Pythonで関数の実行まで含めてやろうとすると、以前OpenAIで動いていたTool Callsのコードをそのまま実行しても動かなかった。ちょっと原因は追えてないけど、なんとなく返ってくるレスポンスがOpenAI SDKで期待している型と異なるように見える。

このへんかな。

kun432

VLMにも対応している

kun432

 まとめとりあえずざっと動かしてみた。あくまでも個人的な印象ではあるが、
text-embeddings-inferenceに比べると、やれることも多いし、モデルもたくさんあるので、ちょっと安定していないのかなぁという感じ
OpenAI互換APIなのであれば、Tool Callsは同じ挙動になってほしい。結構前からIssue挙がってるけど対応されていないのは微妙に感じる。
huggingface_hubのクライアントであればうまく動いたりするのかな？？？試してないけど。

というところ。Dockerイメージが用意されているのは、本番デプロイ等を考えてもいいと思うしお手軽なんだけどな。

このスクラップは4日前にクローズされました