FP8量子化モデルをTensorRTで本番投入する完全ガイド（CLIP実践編）

はじめに：量子化モデルの「本番デプロイ」問題

モデル量子化（Model Quantization）は、ディープラーニングモデルの重みをFP16からFP8に削減し、メモリ使用量を半分近くに抑えつつ推論速度を向上させる技術です。しかし、多くの記事は「量子化チェックポイントの作り方」で終わってしまい、「実際にプロダクションにどう載せるのか？」という実務者の悩みに答えられていません。

本記事では、NVIDIA TensorRT Model OptimizerでFP8量子化したCLIPモデルをONNXにエクスポートし、TensorRTエンジンにコンパイルしてFP16ベースラインと比較する全工程を具体的なコードとコマンドで解説します。

参考資料： 本記事はNVIDIA公式ブログの内容を基に、日本の開発者コミュニティ（Qiita/Zenn）向けに再構成しています。原文はこちらで確認できます。

Step 1: ModelOptチェックポイントをONNXにエクスポート

最初のステップは、ModelOptで生成したFP8チェックポイントをONNX形式に変換することです。ModelOptはget_onnx_bytes_and_metadataというビルトインヘルパーを提供しており、ONNX opset 20+でFP8 QuantizeLinear/DequantizeLinear（Q/DQ）を完全サポートします。

ポイント： ModelOptのエクスポーターは、重み側のQ-DQペアをFP8格納DQ専用チェーンに折り畳み（fold）、ONNXファイルサイズを大幅に削減します。

import torch
from transformers import CLIPModel, CLIPTokenizer
from transformers.models.clip.modeling_clip import CLIPAttention
import modelopt.torch.opt as mto
import modelopt.torch.quantization as mtq
from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata
from modelopt.torch.quantization.plugins.diffusion.diffusers import _QuantAttention

# ONNXエクスポーターが単一forwardのみ見るように薄いラッパークラスを定義
class TextEncoder(torch.nn.Module):
    def __init__(self, m):
        super().__init__(); self.m = m
    def forward(self, x):
        return self.m.get_text_features(x)

class ImageEncoder(torch.nn.Module):
    def __init__(self, m):
        super().__init__(); self.m = m
    def forward(self, x):
        return self.m.get_image_features(x)

def prepare_for_fp8_onnx_export(model):
    """FP8 attention fusionを有効化し、CLIPのfloat scaleをクリア"""
    for _, mod in model.named_modules():
        if isinstance(mod, _QuantAttention):
            mod._disable_fp8_mha = False
        if isinstance(mod, CLIPAttention) and getattr(mod, "scale", None) is not None:
            mod.scale = None

def export(wrapper, dummy, axis_name, out_name):
    """
    ModelOptエクスポーターは重みのQ+DQをFP8格納DQチェーンに折り畳み、
    TRTカスタム演算子をネイティブONNX QDQに書き換える → TRTエンジン構築準備完了
    """
    onnx_bytes, _ = get_onnx_bytes_and_metadata(
        model=wrapper, dummy_input=(dummy,), model_name=out_name,
        dynamic_axes={axis_name: {0: "batch"}}, onnx_opset=20, weights_dtype="fp16",
    )
    OnnxBytes.from_bytes(onnx_bytes).write_to_disk("./onnx_output", clean_dir=False)

# ModelOptチェックポイントからFP8量子化CLIPModelを復元
mto.enable_huggingface_checkpointing()
mtq.QuantModuleRegistry.register({CLIPAttention: "CLIPAttention"})(_QuantAttention)
model = (
    CLIPModel.from_pretrained(modelopt_ckpt, attn_implementation="sdpa", torch_dtype=torch.float16)
    .eval().cuda()
)
prepare_for_fp8_onnx_export(model)

# テキストエンコーダをONNXにエクスポート
tok = CLIPTokenizer.from_pretrained(model_ckpt)
dummy_text = tok(["a photo of a cat"], return_tensors="pt", padding="max_length", max_length=77)["input_ids"].cuda()
export(TextEncoder(model), dummy_text, "text_input", "text_clip_fp8")

# 画像エンコーダをONNXにエクスポート
dummy_image = torch.randn(16, 3, 224, 224, dtype=torch.float16).cuda()
export(ImageEncoder(model), dummy_image, "image_input", "image_clip_fp8")

結果：ONNXファイルサイズ比較

モデル構成要素	FP8 ModelOptチェックポイント	FP16 HuggingFaceチェックポイント	サイズ削減率
CLIPテキストエンコーダONNX	156 MB	237 MB	~34%
CLIP画像エンコーダONNX	292 MB	582 MB	~50%

補足： ONNXファイルサイズの縮小は便宜上のものであり、必須ではありません。TensorRTはエンジン構築時に重み側のQノードをFP8重みに折り畳むためです。ModelOpt ONNXエクスポーターはディスク上のファイルを小さく保つためにONNX側であらかじめ折り畳んでいます。

Step 2: ONNXモデルをTensorRTエンジンにビルド＆プロファイリング

ONNXファイルが準備できたら、TensorRTのtrtexecコマンドラインツールを使用してエンジンをビルドし、ベンチマークを実行します。

重要： --stronglyTypedフラグを使用して、ModelOptがONNXグラフに埋め込んだ精度アノテーションをTensorRTが強制的に従うようにします。これによりFP8重みと活性化が実際のFP8カーネルで実行されます。

# TensorRT環境変数の設定
export PATH=/usr/src/tensorrt/bin:$PATH
export LD_LIBRARY_PATH=/usr/src/tensorrt/lib:$LD_LIBRARY_PATH

# テキストエンコーダエンジンのビルド＆ベンチマーク
trtexec --onnx=text_clip_fp8.onnx \
        --shapes=text_input:128x77 \
        --stronglyTyped \
        --saveEngine=text_clip_fp8.plan

# 画像エンコーダエンジンのビルド＆ベンチマーク
trtexec --onnx=image_clip_fp8.onnx \
        --shapes=image_input:128x3x224x224 \
        --stronglyTyped \
        --saveEngine=image_clip_fp8.plan

FP32型キャスト問題の解決

ModelOptのエクスポーターはattentionスケーリングをFP32ラウンドトリップでラップするため、--stronglyTypedモードではエラーが発生します。以下のスクリプトでFP32初期化子とCast演算子をFP16に再型付けします。

import numpy as np
import onnx
from onnx import TensorProto, numpy_helper, shape_inference

model = onnx.load("clip_fp8.onnx")

# すべてのFP32初期化子をFP16に変換
for init in model.graph.initializer:
    if init.data_type == TensorProto.FLOAT:
        arr = numpy_helper.to_array(init).astype(np.float16)
        init.CopyFrom(numpy_helper.from_array(arr, name=init.name))

# すべてのCast(to=FP32)をCast(to=FP16)に変更
for node in model.graph.node:
    if node.op_type == "Cast":
        to_attr = next(a for a in node.attribute if a.name == "to")
        if to_attr.i == TensorProto.FLOAT:
            to_attr.i = TensorProto.FLOAT16

model = shape_inference.infer_shapes(model, data_prop=True, check_type=False)
# 'strongtyped'接尾辞を付けて保存
onnx.save(model, "clip_fp8_strongtyped.onnx")

ベンチマーク結果（NVIDIA RTX 6000 Ada、TensorRT 10.16、バッチ128）

項目	FP16	FP8	改善率
画像エンコーダエンジンサイズ	588 MB	306 MB	48%削減
テキストエンコーダエンジンサイズ	238 MB	156 MB	34%削減
画像エンコーダレイテンシ	166.2 ms	119.8 ms	1.39倍高速化
テキストエンコーダレイテンシ	13.2 ms	9.1 ms	1.45倍高速化

FP8高速化の秘密：Nsight Deep Learning Designerによる可視化

Nsight Deep Learning DesignerでFP16とFP8のプロファイルを並べると、3つの明確な違いが浮かび上がります。

1. GEMM（行列乗算）区間が約1.8msから0.84msに2倍以上高速化。これはRTX 6000 Ada GPUのFP8 Tensor Coreカーネルによるものです。
2. FP16プロファイルにあった'fusion'レイヤーがFP8では消失。TensorRTがattentionブロック全体を専用FP8 MHAカーネルにルーティングするためです。
3. 精度ドーナツチャートがオレンジ（FP16）から紫（FP8）に変化。量子化された重みと活性化が実際にFP8 Tensor Coreで実行されている証拠です。

注意点： FP8 Tensor CoreはAdaアーキテクチャ（Compute Capability 8.9以上）からのサポートです。RTX 4090、A100、H100などでのみFP8アクセラレーションが有効です。V100やT4環境ではFP16のまま使用する必要があります。データセンター導入前にGPUアーキテクチャを必ず確認してください。

Step 3: TensorRT量子化メカニズムの理解

TensorRTがONNXモデルをインポートする際、QuantizeLinear/DequantizeLinear（Q/DQ）ノードを探します。これらのノードはテンソルがFP8などの低精度データ型に遷移するポイントを示します。

TensorRTは量子化可能なすべてのレイヤーの入力にQ/DQレイヤーペアを要求します。エンジン構築時にオプティマイザがこれらのQ/DQノードを隣接レイヤーに融合（fuse）し、元のレイヤーを低精度テンソルで直接動作する特殊カーネルに置き換えます。このプロセスにより量子化-逆量子化のラウンドトリップが排除され、より高い演算スループットと低いメモリ帯域幅使用を実現します。

実務適用時の注意点と制限

1. GPUアーキテクチャ依存性： FP8はAda（RTX 4000シリーズ）以上またはHopper（H100）以上でのみハードウェアアクセラレーションが可能です。RTX 3090（V100）以下ではソフトウェアエミュレーションとなり、むしろ遅くなる可能性があります。
2. 精度損失： FP8はFP16と比較して表現範囲が狭いため、一部のモデル（特に非常に深いネットワーク）で精度が低下する可能性があります。デプロイ前に必ず検証データセットで精度を測定してください。
3. 動的シェイプ制限： 上記の例では固定バッチ（128）でエンジンをビルドしています。可変バッチが必要な場合は--minShapes、--optShapes、--maxShapesオプションを追加で指定する必要があります。
4. ONNX opsetバージョン： FP8 QDQノードはONNX opset 20以上でのみ完全サポートされます。それ以前のバージョンを使用するとエクスポートに失敗する可能性があります。

次のステップ学習の方向性

• Triton Inference Server連携： --saveEngineで保存した.planファイルをTritonサーバーにデプロイし、REST/gRPCエンドポイントとしてサービングする方法を学びましょう。Metaのプライバシー保護技術ABPの詳細分析では、実際のサービスに適用するアーキテクチャを参考にできます。
• LLM量子化： 本記事はCLIP（ビジョン・言語）モデルを扱っていますが、LLMはTensorRT-LLMという別のパイプラインを必要とします。FP8 LLM量子化とデプロイに興味があれば、TensorRT-LLMチュートリアルを確認してください。
• Nsight Deep Learning Designer応用： ONNXモデル編集、カスタムカーネル挿入、メモリ帯域幅分析などの高度なプロファイリング機能を習得すると、デプロイのボトルネックを正確に特定できます。

まとめ

FP8量子化は単なる「モデル軽量化」を超え、実際のプロダクション環境で1.4倍以上の速度向上とメモリ使用量半減を同時に達成できる強力な技術です。本記事で紹介したModelOpt → ONNX → TensorRTパイプラインを追えば、誰でもFP8量子化モデルを実際のサービスにデプロイできます。

合わせて読みたい記事：

• メッセンジャーでリンクを安全に確認する方法：Metaのプライバシー保護技術ABP詳細分析
• border-radiusの限界を破るCSS corner-shape、実務適用ガイド