模型部署

训练好的模型最终需要部署到生产环境中提供服务。本章将介绍如何导出、优化和部署 Transformers 模型，包括 ONNX 导出、模型量化、推理优化等内容。

部署流程概览

模型从训练到部署通常经历以下阶段：

┌─────────────────────────────────────────────────────────────────┐
│                      模型部署流程                                │
├─────────────────────────────────────────────────────────────────┤
│                                                                  │
│  训练阶段          导出阶段          优化阶段          部署阶段  │
│      │                │                │                │       │
│      ▼                ▼                ▼                ▼       │
│  ┌───────┐       ┌─────────┐      ┌─────────┐      ┌─────────┐ │
│  │ PyTorch│       │  ONNX   │      │  量化   │      │ 服务封装 │ │
│  │ 模型   │──────▶│  导出   │─────▶│  优化   │─────▶│  API    │ │
│  │ 权重   │       │  格式   │      │  剪枝   │      │  接口   │ │
│  └───────┘       └─────────┘      └─────────┘      └─────────┘ │
│                                                                  │
│  开发环境                           生产环境                     │
│                                                                  │
└─────────────────────────────────────────────────────────────────┘

ONNX 导出

ONNX（Open Neural Network Exchange）是一种开放的模型格式，支持跨框架部署。导出为 ONNX 后，可以使用 ONNX Runtime、TensorRT 等高性能推理引擎。

使用 Optimum 导出

Hugging Face 提供的 Optimum 库简化了 ONNX 导出流程：

pip install optimum[onnx]

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer

# 直接加载为 ONNX 模型
model_id = "distilbert-base-uncased-finetuned-sst-2-english"
model = ORTModelForSequenceClassification.from_pretrained(model_id, export=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)

# 推理
inputs = tokenizer("I love this movie!", return_tensors="pt")
outputs = model(**inputs)
print(outputs.logits)

使用 transformers.onnx 导出

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from transformers.onnx import FeaturesManager

# 加载模型和分词器
model_id = "distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForSequenceClassification.from_pretrained(model_id)

# 获取 ONNX 配置
model_type = model.config.model_type
feature = "sequence-classification"
onnx_config = FeaturesManager.get_config(model_type, feature)(model.config)

# 导出为 ONNX
from transformers.onnx import export

onnx_inputs, onnx_outputs = export(
    preprocessor=tokenizer,
    model=model,
    config=onnx_config,
    opset=14,
    output_path="model.onnx"
)

print(f"导出成功！输入: {onnx_inputs}, 输出: {onnx_outputs}")

命令行导出

# 使用 transformers 命令行工具导出
python -m transformers.onnx --model=distilbert-base-uncased-finetuned-sst-2-english --feature=sequence-classification output_dir/

# 支持的特征类型
# - sequence-classification: 序列分类
# - token-classification: Token 分类
# - question-answering: 问答
# - feature-extraction: 特征提取
# - text-generation: 文本生成

使用 ONNX Runtime 推理

import onnxruntime as ort
from transformers import AutoTokenizer
import numpy as np

# 加载 ONNX 模型
session = ort.InferenceSession("model.onnx")

# 获取输入输出信息
input_names = [inp.name for inp in session.get_inputs()]
output_names = [out.name for out in session.get_outputs()]

print(f"输入: {input_names}")
print(f"输出: {output_names}")

# 准备输入
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
text = "This is a great product!"
inputs = tokenizer(text, return_tensors="np")

# ONNX Runtime 推理
outputs = session.run(
    output_names,
    {
        "input_ids": inputs["input_ids"].astype(np.int64),
        "attention_mask": inputs["attention_mask"].astype(np.int64)
    }
)

# 解析结果
logits = outputs[0]
predicted_class = np.argmax(logits, axis=-1)[0]
print(f"预测类别: {predicted_class}")

模型量化详解

量化通过降低模型参数的精度来减少模型大小和提升推理速度，是部署优化的重要手段。Transformers 支持多种量化方法，每种方法都有其特点和适用场景。

量化方法概览

量化方法	位数	特点	适用场景
bitsandbytes	8-bit/4-bit	易用性强，无需校准数据	快速部署、原型验证
AWQ	4-bit	保持精度，推理速度快	高性能推理
GPTQ	2-8 bit	灵活配置，支持多种后端	平衡精度与速度
Quanto	2-8 bit	设备无关，支持多平台	跨平台部署
AQLM	2-4 bit	极致压缩	资源受限环境
HQQ	1-8 bit	无需校准数据	快速量化

bitsandbytes 量化

bitsandbytes 是最常用的量化库之一，支持 8-bit 和 4-bit 量化，使用简单且效果稳定。

8-bit 量化

8-bit 量化使用 LLM.int8() 算法，通过向量级量化处理离群值，在保持模型性能的同时大幅减少显存占用。

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# 方法1：直接使用 load_in_8bit 参数
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    load_in_8bit=True,
    device_map="auto"
)

# 方法2：使用 BitsAndBytesConfig 进行详细配置
bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_threshold=6.0,              # 离群值检测阈值
    llm_int8_skip_modules=None,          # 跳过量化的模块
    llm_int8_enable_fp32_cpu_offload=False,  # CPU 卸载
    llm_int8_has_fp16_weight=False       # 是否保留 fp16 权重
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# 推理
inputs = tokenizer("The future of AI is", return_tensors="pt").to(model.device)
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=50)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

8-bit 量化的关键参数说明：

• llm_int8_threshold：离群值检测阈值。大于此阈值的激活值会被视为离群值，使用 fp16 精度计算。默认值 6.0 适用于大多数模型，对于不稳定的小模型可能需要降低阈值。

• llm_int8_skip_modules：指定不进行量化的模块列表。某些模型的特定模块需要保持原始精度，例如 Whisper 的编码器、LLaVA 的视觉编码器等。

4-bit 量化

4-bit 量化使用 NF4（NormalFloat4）或 FP4 数据类型，配合双重量化技术，可以在几乎不损失精度的情况下将模型大小减少到原来的 1/4。

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

# 4-bit 量化配置
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",           # 量化类型: "fp4" 或 "nf4"
    bnb_4bit_compute_dtype=torch.float16,  # 计算精度
    bnb_4bit_use_double_quant=True,      # 双重量化，进一步压缩
    bnb_4bit_quant_storage=torch.uint8   # 存储类型
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantization_config=bnb_config,
    device_map="auto"
)

print(f"模型显存占用: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")

4-bit 量化的关键参数说明：

• bnb_4bit_quant_type：量化数据类型。NF4（NormalFloat4）专门为正态分布的权重设计，通常比 FP4 效果更好。

• bnb_4bit_compute_dtype：计算时使用的数据类型。通常使用 fp16 或 bf16，在保持精度的同时加速计算。

• bnb_4bit_use_double_quant：是否使用双重量化。开启后会对量化常数再次量化，可以额外节省约 0.5GB 显存。

AWQ 量化

AWQ（Activation-aware Weight Quantization）是一种激活感知的权重量化方法，通过保留对模型性能重要的少量权重，在 4-bit 量化下几乎不损失性能。

加载预量化的 AWQ 模型

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载预量化的 AWQ 模型
model_id = "TheBloke/zephyr-7B-alpha-AWQ"
tokenizer = AutoTokenizer.from_pretrained(model_id)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 推理
inputs = tokenizer("Explain quantum computing in simple terms:", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

AWQ 配置详解

可以通过 AwqConfig 进行详细配置：

from transformers import AutoModelForCausalLM, AwqConfig

# 基本配置
quantization_config = AwqConfig(
    bits=4,                    # 量化位数，固定为 4
    group_size=128,            # 分组大小，推荐 128
    zero_point=True,           # 是否使用零点量化
    version="gemm"             # 内核版本: "gemm" 或 "exllama"
)

model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Mistral-7B-OpenOrca-AWQ",
    quantization_config=quantization_config,
    device_map="auto"
)

启用融合模块提升性能

AWQ 支持融合模块（Fused Modules），可以显著提升推理速度：

from transformers import AutoModelForCausalLM, AwqConfig
import torch

# 启用融合模块
quantization_config = AwqConfig(
    bits=4,
    fuse_max_seq_len=512,      # 最大序列长度，包含输入和生成
    do_fuse=True               # 启用融合
)

model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Mistral-7B-OpenOrca-AWQ",
    quantization_config=quantization_config,
    torch_dtype=torch.float16,
    device_map="auto"
)

性能对比：融合模块可以将解码速度从约 35 tokens/s 提升到约 90 tokens/s，同时显存占用略有减少。

使用 ExLlamaV2 后端

ExLlamaV2 内核提供更快的推理速度：

from transformers import AutoModelForCausalLM, AwqConfig

# 安装: pip install git+https://github.com/casper-hansen/AutoAWQ.git

quantization_config = AwqConfig(
    bits=4,
    version="exllama"          # 使用 ExLlamaV2 内核
)

model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Mistral-7B-Instruct-v0.1-AWQ",
    quantization_config=quantization_config,
    device_map="auto"
)

GPTQ 量化

GPTQ 是一种后训练量化技术，通过逐行量化权重矩阵找到最小化误差的权重版本。权重量化为 int4，但在推理时动态恢复为 fp16。

量化模型

from transformers import AutoModelForCausalLM, AutoTokenizer, GPTQConfig
import torch

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")

# 配置 GPTQ 量化
gptq_config = GPTQConfig(
    bits=4,                    # 量化位数: 2, 3, 4, 8
    dataset="c4",              # 校准数据集
    tokenizer=tokenizer,       # 分词器
    group_size=128,            # 分组大小
    damp_percent=0.1,          # 阻尼系数
    desc_act=False             # 是否按激活值降序量化
)

# 量化模型
model = AutoModelForCausalLM.from_pretrained(
    "facebook/opt-125m",
    device_map="auto",
    quantization_config=gptq_config
)

# 保存量化后的模型
model.save_pretrained("./opt-125m-gptq")
tokenizer.save_pretrained("./opt-125m-gptq")

GPTQ 关键参数说明：

• bits：量化位数。4-bit 是最常用的选择，平衡了压缩率和精度。

• dataset：用于校准的数据集。推荐使用 "c4" 或 "wikitext2"，也可以传入自定义文本列表。

• group_size：分组量化的大小。-1 表示按列量化，128 是推荐值。

• desc_act：是否按激活值大小降序量化。设为 False 可以显著加速推理，但困惑度可能略有上升。

加载 GPTQ 模型

from transformers import AutoModelForCausalLM, GPTQConfig
import torch

# 基本加载
model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Llama-2-7B-GPTQ",
    device_map="auto"
)

# 使用 Marlin 后端加速（仅支持 A100 等 Ampere 架构 GPU）
model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Llama-2-7B-GPTQ",
    device_map="auto",
    quantization_config=GPTQConfig(bits=4, backend="marlin")
)

使用 ExLlama 后端

from transformers import AutoModelForCausalLM, GPTQConfig

# ExLlamaV2 内核
gptq_config = GPTQConfig(
    bits=4,
    exllama_config={"version": 2}
)

model = AutoModelForCausalLM.from_pretrained(
    "TheBloke/Llama-2-7B-GPTQ",
    device_map="auto",
    quantization_config=gptq_config
)

Quanto 量化

Quanto 是一个设备无关的量化库，支持 CPU、GPU 和 Apple Silicon（MPS），适合跨平台部署。

from transformers import AutoModelForCausalLM, QuantoConfig
import torch

# Quanto 配置
quantization_config = QuantoConfig(
    weights="int8",            # 权重量化: "int2", "int4", "int8", "float8"
    activations=None           # 激活量化: None, "int8", "float8"
)

model = AutoModelForCausalLM.from_pretrained(
    "facebook/opt-125m",
    quantization_config=quantization_config,
    device_map="auto"
)

Quanto 的优势在于其通用性，可以在 Apple Silicon 上高效运行。

AQLM 量化

AQLM（Additive Quantization for LLMs）是一种极致压缩方法，可以将模型压缩到 2-bit。

from transformers import AutoModelForCausalLM, AqlmConfig

quantization_config = AqlmConfig(
    in_group_size=8,           # 输入维度分组大小
    out_group_size=1,          # 输出维度分组大小
    num_codebooks=1,           # 码本数量
    nbits_per_codebook=16      # 每个码本的位数
)

model = AutoModelForCausalLM.from_pretrained(
    "ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf",
    quantization_config=quantization_config,
    device_map="auto"
)

KV Cache 量化

在自回归生成任务中，KV Cache 是一个重要的内存瓶颈。对于一个 7B 模型，10,000 个 token 的 KV Cache 约需要 5GB 显存。KV Cache 量化可以显著减少这部分内存占用。

什么是 KV Cache？

在 Transformer 的自回归生成中，每生成一个新 token 都需要重新计算之前所有 token 的 Key 和 Value。KV Cache 通过缓存这些计算结果避免重复计算：

┌─────────────────────────────────────────────────────────────────┐
│                    KV Cache 工作原理                             │
├─────────────────────────────────────────────────────────────────┤
│                                                                  │
│   生成第 K 个 token 时：                                         │
│   ┌──────────────────────────────────────────┐                  │
│   │  Q_K × [K_1, K_2, ..., K_{K-1}]^T       │                  │
│   │              ↓                           │                  │
│   │  Attention Scores                        │                  │
│   │              ↓                           │                  │
│   │  Weighted Sum of [V_1, V_2, ..., V_{K-1}]│                  │
│   └──────────────────────────────────────────┘                  │
│                                                                  │
│   生成第 K+1 个 token 时：                                       │
│   ┌──────────────────────────────────────────┐                  │
│   │  直接从缓存获取 K_1...K_{K-1}, V_1...V_{K-1}│                 │
│   │              ↓                           │                  │
│   │  只需计算 K_K, V_K 并添加到缓存           │                  │
│   │              ↓                           │                  │
│   │  大幅减少重复计算                         │                  │
│   └──────────────────────────────────────────┘                  │
│                                                                  │
│   KV Cache 内存占用 ≈ 2 × 2 × layers × heads × dim × seq_len    │
│                                                                  │
└─────────────────────────────────────────────────────────────────┘

使用 KV Cache 量化

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# 安装依赖: pip install quanto

# 加载模型
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-chat-hf",
    torch_dtype=torch.float16,
    device_map="cuda:0"
)

# 输入
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

# 使用 int4 量化 KV Cache
outputs = model.generate(
    **inputs,
    do_sample=False,
    max_new_tokens=50,
    cache_implementation="quantized",
    cache_config={
        "backend": "quanto",
        "nbits": 4
    }
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

KV Cache 量化配置详解

# 详细的 KV Cache 量化配置
from transformers import CacheConfig

cache_config = CacheConfig(
    backend="quanto",          # 后端: "quanto" 或 "hqq"
    nbits=4,                   # 量化位数: 2, 4, 8
    group_size=64,             # 分组大小
    residual_length=128        # 残差缓存长度
)

outputs = model.generate(
    **inputs,
    cache_implementation="quantized",
    cache_config=cache_config
)

KV Cache 量化参数说明：

• nbits：量化位数。int4 几乎不影响模型质量，int2 会有些许性能下降。

• group_size：分组量化的大小。较小的值精度更高，但开销更大。

• residual_length：残差缓存长度。保留最近 N 个 token 的原始精度，平衡精度和内存。

性能对比

配置	最大序列长度	显存占用	相对速度
fp16 KV Cache	40k tokens	高	基准
int4 KV Cache	128k tokens	约 40%	约 80-100%
int2 KV Cache	200k+ tokens	约 25%	约 60-80%

结合 Flash Attention，在 80GB A100 上，int4 量化 KV Cache 可支持最多 128k tokens 的上下文。

量化方法选择指南

场景	推荐方法	说明
快速原型	bitsandbytes 4-bit	最简单，无需校准
高性能推理	AWQ + ExLlamaV2	最快推理速度
平衡精度与速度	GPTQ + Marlin	灵活配置
跨平台部署	Quanto	支持 CPU/GPU/MPS
长上下文生成	KV Cache 量化	扩展上下文长度
极致压缩	AQLM 2-bit	最小模型体积

推理优化

半精度推理

使用 FP16 或 BF16 可以减少显存占用并加速推理：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# FP16 加载
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.float16,
    device_map="auto"
)

# BF16 加载（需要 Ampere 及更新架构的 GPU，如 A100、RTX 30/40 系列）
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

Flash Attention 2

Flash Attention 通过优化内存访问模式大幅提升注意力计算效率，特别适合长序列场景。

from transformers import AutoModelForCausalLM
import torch

# 启用 Flash Attention 2
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
    device_map="auto"
)

# 推理时自动使用 Flash Attention
inputs = tokenizer("Long text here...", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=100)

Flash Attention 的优势：

• 内存占用从 $O(n^2)$ 降低到 $O(n)$
• 支持更长的上下文长度
• 训练和推理都能受益

使用 BetterTransformer

BetterTransformer 可以自动使用优化的注意力实现：

from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")

# 启用 BetterTransformer
model = model.to_bettertransformer()

# 推理
inputs = tokenizer("Hello world", return_tensors="pt")
outputs = model(**inputs)

KV Cache 优化

对于自回归生成任务，启用 KV Cache 可以显著减少重复计算：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")

inputs = tokenizer("The capital of France is", return_tensors="pt")

# 启用 KV Cache
with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=50,
        use_cache=True,           # 启用 KV Cache
        pad_token_id=tokenizer.eos_token_id
    )

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

批处理推理

批处理可以充分利用 GPU 并行计算能力：

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english"
).to("cuda")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

texts = [
    "I love this product!",
    "This is terrible.",
    "It's okay, nothing special.",
    "Highly recommended!"
]

# 批量处理
inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt").to("cuda")

with torch.no_grad():
    outputs = model(**inputs)
    predictions = torch.argmax(outputs.logits, dim=-1)

for text, pred in zip(texts, predictions):
    label = "POSITIVE" if pred == 1 else "NEGATIVE"
    print(f"{text}: {label}")

模型服务化

Flask API 服务

from flask import Flask, request, jsonify
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

app = Flask(__name__)

# 加载模型
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english"
).to("cuda")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
model.eval()

@app.route("/predict", methods=["POST"])
def predict():
    data = request.json
    texts = data.get("texts", [])
    
    if not texts:
        return jsonify({"error": "No texts provided"}), 400
    
    # 批量推理
    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt").to("cuda")
    
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1)
        predictions = torch.argmax(probs, dim=-1)
    
    results = []
    for text, pred, prob in zip(texts, predictions, probs):
        results.append({
            "text": text,
            "label": "POSITIVE" if pred.item() == 1 else "NEGATIVE",
            "confidence": prob[pred.item()].item()
        })
    
    return jsonify({"results": results})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5000)

FastAPI 服务

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

app = FastAPI()

# 加载模型
model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english"
).to("cuda")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
model.eval()

class PredictRequest(BaseModel):
    texts: List[str]

class PredictionResult(BaseModel):
    text: str
    label: str
    confidence: float

class PredictResponse(BaseModel):
    results: List[PredictionResult]

@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    if not request.texts:
        raise HTTPException(status_code=400, detail="No texts provided")
    
    inputs = tokenizer(
        request.texts,
        padding=True,
        truncation=True,
        return_tensors="pt"
    ).to("cuda")
    
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1)
        predictions = torch.argmax(probs, dim=-1)
    
    results = []
    for text, pred, prob in zip(request.texts, predictions, probs):
        results.append(PredictionResult(
            text=text,
            label="POSITIVE" if pred.item() == 1 else "NEGATIVE",
            confidence=prob[pred.item()].item()
        ))
    
    return PredictResponse(results=results)

使用 Hugging Face Inference Endpoints

import requests

# Hugging Face Inference Endpoints API
API_URL = "https://api-inference.huggingface.co/models/distilbert-base-uncased-finetuned-sst-2-english"
headers = {"Authorization": f"Bearer {YOUR_HF_TOKEN}"}

def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()

# 推理
output = query({"inputs": "I love this movie!"})
print(output)
# [{'label': 'POSITIVE', 'score': 0.9998}]

模型打包与分发

保存完整模型

from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# 保存到本地
save_path = "./my_model"
model.save_pretrained(save_path)
tokenizer.save_pretrained(save_path)

# 打包为压缩文件
import shutil
shutil.make_archive("my_model", "zip", save_path)

上传到 Hugging Face Hub

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from huggingface_hub import login

# 登录
login()

# 加载模型
model = AutoModelForSequenceClassification.from_pretrained("./my_finetuned_model")
tokenizer = AutoTokenizer.from_pretrained("./my_finetuned_model")

# 推送到 Hub
model.push_to_hub("my-username/my-model")
tokenizer.push_to_hub("my-username/my-model")

使用 safetensors 格式

safetensors 是一种安全、高效的模型存储格式：

from transformers import AutoModelForSequenceClassification

# 加载模型并保存为 safetensors
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased")
model.save_pretrained("./model", safe_serialization=True)

# 加载 safetensors 格式的模型
model = AutoModelForSequenceClassification.from_pretrained(
    "./model",
    use_safetensors=True
)

性能基准测试

测量推理延迟

import time
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased-finetuned-sst-2-english"
).to("cuda")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
model.eval()

# 预热
inputs = tokenizer("Warmup text", return_tensors="pt").to("cuda")
with torch.no_grad():
    for _ in range(10):
        _ = model(**inputs)

# 测量
texts = ["This is a test sentence."] * 100
inputs = tokenizer(texts, padding=True, return_tensors="pt").to("cuda")

torch.cuda.synchronize()
start = time.time()

with torch.no_grad():
    for _ in range(10):
        _ = model(**inputs)

torch.cuda.synchronize()
end = time.time()

total_samples = len(texts) * 10
total_time = end - start
latency = total_time / total_samples * 1000  # ms

print(f"平均延迟: {latency:.2f} ms")
print(f"吞吐量: {total_samples / total_time:.2f} samples/s")

测量显存占用

import torch
from transformers import AutoModelForCausalLM

def measure_memory(model_id, dtype=torch.float32):
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=dtype,
        device_map="cuda"
    )
    
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.max_memory_reserved() / 1024**3
    
    print(f"模型: {model_id}")
    print(f"数据类型: {dtype}")
    print(f"已分配显存: {allocated:.2f} GB")
    print(f"预留显存: {reserved:.2f} GB")
    
    del model
    torch.cuda.empty_cache()

measure_memory("gpt2", torch.float32)
measure_memory("gpt2", torch.float16)

部署最佳实践

1. 模型选择

根据实际需求选择合适的模型规模：

场景	推荐模型	说明
实时交互	DistilBERT, GPT-2 small	低延迟优先
批量处理	BERT-base, LLaMA-7B	吞吐量优先
高精度需求	RoBERTa-large, LLaMA-70B	准确性优先

2. 资源管理

# 限制 GPU 显存使用
import torch
torch.cuda.set_per_process_memory_fraction(0.8, 0)  # 使用 80% 显存

# 自动设备选择
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

3. 错误处理

import torch
from transformers import AutoModelForCausalLM

def load_model_safely(model_id, fallback_id=None):
    try:
        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        return model
    except torch.cuda.OutOfMemoryError:
        print(f"显存不足，尝试使用 CPU 加载 {model_id}")
        return AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu")
    except Exception as e:
        if fallback_id:
            print(f"加载 {model_id} 失败: {e}，使用备用模型 {fallback_id}")
            return AutoModelForCausalLM.from_pretrained(fallback_id)
        raise

4. 健康检查

from fastapi import FastAPI
import torch

app = FastAPI()

@app.get("/health")
async def health_check():
    checks = {
        "status": "healthy",
        "cuda_available": torch.cuda.is_available(),
        "cuda_device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0
    }
    
    if torch.cuda.is_available():
        checks["cuda_memory_allocated"] = f"{torch.cuda.memory_allocated() / 1024**3:.2f} GB"
        checks["cuda_memory_reserved"] = f"{torch.cuda.memory_reserved() / 1024**3:.2f} GB"
    
    return checks

小结

模型部署是将训练好的模型投入生产的关键环节。本章涵盖了：

1. ONNX 导出：跨平台部署的标准格式
2. 模型量化：bitsandbytes、AWQ、GPTQ、Quanto 等多种量化方法
3. KV Cache 量化：扩展长上下文生成能力
4. 推理优化：半精度、Flash Attention、批处理等技术
5. 模型服务化：Flask、FastAPI 等 API 服务封装
6. 性能测试：延迟和吞吐量的测量方法

在实际部署中，需要根据具体场景在延迟、吞吐量、准确性之间权衡，选择合适的优化策略。

参考资源

• ONNX 官方文档
• Hugging Face Optimum
• ONNX Runtime
• bitsandbytes
• Flash Attention
• KV Cache 量化博客
• AWQ 量化指南
• GPTQ 量化指南