ml-inference-optimization

ML 推論の最適化

概要

推論レイテンシの短縮、スループットの向上、メモリ使用量の削減、コストの低減、最適化済みモデルの安全なデプロイには、このスキルを使用します。最適化はベンチマークに基づいて行う必要があります。モデルを変更する前に、ワークロード、入力形状、同時実行数、SLO、ハードウェア、ランタイム、数値許容誤差、品質指標を定義してください。

最適化のワークフロー

1. 現実的なデータ、前処理、後処理、ネットワークのオーバーヘッド、ウォームアップ、同時実行数を含めてベースラインを確立します。
2. ボトルネックをプロファイリングします。CPU 前処理、モデル計算、メモリ帯域幅、GPU 使用率、シリアライズ、キューイング、検索、下流への呼び出しなどを確認します。
3. リスクの最も低い最適化から適用します。バッチ処理、コンパイル、精度、ランタイムのチューニング、圧縮の順に進めます。
4. 変更のたびに、精度、キャリブレーション、公平性のスライス、堅牢性、数値安定性を再検証します。
5. p50、p95、p99 のレイテンシ、スループット、メモリ、コールドスタート、エラー率、予測 1,000 件あたりのコストを測定します。

エクスポートとコンパイル

エクスポート方法はデプロイ先に応じて選択します。代表的な入力を使って、エクスポート後の出力をフレームワークの出力と比較してください。出力の種類ごとに許容誤差を定義します。ベンチマークには前処理と後処理も含めてください。多くの不具合はモデル本体の外で発生します。

PyTorch から ONNX へのエクスポートのひな形

import torch
import torch.nn as nn

def export_to_onnx(model: nn.Module, dummy_input: torch.Tensor, onnx_path: str):
    model.eval()
    
    # Define input and output names, and dynamic axes for batch size and sequence length
    input_names = ["input_ids"]
    output_names = ["logits"]
    dynamic_axes = {
        "input_ids": {0: "batch_size", 1: "seq_len"},
        "logits": {0: "batch_size"}
    }
    
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,
        opset_version=17, # Modern stable opset version
        do_constant_folding=True,
        input_names=input_names,
        output_names=output_names,
        dynamic_axes=dynamic_axes
    )

PyTorch torch.compile 最適化のセットアップ

torch.compile は Inductor コンパイラーを活用し、最新の GPU で大幅な高速化を実現します。

import torch

def compile_model(model: torch.nn.Module, mode: str = "default"):
    # Modes:
    # - "default": Balance compile time and inference speed
    # - "reduce-overhead": Reduce framework overhead (uses CUDA graphs, good for small batch sizes)
    # - "max-autotune": Perform exhaustive kernel tuning (longest compile time, fastest runtime)
    
    optimized_model = torch.compile(
        model, 
        mode=mode, 
        fullgraph=False,
        dynamic=True # Handles variable sequence lengths/batch sizes efficiently
    )
    return optimized_model

バッチ処理と同時実行

動的バッチ処理は、GPU エンドポイントのサービング性能を最も大きく向上させる方法であることがよくあります。最大バッチサイズとキューの待ち時間は合わせて調整してください。バッチサイズを大きくするとスループットは向上しますが、テールレイテンシが悪化する可能性があります。LLM では、プリフィル段階とデコード段階を区別してください。メモリ使用量は、シーケンス長、KV キャッシュ、同時リクエスト数に大きく左右されます。適切な場合は、継続的バッチ処理または LLM 専用のサービングランタイムを使用します。

CPU サービスでは、メモリ帯域幅、GIL／ランタイムの制約、またはスレッドの過剰な割り当てが支配的になるまでは、同時実行数を増やすと使用率が向上することがあります。BLAS／OpenMP／ONNX Runtime のスレッド数を固定し、デフォルト設定をそのまま受け入れるのではなく、ベンチマークを実施してください。

Triton Inference Server の設定

Triton は、高性能なマルチフレームワークサービングの業界標準です。

Triton config.pbtxt テンプレート

name: "text_classifier"
platform: "onnxruntime_onnx"
max_batch_size: 32

input [
  {
    name: "input_ids"
    data_type: TYPE_INT64
    dims: [ -1 ] # Dynamic sequence length
  }
]

output [
  {
    name: "logits"
    data_type: TYPE_FP32
    dims: [ 5 ] # Number of classes
  }
]

# Enable Dynamic Batching
dynamic_batching {
  max_queue_delay_microseconds: 5000 # Wait up to 5ms for concurrent requests
  preferred_batch_size: [ 8, 16, 32 ]
}

# GPU Instance Groups
instance_group [
  {
    count: 2 # Host 2 model execution instances per GPU
    kind: KIND_GPU
  }
]

精度と量子化

方法	適した用途	注意点
fp16/bf16	GPU／アクセラレーターのスループット向上とメモリ削減	数値に敏感な演算を検証してください
INT8 ポストトレーニング量子化	キャリブレーションデータを用意できる CNN、トランスフォーマー、表形式データ向けニューラルネット	キャリブレーションセットは本番環境を代表するものである必要があります
量子化を考慮した学習	精度が重視される低精度デプロイ	学習が複雑になります
INT4/GPTQ/AWQ/bitsandbytes	大規模言語モデルのメモリ削減	推論、まれなトークン、キャリブレーション、ツール使用の性能が低下する可能性があります
重みのみの量子化	活性化の精度を高く保つ LLM サービング	高速化の度合いはカーネル／ハードウェアに依存します

量子化によってモデルの振る舞いは変化します。タスク指標、キャリブレーション、スライス指標、安全性チェックを再実行してください。RAG または埋め込みモデルでは、量子化後に検索再現率と埋め込み類似度の分布を検証してください。

プルーニング、スパース性、蒸留

プルーニングでは、重み、チャネル、ヘッド、レイヤーを削除します。非構造化スパース性は、ハードウェア／ランタイムのサポートがなければ推論速度が向上しない場合があります。構造化プルーニングのほうが、レイテンシを低減しやすい傾向があります。知識蒸留では、小さな生徒モデルを学習させ、大きな教師モデルの振る舞いを模倣させます。

TensorRT コンパイルレシピ（trtexec）

NVIDIA TensorRT でコンパイルされたエンジンは、GPU で最大のスループットを実現します。

1. 動的シェイプを使用して FP16 エンジンを構築する

trtexec --onnx=model.onnx \
        --saveEngine=model.engine \
        --fp16 \
        --minShapes=input_ids:1x1 \
        --optShapes=input_ids:16x128 \
        --maxShapes=input_ids:32x512

2. キャリブレーションデータを使用して INT8 エンジンを構築する

trtexec --onnx=model.onnx \
        --saveEngine=model_int8.engine \
        --int8 \
        --calib=calibration_profile.cache \
        --minShapes=input_ids:1x1 \
        --optShapes=input_ids:16x128 \
        --maxShapes=input_ids:32x512

サービングランタイム

ランタイム	使用する場面
Triton Inference Server	複数のフレームワークに対応した GPU／CPU サービング、動的バッチ処理、アンサンブル、モデルリポジトリ、高スループットが必要な場合
TorchServe	ハンドラーと管理 API を使用して PyTorch モデルをサービングする場合
TensorFlow Serving	TensorFlow SavedModel のサービングとバージョン管理されたモデル管理が必要な場合
ONNX Runtime Server/custom API	軽量なサービス制御でポータブルな ONNX 推論を行う場合
BentoML	モデルを API としてパッケージ化し、デプロイの柔軟性と Python 中心のワークフローが必要な場合
Seldon/KServe	Kubernetes ネイティブのモデルサービング、カナリアリリース、推論グラフ、エンタープライズ運用が必要な場合
vLLM / TGI	継続的バッチ処理、PagedAttention、KV キャッシュ管理、テンソル並列による LLM サービングが必要な場合

モデルリポジトリのバージョン管理とレディネスチェックを使用してください。ヘルスチェックは負荷の高い推論処理と分離してください。GPU サービングでは、容量が実証されていない限り、1 台のデバイスにメモリを大量に消費するモデルを複数ロードしないでください。

メモリの最適化

Transformer では、メモリ使用量は重みと KV キャッシュによって決まります。量子化、テンソル並列処理、最大シーケンス長の短縮、PagedAttention、キャッシュの削除、バッチサイズ／同時実行数の縮小、または蒸留モデルによってメモリ使用量を削減できます。拡散モデルやビジョンモデルでは、レイテンシが許容される場合、アテンションスライシング、VAE タイリング、fp16/bf16、メモリ効率の高いアテンション、モデルのオフロードを使用してください。

エッジデプロイ

エッジ環境の制約には、バイナリサイズ、RAM、熱によるスロットリング、オフライン動作、断続的な更新、プライバシー、ハードウェアアクセラレーションがあります。対象に応じて TFLite/Core ML/OpenVINO/ONNX を選択してください。実機でテストしてください。

ベンチマークでの落とし穴

• 本番環境では同時にトラフィックがあるにもかかわらず、ウォームアップ済みの単一リクエストのレイテンシだけを測定する。
• トークン化、画像のデコード、検索、シリアライズ、ネットワーク時間を無視する。
• p95/p99 ではなく平均レイテンシを報告する。
• 量子化モデルとベースラインモデルを異なるデータで比較する。
• ある形状用に TensorRT エンジンを構築し、サポートされていない動的形状をサービングする。
• スループットが向上するまでバッチサイズを増やし、SLO を満たせなくする。

参考資料

• NVIDIA Triton Inference Server のドキュメント：https://docs.nvidia.com/deeplearning/triton-inference-server/
• NVIDIA TensorRT のドキュメント：https://docs.nvidia.com/deeplearning/tensorrt/
• ONNX Runtime のドキュメント：https://onnxruntime.ai/docs/
• PyTorch のエクスポートおよびコンパイルに関するドキュメント：https://pytorch.org/docs/stable/
• TensorFlow Lite のドキュメント：https://www.tensorflow.org/lite
• BentoML のドキュメント：https://docs.bentoml.com/