Posted inAI

DockerでText Generation Inference（TGI）をデプロイして高性能LLMサービングを実現する

AI tutorial - IT technology blog
AI tutorial - IT technology blog

本番環境のボトルネック：シンプルな推論が限界を迎えるとき

個人実験でLlama 3のような大規模言語モデル（LLM）をデプロイするのは簡単です。Transformersライブラリでモデルを読み込み、FastAPIエンドポイントでラップすれば動作します。しかし、私のチームはこのセットアップがスケールしないことを痛い経験から学びました。半年前、RAGベースの社内ツールをリリースしましたが、たった50人の同時ユーザーの負荷であっけなく崩壊してしまったのです。

症状は深刻でした。レイテンシは30秒を超えてスパイクし、タイムアウトエラーが常態化しました。サーバーが1つのリクエストを処理するのに必死になっている間、ユーザーは白い画面を見つめ続け、やっと返ってくる応答は大量のテキストが一気に表示されるだけ。標準的なPythonベースの推論では、本番トラフィックをうまく捌けないことが明らかになりました。レスポンシブなユーザー体験に必要な効率性が根本的に欠けているのです。

原因の究明：標準サービングが詰まる理由

遅延を解消するには、LLMがデータを処理する仕組みを根本から見直す必要がありました。従来のサーバーは順次処理または静的バッチングを使用します。静的バッチでは、サーバーは一定数のリクエストが集まるまで待機し、それらをまとめて1つの行列演算として処理します。これは非効率です。あるユーザーが俳句を求め、別のユーザーが500語のエッセイを求めた場合、短いリクエストは長いリクエストに引きずられて待たされます。高価なGPUサイクルが無駄になるのです。

標準的なPython実装はGIL（グローバルインタープリタロック）の問題も抱えており、このオーバーヘッドがGPUのフル活用を妨げます。トークンストリーミングがなければ、モデルの出力が生成される途中でユーザーに見せることができず、体感的な待ち時間がさらに長く感じられます。私たちには継続的バッチング（Continuous Batching）が必要でした。この技術は、既存リクエストのトークンが生成された瞬間に新しいリクエストをバッチに挿入し、GPUを常にフル稼働させ続けます。

適切なスタックの選定：なぜTGIなのか

本番スタックを決定する前に、いくつかの専用エンジンを評価しました。

  • FastAPI + Transformers：構築は簡単ですが、PagedAttentionのような高同時接続の最適化が欠如しています。
  • vLLM：PagedAttentionで非常に高速かつ人気がありますが、当時は特定のニッチなモデルにおいてHugging Faceエコシステムとの統合感がやや薄く感じられました。
  • Text Generation Inference（TGI）：Hugging Faceが自社の本番APIのために専用に開発したもので、Rust、C++、Pythonで書かれた強力なエンジンです。

TGIはFlash Attention、PagedAttention、そしてAWQやbitsandbytesといった量子化手法をネイティブでサポートしています。TGIへ移行したことで、ハードウェアコストを40%削減し、総スループットを2倍に増加させました。以前は50ユーザーで限界を迎えた同じハードウェアで、200以上の同時リクエストを処理できるようになりました。

DockerでTGIをデプロイする：ステップバイステップガイド

TGIを安定して運用するにはDockerが最も確実な方法です。複雑なNVIDIAドライバ、CUDA 12.1カーネル、Rustの依存関係をひとつのポータブルなコンテナにまとめることで、「自分のマシンでは動く」問題を根絶できます。

1. ハードウェアの前提条件

十分なVRAMを持つNVIDIA GPUが必要です。Llama 3（8B）モデルの場合、RTX 4090やA10Gのように最低16GBのVRAMを目安にしてください。DockerがハードウェアにアクセスできるよU、NVIDIA Container Toolkitがインストールされていることを確認してください。

# DockerがGPUを認識できるか確認する
docker run --rm --gpus all nvidia/cuda:12.1.0-base-ubuntu22.04 nvidia-smi

2. TGIコンテナの起動

Hugging Faceの公式イメージを使ってLlama-3-8B-Instructを実行します。ゲートモデルを使用する場合は、Hugging Face HubのトークンをあらかじめU意しておいてください。

model="meta-llama/Meta-Llama-3-8B-Instruct"
volume=$PWD/data
token="your_hf_token_here"

docker run --gpus all --shm-size 1g -p 8080:80 \
    -v $volume:/data \
    -e HUGGING_FACE_HUB_TOKEN=$token \
    ghcr.io/huggingface/text-generation-inference:2.0 \
    --model-id $model \
    --max-batch-prefill-tokens 2048 \
    --max-total-tokens 4096

設定の解説：

  • --shm-size 1g：GPU間の高速通信のための共有メモリを確保します。
  • --max-total-tokens：入力と出力の合計長の上限を設定します。
  • --max-batch-prefill-tokens：最初のプロンプト処理フェーズで処理するトークン数を制限し、OOM（メモリ不足）エラーを防ぎます。

3. トークンストリーミングの有効化

TGIはServer-Sent Events（SSE）を使うときに真価を発揮します。UIがテキストを1文字ずつリアルタイムで表示できるようになります。以下はストリームを受信するPythonのコード例です。

import requests
import json

def stream_llm_response(prompt):
    url = "http://localhost:8080/generate_stream"
    data = {
        "inputs": prompt,
        "parameters": {"max_new_tokens": 500, "temperature": 0.7}
    }

    response = requests.post(url, json=data, stream=True)
    
    for line in response.iter_lines():
        if line:
            decoded = line.decode('utf-8')
            if decoded.startswith("data:"):
                json_data = json.loads(decoded[5:])
                print(json_data['token']['text'], end="", flush=True)

stream_llm_response("継続的バッチングとは何ですか？")

本番環境の強化：現場で得た教訓

6ヶ月間クラスターでTGIを運用して、いくつかの重要な最適化のコツを学びました。パフォーマンスが頭打ちになったら、以下の3つの観点を確認してください。

量子化でモデルを圧縮する

VRAMが逼迫しているときは量子化を活用しましょう。コマンドに--quantize bitsandbytes-nf4を追加するだけで、メモリ使用量を大幅に削減できます。8Bモデルの場合、VRAMフットプリントを約15GBから6GB未満に圧縮できます。ロジックの精度をほとんど損なうことなく、より安価なハードウェアで大きなモデルを動かせるようになります。QLoRAを使ったファインチューニングと組み合わせれば、コンシューマ向けGPUでもカスタムモデルを効率よく運用できます。

テンソル並列処理でスケールする

13Bを超えるパラメータを持つモデルは、コンシューマ向けGPU1枚に収まらないことがほとんどです。TGIはマルチGPU構成をシンプルに実現します。--num-shardフラグ（例：--num-shard 2）を使うことで、2枚のGPUにモデルを分散できます。ワークロード分散の複雑な計算はTGIが自動的に処理します。

メトリクスを監視する

TGIにはPrometheus向けの/metricsエンドポイントが組み込まれています。tgi_request_queue_sizeを注意深く監視してください。キューが常に高い状態が続いている場合、GPUが飽和しているサインです。本番環境でのAIモデル監視を体系的に行うことで、問題を早期に発見し、ロードバランサーの背後に別のTGIインスタンスを追加するタイミングを的確に判断できます。

より良いAIインフラを構築する

PythonスクリプトからTGIのような専用エンジンへの移行は、AIエンジニアにとって大きな前進です。DockerはこのインフラをU現性高く、スケールしやすい形で提供します。Rustベースの速度と継続的バッチングを組み合わせることで、TGIはオープンソースモデルを大規模に安定してサービングするための強力な手段となります。max-batch-sizeのチューニングとGPU使用率の監視に注力することで、ハードウェアから最大限のパフォーマンスを引き出せるでしょう。

Share:
Tags: