🖥️

TensorRT-LLMによるRTX 5090でのLLMのNVFP4量子化・推論

に公開

はじめに

この記事は、TensorRT-LLMを使ってRTX 5090上でLLMのNVFP4量子化と推論を実際に試し、速度を検証してみたという記事です。

最近RTX 5090を購入したので、せっかくなのでBlackwellから対応しているNVFP4の推論をローカルで試したいと思い、この記事を書きました。

NVFP4についてはこちらの記事をご確認ください。

https://developer.nvidia.com/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference

概要

TensorRT-LLMはNVIDIA製のLLMの高速推論のためのフレームワークです。

vLLMやSGLangと同じような高速な推論エンジンですが、これらと違って完全なオープンソースではなく、一部の実装がクローズドになっています。(そもそもTensorRT自体がクローズド)一方、NVIDIA謹製なだけあってGPU側の最適化への対応が比較的早いです。

例えば、vLLMではSM100(B200)でのNVFP4への対応が最近入りましたが、SM120(RTX PRO 6000 BlackwellやRTX 5090など)では未対応です。一方、TensorRT-LLMでは2月ごろからSM100への対応が入っており、SM120でのNVFP4推論も対応済みです。

今回はこのTensorRT-LLMを使い、手元のRTX 5090上でNVFP4量子化と推論を実際に試してみました。

量子化・推論の手順

実際にNVFP4への量子化を行い、推論を実行するまでの手順を以下に示します。今回は環境としてWindows 11のWSL2上のUbuntu 24.04で実行しました。

1. TensorRT-LLMの環境の準備

まずはTensorRT-LLMが動く環境を準備します。pipなどでインストールもできますが、ここでは簡単にDockerイメージを使って準備します。

DockerイメージはNGCで提供されているので、ここから最新のものを利用します。

https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release

今回は1.0.0rc2のイメージを利用します。(最新で上手く動かない場合は少し前のバージョンを利用してください)
以下のようなコマンドでDockerコンテナを動かします。

docker run --rm -it --ipc=host --shm-size=2g --ulimit memlock=-1 \
  --ulimit stack=67108864 --gpus=all -p 8000:8000 -v ~/models:/models \
  nvcr.io/nvidia/tensorrt-llm/release:1.0.0rc2

試しにtrtllm-serve --help等のコマンドを打って問題なく動けば環境は問題ないと思います。

2. モデルのダウンロード

Hugging Faceから推論に利用したいモデルをダウンロードします。普通にhuggingface-cliでダウンロードすれば問題ありません。今回は以下のMistral-Nemoベースの12Bモデルを利用します。

https://huggingface.co/Aratako/NemoAurora-RP-12B

以下のようなコマンドでダウンロードしておきます。

huggingface-cli download Aratako/NemoAurora-RP-12B --local-dir /models/NemoAurora-RP-12B

3. キャリブレーションデータの準備

(この項目はオプションなので、飛ばしても動作には問題ありません。)

量子化の際のキャリブレーションデータを準備します。これは必須ではなく指定しなければデフォルトでcnn_dailymailでキャリブレーションされますが、今回はタスクに近いデータとして実際にこのモデルの学習に利用したデータを使ってキャリブレーションしてみます。

キャリブレーションデータはローカルパスを指定すると以下のように読み込まれるので、これに対応するようなデータを用意しておきます。

https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/quantization/quantize_by_modelopt.py#L410-L418

以下のようなスクリプトでキャリブレーション用のデータを準備しておきます。これをnano等で書き出して実行しておきます。

キャリブレーションデータの準備用スクリプト
import os

from datasets import concatenate_datasets, load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Aratako/NemoAurora-RP-12B")

ds1 = load_dataset(
    "Aratako/Synthetic-Japanese-Roleplay-SFW-DeepSeek-R1-0528-10k-formatted",
    split="train",
)
ds2 = load_dataset(
    "Aratako/Synthetic-Japanese-Roleplay-NSFW-DeepSeek-R1-0528-10k-formatted",
    split="train",
)
ds = concatenate_datasets([ds1, ds2])
ds = ds.shuffle(seed=42).select(range(1000))

ds = ds.map(
    lambda x: {
        "text": tokenizer.apply_chat_template(
            x["messages"],
            tokenize=False,
        ),
    },
    batched=True,
)

os.makedirs("calib_data", exist_ok=True)

ds.to_json(
    "calib_data/train.jsonl",
    orient="records",
    lines=True,
    force_ascii=False,
)

4. 量子化の実行

NVFP4への量子化を実行します。量子化はNVIDIAのModelOptで行われます。

以下のようなコマンドでNVFP4への量子化を行います。ここではKV CacheもFP8に量子化しています。ステップ3をスキップした方はcalib_datasetの指定を削除してください。

# model_dir: 量子化対象のモデル
# dtype: モデル読み込みの際のdtype
# qformat: 量子化の形式
# output_dir: 量子化後のcheckpointの出力先
# calib_size: キャリブレーションに使うデータの量
# calib_dataset: キャリブレーションに使うデータ(省略可能)
# kv_cache_dtype: KV Cacheのdtype
python /app/tensorrt_llm/examples/quantization/quantize.py \
  --model_dir /models/NemoAurora-RP-12B \
  --dtype bfloat16 \
  --qformat nvfp4 \
  --output_dir /models/quantized \
  --calib_size 512 \
  --calib_dataset /app/tensorrt_llm/calib_data \
  --kv_cache_dtype fp8

5. TensorRTエンジンのbuild

trtllm-buildコマンドでTensorRTエンジンをbuildします。これにより、TensorRT-LLMで最適化された推論エンジンの形式に変換されます。

以下のようなコマンドでengineのbuildを行います。各種設定は特に最適化したわけではないので、気になる方は色々と調べてみてください。

# checkpoint_dir: 先ほど量子化したcheckpointのパス
# output_dir: engineの出力先
# gemm_plugin: GEMM pluginの設定
# kv_cache_type: KV Cacheのタイプ
# max_batch_size: 一度にschedule出来る最大のバッチサイズ
# max_seq_len: 取り扱える最大シーケンス長
# max_num_tokens: バッチ内の最大合計トークン数
# use_paged_context_fmha: KV Cache Reuseなどの有効化
# use_fp8_context_fmha: FP8 Context FMHAの有効化
# multiple_profiles: 複数のoptimization profilesを有効に
trtllm-build \
  --checkpoint_dir /models/quantized \
  --output_dir /models/engine \
  --gemm_plugin auto \
  --kv_cache_type paged \
  --max_batch_size 64 \
  --max_seq_len 32768 \
  --max_num_tokens 1024 \
  --use_paged_context_fmha enable \
  --use_fp8_context_fmha enable \
  --multiple_profiles enable

6. 推論サーバの起動

trtllm-serveコマンドで上で作成したエンジンを使ったOpenAI Compatibleな推論サーバを立ち上げます。

以下のようなコマンドで推論サーバを起動します。

# Chunked Prefillを有効に
echo "enable_chunked_prefill: true" > extra_llm_api_config.yml
# 推論サーバの起動
trtllm-serve /models/engine --tokenizer /models/NemoAurora-RP-12B --max_batch_size 64 --max_seq_len 32768 --max_num_tokens 1024 --host 0.0.0.0 --extra_llm_api_options /app/tensorrt_llm/extra_llm_api_config.yml

7. 推論の実行

実際に推論を試してみます。別のターミナルから以下のようにcurlでリクエストを送ってみます。

curl http://localhost:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "Aratako/NemoAurora-RP-12B",
        "messages": [
            {"role": "user", "content": "こんにちは"}
        ]
    }'

無事レスポンスが返ってきました。

{"id":"chatcmpl-95ba9e29d45d4e2c8f460d7fd80c8ed2","object":"chat.completion","created":1752380136,"model":"Aratako/NemoAurora-RP-12B","choices":[{"index":0,"message":{"role":"assistant","content":"おかえりなさいませ。今日はお忙しかったですか?","reasoning_content":null,"tool_calls":[]},"logprobs":null,"finish_reason":"stop","stop_reason":null,"disaggregated_params":null}],"usage":{"prompt_tokens":11,"total_tokens":25,"completion_tokens":14},"prompt_token_ids":null}

OpenAI Compatibleなので、普通にOpenAIライブラリ等からも推論が可能です。

推論速度のベンチマーク

無事NVFP4での推論ができたので、速度の検証も行ってみます。ここではFP8と比較してみます。なお、FP8のモデルは上述した手順中の量子化の際の--qformatfp8にするだけで作成可能です。

今回はベンチマークツールとして、Hugging Faceが公開しているinference-benchmarkerというものを利用してみます。

https://github.com/huggingface/inference-benchmarker

手順

以下のような手順でinference-benchmarkerによるベンチマークが実行できます。

1. リポジトリのクローン

inference-benchmarkerのリポジトリをクローンして手元に持ってきます。

git clone https://github.com/huggingface/inference-benchmarker.git

2. Dockerイメージのbuild

inference-benchmarkerのDockerイメージをbuildします。

cd inference-benchmarker
docker build -t inference_benchmarker_latest .

3. ベンチマークの実行

あらかじめtrtllm-serveで推論サーバを立ち上げた後、先ほどのDockerイメージを使って実際にベンチマークを実行します。
各種オプションについては元リポジトリをご確認ください。

docker run --network host \
  -v ~/inference-benchmarker-results:/opt/inference-benchmarker/results \
  inference_benchmarker_latest inference-benchmarker --no-console \
  --url http://host.docker.internal:8000/v1 \
  --max-vus 100 --duration 120s --warmup 30s --benchmark-kind rate \
  --rates 1.0 --rates 10.0 --rates 30.0 --rates 100.0 \
  --prompt-options "num_tokens=200,max_tokens=220,min_tokens=180,variance=10" \
  --decode-options "num_tokens=200,max_tokens=220,min_tokens=180,variance=10" \
  --model-name "Aratako/NemoAurora-RP-12B" \
  --tokenizer-name "Aratako/NemoAurora-RP-12B"

ベンチマーク結果

NVFP4とFP8でそれぞれベンチマークを行った結果、以下のような結果になりました。

それぞれの指標は以下のようなものです。

  • QPS(Queries Per Second): 推論サーバが1秒間あたりに処理できるリクエスト数。大きいほどサーバのリクエスト処理性能が高い。
  • TTFT(Time To First Token): リクエスト送信から最初のトークンを受け取るまでの時間。小さいほどサーバのレスポンスタイムが早い。
  • ITL(Inter-Token Latency): 各トークン1つ1つが生成される間隔の時間。小さいほどストリーミング応答が滑らかになる。
  • Throughput: 全ての処理の中で1秒あたりに生成された合計のトークン数。大きいほどサーバが大量のテキストを高速に生成できる。

1. QPS

10 req/sあたりまではどちらも余裕があり横並びですが、より高負荷な条件においてはNVFP4の方が多くのリクエストをさばけていることがわかります。

NVFP4 vs FP8(QPS)

Benchmark NVFP4 (req/s) FP8 (req/s)
warmup 0.58 0.62
constant @ 1 req/s 1.01 0.99
constant @ 10 req/s 9.78 9.85
constant @ 30 req/s 24.53 19.07
constant @ 100 req/s 25.56 18.68

2. TTFT

余裕を持ってリクエストをさばけている10 req/sまではNVFP4の方がFP8よりも10%程度TTFTが早く、優位に立っています。
逆に30 req/s以降はFP8の方がTTFTが早くなっていますが、QPSではNVFP4が上回っていたことを考えると、NVFP4の方がMemory Footprintが軽くKV Cacheを多く持てるため実効バッチサイズが大きくなり、Throughputを向上させる代わりにTTFTがやや悪くなっているようなことが推測できます。

NVFP4 vs FP8(TTFT)

Benchmark NVFP4 (ms) FP8 (ms)
warmup 60.30 64.55
constant @ 1 req/s 53.62 57.78
constant @ 10 req/s 58.16 66.29
constant @ 30 req/s 2355.64 1818.88
constant @ 100 req/s 2426.88 1977.71

3. ITL

すべての条件においてNVFP4が勝っている結果になりました。NVFP4の30 req/s以上がなぜかそれ以下よりも良いという結果になっていますが、これはちょっと良く分かりません。(少し考えられない結果なので普通にバグ?)

NVFP4 vs FP8(ITL)

Benchmark NVFP4 (ms) FP8 (ms)
warmup 11.31 13.79
constant @ 1 req/s 11.51 11.87
constant @ 10 req/s 12.96 15.48
constant @ 30 req/s 7.45 19.71
constant @ 100 req/s 7.49 19.90

4. Throughput

基本的にはNVFP4の方が優位に立つ結果となりました。特に30 req/s以上では顕著で、高負荷条件下ではNVFP4の方がより多くのリクエストを同時にさばけていることがここからも分かります。

NVFP4 vs FP8(Throughput)

Benchmark NVFP4 (tok/s) FP8 (tok/s)
warmup 85.29 75.39
constant @ 1 req/s 158.53 154.48
constant @ 10 req/s 1591.56 1602.03
constant @ 30 req/s 3840.57 3000.33
constant @ 100 req/s 4007.33 3019.02

まとめ

この記事では、TensorRT-LLMを使ってRTX 5090上でLLMのNVFP4量子化と推論を実際に試し、速度を検証してみました。

ベンチマークの結果、全体的にNVFP4の方が処理が早く、特に高負荷条件において処理性能の向上が明確にみられるという結果になりました。今後さらにカーネルの最適化が進むにつれ、よりNVFP4の推論が有効に働くようになることが期待されます。

Discussion