mdx.jpのA100 x2でLLM-jp-4 32B ThinkingをvLLMサーブし、OpenAI互換APIから使う

LLM-jp-4-32b-a3b-thinking の公式版を、mdx.jp 上の A100 40GB x2 サーバで動かし、最終的に vLLM の OpenAI 互換 API として使えるところまで確認しました。

結論から書くと、Transformers + device_map="auto" では生成時に OOM しましたが、vLLM + tensor_parallel_size=2 に切り替えるとサーブ成功しました。ただし、thinking モデルを OpenAI 互換 API でそのまま使うと analysis が露出するため、クライアント側で assistant final を抜く後処理が必要でした。

背景

やりたかったことは以下です。

LLM-jp-4-32b-a3b-thinking の公式版を使う
A100 サーバ上で安定して動かす
できれば OpenAI互換API で他クライアントから叩けるようにする
日本語 OCR テキストから表情・感情表現を抽出するバッチへつなぐ

ローカル Mac では GGUF + Ollama で試していましたが、32B-A3B の品質評価としては公式実装でも見ておきたかったため、GPU サーバ側で再検証しました。

環境

実行基盤: mdx.jp
OS: Ubuntu 22.04 LTS
GPU: NVIDIA A100-SXM4-40GB x 2
CUDA Driver: 590.48.01
CUDA Version: 13.1
仮想ディスク: 360GB
モデル: llm-jp/llm-jp-4-32b-a3b-thinking

nvidia-smi では以下のように 2 枚見えていました。

|   0  NVIDIA A100-SXM4-40GB          On  | ... | 0MiB / 40960MiB |
|   1  NVIDIA A100-SXM4-40GB          On  | ... | 0MiB / 40960MiB |

今回は mdx.jp 上で、Ubuntu 22.04、A100 40GB x2、仮想ディスク 360GB の仮想マシンを使っています。モデル本体のダウンロードや Hugging Face キャッシュを考えると、360GB あるのはかなり助かりました。

最初の試行: Transformers でそのまま動かす

まずは Hugging Face の公式モデルカードに近い形で、Transformers から直接ロードしました。

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODEL_ID = "llm-jp/llm-jp-4-32b-a3b-thinking"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

モデルのダウンロードと重みロード自体は成功しました。

Download complete: 100%|...| 64.3G/64.3G
Loading weights: 100%|...| 355/355

ところが、生成に入ったところで OOM しました。

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 8.79 GiB.
GPU 0 has a total capacity of 39.49 GiB of which 6.11 GiB is free.

これは「ロードは入るが、生成時のワーキングメモリが足りない」パターンです。device_map="auto" では 2GPU に分散されても、MoE 部分の計算で GPU 0 に偏りが出ていました。

この時点での判断は次の通りでした。

このサーバで不可能なのではない
Transformers の素朴な載せ方では厳しい
サーブ前提なら vLLM に切り替えたほうが早い

vLLM へ切り替える

今回は Docker + vLLM を使いました。起動コマンドは以下です。

docker run --runtime nvidia --gpus all \
  --ipc=host \
  -p 8000:8000 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e HUGGING_FACE_HUB_TOKEN=$HF_TOKEN \
  vllm/vllm-openai:latest \
  --model llm-jp/llm-jp-4-32b-a3b-thinking \
  --trust-remote-code \
  --tensor-parallel-size 2

ポイントは以下です。

--trust-remote-code
--tensor-parallel-size 2
Hugging Face キャッシュをマウント

vLLM 起動ログ

起動ログの重要なところだけ抜くとこうです。

Resolved architecture: Qwen3MoeForCausalLM
Using max model len 65536
tensor_parallel_size=2

重みロード後は、各 GPU でおおむね次のような使用量になりました。

Model loading took 29.99 GiB memory
GPU KV cache size: 149,168 tokens
Maximum concurrency for 65,536 tokens per request: 2.28x

ここまで行けば、Transformers での OOM を超えて、実用的なサーブに入れます。

OpenAI 互換 API の確認

まずはモデル一覧を確認しました。

curl http://localhost:8000/v1/models

返ってきたレスポンスは以下の通りです。

{
  "object": "list",
  "data": [
    {
      "id": "llm-jp/llm-jp-4-32b-a3b-thinking",
      "object": "model",
      "max_model_len": 65536
    }
  ]
}

この時点で、vLLM サーバ自体は成功しています。

ただし thinking モデルはそのままだと扱いにくい

次に chat completion を投げると、返ってきたのは期待した日本語回答そのものではなく、内部的な analysis を含む長いテキストでした。

たとえば次のようなプロンプトです。

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "llm-jp/llm-jp-4-32b-a3b-thinking",
    "messages": [
      {"role": "system", "content": "あなたは日本語の研究支援を行うアシスタントです。"},
      {"role": "user", "content": "『ず笑ふて』という表現は、どのような感情に近いですか。簡潔に答えてください。"}
    ],
    "max_tokens": 512,
    "temperature": 0
  }'

返答は概ね以下のような形でした。

analysis ...
assistant final 悲しみや失望に近い感情です。

つまり、thinking モデルを OpenAI 互換 API にそのまま載せると、最終回答の前に内部 analysis が露出します。

tokenizer.parse_response は使えなかった

最初は Hugging Face 側の tokenizer を読み込み、parse_response() を使えば素直に最終回答を取れると考えました。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "llm-jp/llm-jp-4-32b-a3b-thinking",
    trust_remote_code=True,
)

print(tokenizer.parse_response(raw_text))

しかし、実際には以下のようにほとんど情報が取れませんでした。

{'role': 'assistant'}

つまり、vLLM が返す OpenAI 互換の文字列形式と、tokenizer.parse_response() が期待している形式が一致していませんでした。

実用上の回避策: assistant final を抜く

そこで、当面の回避策として assistant final 以降を自前で取り出すことにしました。

import re

match = re.search(r"assistant final\s*(.*)$", raw, flags=re.DOTALL)
if match:
    final_text = match.group(1).strip()
else:
    final_text = raw.strip()

この方法なら、最終回答だけはきちんと取れます。

たとえば先ほどの例では、最終的に以下が抽出できました。

悲しみや失望に近い感情です。

OCR テキスト抽出への応用

この構成をもとに、OCR テキスト 1 ページを読んで表情・感情表現を抽出するスクリプトも作りました。

流れはこうです。

.txt を読む
vLLM の chat/completions に投げる
assistant final を抜く
その中の JSON 配列だけを raw_decode で抽出する
pages を補完して保存する

ここで重要だったのは、モデルが JSON 配列の後ろに説明文を付けることがある点です。最初は正規表現で \[.*\] を抜いていましたが、これだと貪欲に取りすぎて json.decoder.JSONDecodeError: Extra data が起きました。

最終的には json.JSONDecoder().raw_decode() で最初に現れる配列だけを取るようにしました。

ここまでの結論

今回の取り組みをまとめると、以下です。

A100 40GB x2 で LLM-jp-4-32b-a3b-thinking の公式版は動く
Transformers + device_map="auto" では生成時に OOM した
vLLM + tensor_parallel_size=2 ならサーブできた
OpenAI 互換 API 自体も成立する
ただし thinking モデルは内部 analysis が露出する
tokenizer.parse_response() は今回の vLLM 出力には乗らなかった
当面は assistant final を抜くクライアント側後処理が必要

つまり、インフラとしてはかなり前進しました。一方で、アプリケーションとしてそのまま使うには、thinking モデル特有のレスポンス整形が必要です。

次にやりたいこと

次の課題は以下です。

OCR 1ページだけでなく、複数ページをまとめて回す
表情抽出プロンプトの精度調整
assistant final 依存を減らす
base モデルとの比較
vLLM 側での stop 条件やテンプレート調整

特に「thinking モデルを API サーブしてそのまま業務処理に使う」部分は、まだ少し詰める余地があります。

参考

LLM-jp-4 32B A3B Thinking: https://huggingface.co/llm-jp/llm-jp-4-32b-a3b-thinking
vLLM OpenAI Compatible Server: https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
vLLM Docker: https://docs.vllm.ai/en/latest/deployment/docker.html

背景#

環境#

最初の試行: Transformers でそのまま動かす#

vLLM へ切り替える#

vLLM 起動ログ#

OpenAI 互換 API の確認#

ただし thinking モデルはそのままだと扱いにくい#

tokenizer.parse_response は使えなかった#

実用上の回避策: assistant final を抜く#

OCR テキスト抽出への応用#

ここまでの結論#

次にやりたいこと#

参考#

背景

環境