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 推理服务器配置

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、transformer、表格神经网络	校准集必须能代表生产数据
量化感知训练	对精度敏感的低精度部署	增加训练复杂度
INT4/GPTQ/AWQ/bitsandbytes	减少大语言模型的内存占用	可能影响推理能力、罕见 token、校准和工具使用
仅权重量化	激活值保持较高精度的 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 服务，除非已验证设备容量足够，否则不要在每台设备上加载多个占用大量内存的模型。

内存优化

对于 Transformer，内存占用由权重和 KV 缓存决定。可通过量化、张量并行、缩短最大序列长度、分页注意力、缓存淘汰、减小批次大小或并发量，或使用蒸馏模型来降低内存占用。对于扩散模型和视觉模型，可使用注意力切片、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/