从评估走到优化:AI性能基准测试笔记

AI性能基准测试笔记我没按教科书顺序做。

先解决眼前的阻塞,再回头补原理。

基础设施和准备

测什么之前,得先把环境理干净。测试结果怪,先查环境——模型背锅的情况我见识过太多次。

最基础的几点:

# GPU 预热,避免冷启动影响
python -c "import torch; torch.cuda.empty_cache(); _ = torch.randn(1000, 1000).cuda() @ torch.randn(1000, 1000).cuda()"

# 设置 CUDA 环境变量,避免自动寻找优化带来的不一致
export TORCH_CUDA_ARCH_LIST="8.6;9.0"
export CUDA_VISIBLE_DEVICES=0

我吃过两次亏:一次是在多卡服务器上跑了脚本,结果第一个模型占了一卡,第二个模型自动占了另一卡,但两卡的显存状态不一致,导致延迟波动很大;另一次是忘了预热,直接跑推理,第一个样本的延迟是后面的三倍,差点误判模型性能不行。

环境的一致性比想象中重要,尤其是做对比测试时。

测什么和怎么测

光跑一次推理看不出什么,真正有意义的测试通常是这几个维度:

推理延迟

最直接的就是单次推理延迟,包括prefill和decode两个阶段。

import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def measure_latency(model, tokenizer, prompt, n_runs=10, warmup_runs=3):
    # 预热
    for _ in range(warmup_runs):
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            _ = model.generate(**inputs, max_new_tokens=50)

    # 实际测试
    latencies = []
    for _ in range(n_runs):
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        torch.cuda.synchronize()
        start = time.perf_counter()
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=100)
        torch.cuda.synchronize()
        end = time.perf_counter()
        latencies.append((end - start) * 1000)  # ms

    return {
        "mean": sum(latencies) / len(latencies),
        "p50": sorted(latencies)[len(latencies)//2],
        "p95": sorted(latencies)[int(len(latencies)*0.95)],
        "p99": sorted(latencies)[int(len(latencies)*0.99)],
    }

model = AutoModelForCausalLM.from_pretrained("...")
tokenizer = AutoTokenizer.from_pretrained("...")
results = measure_latency(model, tokenizer, "讲个笑话", n_runs=20)

注意这里用了p95和p99,单纯看平均延迟会掩盖很多问题。有个模型平均延迟看着不错,但p99是平均的两倍,说明有偶发的推理卡顿,这在实时场景下很致命。

吞吐量

单次推理延迟是低并发场景下的指标,真正要看量还得测吞吐量。

from concurrent.futures import ThreadPoolExecutor
import numpy as np

def measure_throughput(model, tokenizer, prompts, batch_size=4, n_threads=2, n_runs=5):
    def run_batch(batch):
        inputs = tokenizer(batch, padding=True, truncation=True, return_tensors="pt").to(model.device)
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=50)
        return [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]

    def split_into_batches(items, size):
        for i in range(0, len(items), size):
            yield items[i:i + size]

    throughput_list = []
    for _ in range(n_runs):
        start = time.time()
        with ThreadPoolExecutor(max_workers=n_threads) as executor:
            results = list(executor.map(run_batch, split_into_batches(prompts, batch_size)))
        end = time.time()
        throughput = len(prompts) / (end - start)
        throughput_list.append(throughput)

    return {
        "mean": np.mean(throughput_list),
        "std": np.std(throughput_list),
    }

prompts = ["介绍一下你自己"] * 100
throughput = measure_throughput(model, tokenizer, prompts, batch_size=8, n_threads=4)

这里的关键是找到最佳批处理大小。太小的话GPU利用率不够,太大的话内存装不下或者延迟暴涨。我当时有个模型,batch_size=4时吞吐量不错,但一旦超过8,延迟就开始指数级上升,因为显存碎片化导致频繁换页。

内存占用

内存占用分两部分:模型加载时的显存占用和推理时的峰值占用。

import torch

def measure_memory(model, tokenizer, prompt, max_new_tokens=100):
    torch.cuda.reset_peak_memory_stats()
    torch.cuda.empty_cache()

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # 模型加载后的显存占用
    model_memory = torch.cuda.memory_allocated() / 1024**3  # GB

    # 推理时的峰值占用
    with torch.no_grad():
        _ = model.generate(**inputs, max_new_tokens=max_new_tokens)
    peak_memory = torch.cuda.max_memory_allocated() / 1024**3  # GB

    return {
        "model_memory": model_memory,
        "peak_memory": peak_memory,
        "inference_memory": peak_memory - model_memory,
    }

memory = measure_memory(model, tokenizer, "写一段Python代码")

量化模型要特别注意这个指标。4bit量化后模型占用是原本的1/4左右,但推理峰值内存可能只降了30%,因为attention cache占的显存没变。

踩过的坑和解决方案

批处理效率陷阱

一开始我以为批量推理吞吐量一定比单个快,结果有个模型反而更慢。后来发现那个模型的实现没做好批处理attention,内部的loop反而抵消了并行优势。

解决方案很简单:在评估吞吐量时,覆盖不同batch_size,画出曲线。

batch_sizes = [1, 2, 4, 8, 16, 32]
throughput_curve = []

for bs in batch_sizes:
    result = measure_throughput(model, tokenizer, prompts, batch_size=bs, n_runs=3)
    throughput_curve.append(result["mean"])

# 找到throughput不再明显增长的临界点
optimal_batch = batch_sizes[np.argmax(throughput_curve)]

有个7B参数的模型,batch_size=1时throughput是1.2 tokens/s,batch_size=4时跳到3.8 tokens/s,但batch_size=8时反而掉到3.5 tokens/s,因为显存开始换页。最后在真实服务里选batch_size=4作为默认配置。

把这三个实测点画成曲线,能直观看到吞吐并非随 batch 单调上升:

7B 模型在不同 batch_size 下的推理吞吐量曲线,batch_size=4 为峰值

曲线在 batch_size=4 达到峰值后回落,说明继续加大 batch 只会因显存换页拖慢整体吞吐。

量化精度评估

为了省显存和加快推理,很多人直接上4bit量化,但很少人认真评估精度损失。

我做的一件事是:用数据集跑一套标准评估(比如MMLU、C-Eval),对比量化前后的准确率差异。

def evaluate_model_accuracy(model, tokenizer, eval_dataset):
    correct = 0
    total = len(eval_dataset)

    for item in eval_dataset:
        prompt = item["question"] + "\n" + item["choices_str"]
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=10)

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

        if item["answer"] in generated:
            correct += 1

    return correct / total

# 评估原始模型和量化模型
accuracy_original = evaluate_model_accuracy(original_model, tokenizer, eval_dataset)
accuracy_quantized = evaluate_model_accuracy(quantized_model, tokenizer, eval_dataset)

accuracy_drop = accuracy_original - accuracy_quantized

有个CodeLlama的7B模型,4bit量化后推理速度提升60%,但代码生成任务的准确率掉了12%,权衡之后还是用了8bit量化,速度提升40%,准确率只掉3%。

CPU vs GPU的推理效率

以为GPU推理一定比CPU快,结果在低并发场景下,CPU推理反而更划算。GPU有启动开销和显存传输成本,如果你每秒只有几个请求,GPU的优势还没发挥出来就结束了。

我做了个测试:在同一个模型上,分别用CPU和GPU推理不同并发量的请求。

import matplotlib.pyplot as plt

def test_cpu_vs_gpu_concurrency(concurrency_levels):
    results = {"cpu": [], "gpu": []}

    for concurrency in concurrency_levels:
        # GPU测试
        gpu_model = AutoModelForCausalLM.from_pretrained(...).cuda()
        gpu_result = measure_throughput(gpu_model, tokenizer, prompts, batch_size=concurrency)
        results["gpu"].append(gpu_result["mean"])

        # CPU测试
        cpu_model = AutoModelForCausalLM.from_pretrained(...).cpu()
        cpu_result = measure_throughput(cpu_model, tokenizer, prompts, batch_size=concurrency)
        results["cpu"].append(cpu_result["mean"])

    return results

concurrency_levels = [1, 2, 4, 8, 16, 32]
results = test_cpu_vs_gpu_concurrency(concurrency_levels)

plt.figure(figsize=(10, 6))
plt.plot(concurrency_levels, results["cpu"], label="CPU", marker='o')
plt.plot(concurrency_levels, results["gpu"], label="GPU", marker='s')
plt.xlabel("Concurrency")
plt.ylabel("Throughput (tokens/s)")
plt.legend()
plt.savefig("cpu_vs_gpu_throughput.png")

结果很清楚:并发低于4时,CPU更快;超过8后,GPU碾压。根据这个测试结果,我在部署时做了个策略:低并发时段用CPU推理,高并发时段自动切GPU。

优化策略和效果

优化不是万能药,但要找准瓶颈。

模型蒸馏

蒸馏是最常见的方法,用一个大的Teacher模型教一个小的Student模型。效果好的前提是Teacher和Student的任务场景一致,否则Student学不到重点。

我当时试过一次,拿通用大模型当Teacher去蒸馏一个专门用于代码生成的Student模型,结果Student模型在通用任务上表现还行,但代码生成能力反而不如原始小模型。后来换了个代码生成模型当Teacher,效果才上来。

模型剪枝

剪枝是去掉不重要的权重参数。理论上能减少模型大小和计算量,但实际效果参差不齐。

我试过对一个7B模型做50%非结构化剪枝,模型大小确实减了一半,但推理速度没快多少,因为现在的推理库对非结构化剪枝优化有限。后来改做结构化剪枝,直接砍掉整个attention头,效果明显很多。

量化策略

前面提过,4bit量化可能太激进。现在常用的策略是:关键层保持高精度,非关键层可以降精度。

from transformers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",  # normal float 4
)

model = AutoModelForCausalLM.from_pretrained(
    "model-name",
    quantization_config=quantization_config,
    device_map="auto",
)

有个经验是:attention的QKVO层可以降到4bit,但最后的输出层最好保持fp16,否则输出质量会明显下滑。

推理框架选择

PyTorch原生推理还行,但专业的推理库能榨出更多性能。常见的有ONNX Runtime、TensorRT、vLLM等。

我对比过vLLM和原生PyTorch的throughput差异,在同个7B模型上,vLLM比原生PyTorch快2-3倍,尤其是在高并发场景下。

# vLLM部署示例
python -m vllm.entrypoints.api_server \
    --model meta-llama/Llama-2-7b-chat-hf \
    --quantization awq \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.9

但vLLM也不是银弹,它在低延迟场景下的表现不如TensorRT,而且对模型的兼容性要求高。如果你的模型不在支持列表里,vLLM直接用不了。

持续监控和告警

性能优化不是一次性的,模型状态、用户请求特征、硬件负载都会变化,需要持续监控。

我通常会跟踪这几个指标:

  • 请求延迟的p50、p95、p99
  • GPU利用率
  • 显存占用
  • 错误率和超时率
import prometheus_client

latency_histogram = prometheus_client.Histogram('model_inference_latency_seconds', 'Model inference latency')
throughput_gauge = prometheus_client.Gauge('model_throughput_tokens_per_second', 'Model throughput')

@latency_histogram.time()
def inference_with_metrics(prompt):
    start = time.time()
    result = model.generate(**prompt)
    latency = time.time() - start
    throughput_gauge.set(len(result) / latency)
    return result

告警规则设在p95延迟超过阈值、GPU利用率长期过低或过高、显存占用超过安全边界等。有一次GPU利用率突然掉到10%,仔细排查发现是推理脚本在等待数据库响应,后端瓶颈反而卡住了推理吞吐。

最后一点思考

性能基准测试是个不断逼近真相的过程,但没有永远正确的答案。你今天的测试结果,到明天可能因为模型微调、硬件升级、用户行为变化而失效。

重要的是建立起一套系统的测试方法,环境一变就能快速复现和评估。跟写单元测试一样:测性能也是为了下次改模型、改配置时有对照,别指望测一次就管半年。

最后提醒一句:性能优化的尽头是成本效益分析。模型快50%但成本翻倍,用户延迟感知不到差异,这种优化就意义不大。真正有价值的是那些让边际成本足够低的优化,而不是单纯追求数字好看。

版权声明: 本文首发于 指尖魔法屋-从评估走到优化:AI性能基准测试笔记https://blog.thinkmoon.cn/post/202-ai-benchmark-evaluation-optimization/) 转载或引用必须申明原指尖魔法屋来源及源地址!