AI模型部署：推理不够用了之后

作者: 醉月思发布日期: 2024-03-11 分类: 技术实践

先说个最实在的：项目背景是个多模态问答系统，推理环境是 4 张 A100 80GB，PyTorch 2.1.1，CUDA 12.1。

QPS 要求不高，但延迟不能超过 2 秒，单次推理显存占用要控制在 50GB 以内（因为还有个图像预处理模型在跑）。

一开始用的是 TorchServe

当时选 TorchServe 的原因也很简单：官方出品、PyTorch 原生支持、社区文档多。本地搭个 demo 确实快：

torchserve --start --ncs --model-store model_store \
  --models my-model=BERT/model.mar \
  --ts-config config.properties

配置文件也很直观：

inference_address=http://0.0.0.0:8080
management_address=http://0.0.0.0:8081
number_of_netty_threads=32
job_queue_size=10
model_store=model_store

但真上到生产环境就出问题了。首先是并发一上来，延迟就开始飙升。观察了一下 nvidia-smi，GPU 利用率在 30% 到 70% 之间剧烈波动，说明调度上出了问题。

TorchServe 的批处理机制默认是开启的，但它的实现逻辑是：要么等批次凑满（默认 batch_size=8），要么等超时时间到了（default max_batch_delay=100ms）。如果请求稀疏，经常是等半天凑不齐一个 batch，反而把延迟拖大了。

当时改了一下配置：

min_workers=4
max_workers=8
max_batch_delay=50
batch_size=16
response_timeout=120

结果另一个问题又出来了：某些长文本推理请求会卡住整个 batch，导致后续短请求也被拖慢。TorchServe 没有请求优先级机制，所有请求都平等排队，这在生产环境里其实不太现实。

还有一个实际遇到的坑：TorchServe 的模型加载是懒加载的，也就是说第一个请求来时才开始初始化模型。这导致第一次调用延迟特别高，实测能达到 3-4 秒。虽然它有预热机制，但预热请求也只是走一遍推理流程，不会真正预热所有可能的 batch size。

后来加了主动预热：

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def warmup():
    model = AutoModelForCausalLM.from_pretrained("model_path", device_map="cuda")
    tokenizer = AutoTokenizer.from_pretrained("model_path")
    dummy_input = tokenizer(["Hello world"] * 16, return_tensors="pt", padding=True)
    with torch.no_grad():
        _ = model.generate(**dummy_input, max_length=10)

但这个预热脚本要放在模型 archive 里，TorchServe 加载时自动执行，每次模型更新都得重新打包。整个过程有点繁琐，不像一个生产级别的方案。

后来切到 vLLM

TorchServe 跑了一个月，性能瓶颈越来越明显，就开始调研替代方案。vLLM 当时在圈子里名声不错，号称能大幅提升 LLM 推理吞吐量。

vLLM 的核心思路是 PagedAttention，它把 KV cache 分成固定的块（block），像操作系统分页一样动态管理。这对长序列和变长序列的推理特别友好，不需要像传统实现那样为每个序列预留固定大小的连续显存空间。

部署 vLLM 也简单：

python -m vllm.entrypoints.openai.api_server \
  --model /path/to/model \
  --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.9 \
  --max-model-len 4096 \
  --host 0.0.0.0 \
  --port 8000

但实际用起来也踩了不少坑。

第一个问题是显存利用率。vLLM 默认 gpu-memory-utilization=0.9，理论上是充分利用 GPU，但实际上它还预留了一部分空间给 KV cache 动态分配。如果模型本身就很大（比如 70B 模型），这个预留空间就不够用了，推理到一半就会报 OOM。

解决方法是手动调低参数：

--gpu-memory-utilization 0.85 \
--max-model-len 3072 \
--swap-space 4  # 使用 4GB 的 swap 空间

但 swap 空间会带来性能损失，实测延迟会增加 15-20%，所以只能在低峰期使用。

第二个问题是请求排队。vLLM 默认使用 Scheduler 管理请求队列，但它没有优先级机制，所有请求都是 FIFO。如果某个长请求在排队，后续的短请求也必须等。

后来改用了 continuous batching（vLLM 默认开启），它允许在推理过程中动态加入新请求，但实际效果取决于请求模式。如果大部分请求都是长文本，continuous batching 的优势就不明显。

还有一个问题是多实例部署。vLLM 的 API server 默认是单实例的，要扩展得手动起多个进程：

CUDA_VISIBLE_DEVICES=0,1 python -m vllm.entrypoints.openai.api_server --model model --tensor-parallel-size 2 --port 8000 &
CUDA_VISIBLE_DEVICES=2,3 python -m vllm.entrypoints.openai.api_server --model model --tensor-parallel-size 2 --port 8001 &

然后前面再套一层 Nginx 做负载均衡。但这就回到了原来的问题：多个实例之间没有共享状态，每个实例都要加载一份模型，显存开销翻倍。

vLLM 确实在吞吐量上比 TorchServe 好很多，同等 QPS 下延迟降低了 40-50%，但在资源利用率和请求调度上仍然不够灵活。

最后搞了 TensorRT

把 TorchServe 和 vLLM 都试了一圈后，开始考虑更激进的优化方案：TensorRT。

TensorRT 的优势是它能对模型进行深度优化，包括层融合、精度优化、内核自动调优等。但它的学习曲线也最陡，而且不是所有模型都能直接转换。

以 BERT 为例，转换过程大概是这样：

import torch
from transformers import BertTokenizer, BertModel
from torch2trt import TRTModule

# 加载原始模型
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased').cuda().eval()

# 创建 TensorRT 引擎
x = torch.ones(1, 128).cuda()  # dummy input
model_trt = TRTModule()
model_trt = torch2trt.torch2trt(model, [x], fp16_mode=True)

# 保存引擎
torch.save(model_trt.state_dict(), 'bert_fp16.trt')

但实际转换过程中遇到两个问题：

第一，TensorRT 对动态输入的支持有限。BERT 的序列长度是可变的，但 TensorRT 要求在转换时固定输入形状。虽然它后来加了 dynamic shape 支持，但使用起来比较麻烦：

input_shapes = {
    "input_ids": [[1, 32], [1, 128], [1, 512]],  # min, opt, max
    "attention_mask": [[1, 32], [1, 128], [1, 512]],
}

而且 dynamic shape 的推理性能通常不如固定形状，尤其是在 A100 这种对形状敏感的 GPU 上。

第二，TensorRT 对某些自定义算子支持不好。如果你的模型里有特殊的层（比如自定义的 attention 机制），转换过程很容易失败，或者转换后推理结果不对。这时候要么重写算子用 TensorRT 支持的实现，要么就放弃这部分优化。

我们后来采取了一个折中方案：把模型中结构比较规整的部分用 TensorRT 优化，剩下那部分仍用 PyTorch 推理，推理时在两者之间切换：

def hybrid_inference(input_ids):
    # 前半部分用 TensorRT
    with torch.no_grad():
        embeddings = model_trt.embeddings(input_ids)
        encoder_outputs = model_trt.encoder(embeddings)

    # 后半部分用 PyTorch
    with torch.no_grad():
        logits = model_torch.head(encoder_outputs)

    return logits

这个方案的好处是兼顾了优化和灵活性，但实现起来比较复杂，需要仔细处理两个框架之间的数据格式和设备切换。

GPU 调度这块儿

折腾完推理优化，GPU 资源调度又成了下一个问题。

一开始用的是 Kubernetes 的 GPU device plugin，它能让 Pod 申请整张 GPU 卡。但这个方案有两个问题：

第一，GPU 利用率低。如果模型只需要半张卡的资源，剩下的半张就浪费了。虽然可以用 MPS (Multi-Process Service) 多进程共享一张卡，但 MPS 对内存隔离和性能隔离的支持有限，多个进程互相干扰的概率很高。

第二，无法做细粒度的资源限制。Kubernetes 只能控制 Pod 能用几张 GPU，不能控制显存、计算单元的具体配额。

后来试了 NVIDIA 的 MIG (Multi-Instance GPU)，它把一张 A100 切成多个实例，每个实例有独立的显存和计算核心：

# 创建 MIG 实例
nvidia-smi mig -cgi 1g.5gb -C

# 查看实例
nvidia-smi mig -lgi

但 MIG 也有局限：一是只支持部分 GPU 型号（A100、A30、H100 等），二是切分后的实例性能不如完整 GPU，三是配置起来比较复杂，需要调整容器运行时和调度器。

最后采用了分层调度策略：核心推理服务独占 GPU，辅助服务（比如预处理、后处理）使用 CPU 或者共享 GPU。具体来说，把 4 张 A100 分成两组，前 3 张跑主推理模型，第 4 张跑预处理模型和批处理逻辑。

这样既保证了核心服务的性能，又充分利用了资源。但这个策略是针对当前场景的，如果后续业务变了，可能还要重新调整。

批量推理的坑

批量推理听起来很简单：把多个请求打包成 batch，一次性喂给模型。但实际操作起来问题很多。

首先是序列长度对齐。如果 batch 里有长有短，短序列需要补零，但补零的位置不需要计算 attention，这会浪费计算资源。vLLM 的 PagedAttention 部分解决了这个问题，但它对短序列的处理仍然不够高效。

其次是显存碎片。频繁的 batch 分配和释放会导致显存碎片化，长期运行后可能遇到"明明总显存够用，但就是分配不出大块连续空间"的问题。解决方法是定期重启服务，或者实现一个简单的显存池：

class KVCachePool:
    def __init__(self, capacity, block_size):
        self.capacity = capacity
        self.block_size = block_size
        self.blocks = [None] * (capacity // block_size)
        self.used = [False] * len(self.blocks)

    def allocate(self, n_blocks):
        for i in range(len(self.blocks) - n_blocks + 1):
            if not any(self.used[i:i+n_blocks]):
                self.used[i:i+n_blocks] = [True] * n_blocks
                return i
        raise MemoryError("Not enough contiguous blocks")

    def free(self, start_idx, n_blocks):
        self.used[start_idx:start_idx+n_blocks] = [False] * n_blocks

这个池化方案是粗粒度的，不能完全避免碎片，但至少能缓解问题。

还有一个实际遇到的坑：错误处理。批量推理时，如果 batch 中某个请求失败，整个 batch 都要重试，这会浪费计算资源。后来改成单个请求隔离：

def batch_inference_with_fallback(requests):
    results = []
    failed_indices = []

    try:
        results = model.inference(requests)
    except Exception as e:
        # 如果 batch 失败，逐个重试
        for i, req in enumerate(requests):
            try:
                results[i] = model.inference([req])[0]
            except Exception as e:
                results[i] = fallback_response(req)
                failed_indices.append(i)

    return results, failed_indices

这个方案增加了一些延迟，但至少不会因为一个坏请求影响整个 batch。

监控这块不能少

折腾了这么多，最后发现如果监控不到位，出了问题很难定位。

最开始只监控了基础的 QPS、延迟、GPU 利用率，但这些指标有时候会互相误导。比如 GPU 利用率高并不一定意味着服务状态好，可能只是计算密集型请求多，但吞吐量却不高。

后来加了更多细粒度的指标：

KV cache 命中率（衡量缓存策略效果）
显存碎片率（衡量内存分配健康度）
批次大小分布（衡量请求模式和调度效果）
每层推理时间（定位性能瓶颈具体在哪）

监控工具一开始用的是 Prometheus + Grafana，但后来发现它的实时性不够好，出了问题要等 10-15 秒才能看到指标变化。对于推理服务来说，这个延迟有点长，问题可能已经扩散了。

现在用的是 NVIDIA 的 Nsight Systems，它能记录每个 kernel 的执行时间和资源占用，虽然分析起来比较麻烦，但能精确到微秒级别：

nsys profile --trace=cuda,nvtx --stats=true --output=report.qdrep \
  python inference.py

生成的报告可以导入 Nsight Systems 查看，能看到每个操作的 GPU 时间线和 CPU 时间线，非常直观。

但这个工具也有缺点：一是会带来明显的性能开销（实测 QPS 下降 30-40%），二是只能用于 profiling，不能长期监控。所以现在是定期用 Nsight 做性能分析，日常监控还是用 Prometheus。

总结一下

折腾完这一圈，其实没什么"一招鲜"的方案，每一步都是 trade-off。

TorchServe 适合快速上手的场景，但性能调优空间有限；vLLM 在吞吐量上有优势，但资源管理不够灵活；TensorRT 能榨干每一分性能，但开发和维护成本很高。

GPU 调度这块，没有银弹，只能根据业务特点选择合适的隔离策略。MIG 听起来美好，但不是所有环境都支持；MPS 适合轻量级服务，但干扰问题难以避免。

批量推理和监控也是一样，没有标准答案，只能在实践中不断调整参数和策略。

最后说句实话：模型部署这事儿，90% 的时间在解决各种边界情况和资源争抢问题，真正"优化推理"的时间可能只占 10%。但如果你能把这些边界处理好，剩下的那部分优化其实反而没那么难。

如果你也在搞模型部署，建议先把监控做好，再逐步优化。不然你可能优化了一个月，最后发现瓶颈根本不在推理本身，而在某个不起眼的地方。

参考资源：

版权声明: 本文首发于指尖魔法屋-AI模型部署：推理不够用了之后（https://blog.thinkmoon.cn/post/146-ai-model-deployment-inference-production-practice/）转载或引用必须申明原指尖魔法屋来源及源地址！

AI vLLM：PagedAttention不够用了之后

先说背景：我们团队去年要做个企业级问答系统，模型用的 Qwen2-72B。

📁 技术实践

📅 2024-05-11

AI模型服务性能踩坑记录

线上 Llama2-7B 对话服务跑在 AWS g4dn.xlarge（单卡 T4）上，用户抱怨回答慢：首 token 延迟 2–3 秒，完整回…