多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

从零手搓AI推理服务:显存优化、动态批处理与线上排障实战

从零手搓AI推理服务:显存优化、动态批处理与线上排障实战 1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件调一下API然后跑通一个Demo就觉得自己已经掌握了。我刚开始也是这么想的直到有一次线上推理服务在高峰期直接雪崩日志里全是显存溢出和请求超时我才意识到——那些被封装好的高级接口在关键时刻根本救不了你。ai-engineering-from-scratch这个标题说的就是从最底层开始把AI工程里那些被隐藏起来的环节一个一个亲手搭起来。这篇文章适合谁看如果你已经会用Python写点脚本知道什么是张量但每次遇到模型部署、显存优化、推理加速这些词就心里发虚那这篇内容就是为你准备的。我不会只给你一堆代码让你复制粘贴而是把每个设计决策背后的“为什么”讲清楚。比如为什么推理时要固定batch size为什么量化不是万能药为什么你的GPU利用率永远上不去。这些问题的答案在官方文档里往往一笔带过但在实际工程中每一个都能让你熬夜到凌晨三点。我打算从最基础的推理服务搭建开始一步步走到性能调优和线上排障。整个过程不依赖任何重型框架核心逻辑用Python和少量C扩展实现目的是让你看清AI工程的全貌。你不需要有分布式系统的经验但最好对Linux命令和Python多进程有点概念。如果你准备好了我们就从第一个坑开始。2. 推理服务的最小可行骨架从HTTP请求到张量计算2.1 为什么不用现成的模型服务框架市面上有很多模型服务框架开箱即用一行命令就能启动一个推理端点。但我坚持从零手写一个最小服务原因有三个。第一那些框架为了通用性抽象层数太多一个请求进来要经过路由、预处理、批处理调度、模型执行、后处理五六个模块每个模块都有配置项一旦出问题排查链路极长。第二框架的默认参数往往针对通用场景比如动态批处理窗口默认10毫秒在高并发低延迟场景下这个值就是灾难。第三也是最关键的当你亲手写过一遍请求解析、内存分配、线程调度之后再看那些框架的源码你会有一种“原来如此”的顿悟感。我选用的技术栈很简单Python标准库的http.server做HTTP层numpy做数据搬运模型本身用ONNX Runtime或者自己导出的TorchScript。为什么不直接上FastAPI因为FastAPI的异步模型和GPU推理的同步阻塞特性之间存在微妙的冲突新手很容易写出看似异步实则串行的代码。用最原始的http.server虽然性能差但能让你清楚看到每个请求的生命周期。2.2 请求解析与张量构造的隐藏成本一个典型的推理请求进来JSON里包含输入数据比如一张图片的base64编码或者一个文本序列的token ID列表。很多人直接json.loads然后np.array就完事了但这里有两个隐藏成本。第一base64解码是CPU密集操作一张1080p的JPEG图片解码成RGB数组大约需要15到30毫秒如果并发上来CPU会先于GPU成为瓶颈。第二np.array从Python列表构造时如果数据类型不匹配会发生隐式类型转换比如你的模型期望float32但JSON里的数字被解析成float64这个转换在数据量大时能吃掉几毫秒。我的做法是在服务启动时就预分配好输入缓冲区。对于固定尺寸的输入直接创建一个np.empty数组请求到来时用np.frombuffer或者切片赋值填充避免反复分配内存。对于变长输入维护一个内存池按2的幂次方分级管理。这些技巧在常规教程里很少提但它们是推理服务稳定性的基石。import numpy as np import base64 from io import BytesIO from PIL import Image class InputBuffer: def __init__(self, shape, dtypenp.float32): self.buffer np.empty(shape, dtypedtype) self.shape shape def fill_from_base64(self, b64_str): img_bytes base64.b64decode(b64_str) img Image.open(BytesIO(img_bytes)).convert(RGB) img img.resize((self.shape[2], self.shape[3])) arr np.asarray(img, dtypenp.float32) / 255.0 arr np.transpose(arr, (2, 0, 1)) np.copyto(self.buffer[0], arr) return self.buffer上面这段代码里np.copyto比直接赋值更安全因为它会检查形状和类型是否匹配。Image.resize默认使用双线性插值如果你对精度有要求可以换成Image.BICUBIC但速度会慢一倍。这些细节在批量推理时影响巨大。2.3 模型加载与执行提供者的选择模型加载阶段我强烈建议使用ONNX Runtime而不是直接加载PyTorch模型。原因在于ONNX Runtime的执行提供者Execution Provider机制允许你灵活切换后端。比如在NVIDIA GPU上你可以用CUDA EP在AMD GPU上用ROCM EP在CPU上用OpenVINO EP。更重要的是ONNX Runtime内置了图优化比如算子融合、常量折叠、内存复用这些优化在PyTorch eager模式下是没有的。但这里有个坑ONNX Runtime的默认优化级别是ORT_ENABLE_ALL它会尝试把所有能融合的算子都融合掉。在某些模型上过度融合会导致精度下降尤其是涉及LayerNorm和Softmax的模型。我的经验是先用ORT_ENABLE_BASIC跑一遍对比输出差异如果差异在可接受范围内比如1e-5再逐步提高优化级别。另外intra_op_num_threads这个参数控制算子内部的并行线程数默认是CPU核心数但在GPU推理场景下这个值设得太高反而会因为线程切换开销导致性能下降。我通常设为物理核心数的一半。import onnxruntime as ort options ort.SessionOptions() options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_BASIC options.intra_op_num_threads 4 options.inter_op_num_threads 1 session ort.InferenceSession( model.onnx, sess_optionsoptions, providers[CUDAExecutionProvider, CPUExecutionProvider] )注意providers列表的顺序ONNX Runtime会按顺序尝试如果CUDA不可用就回退到CPU。这个回退机制在生产环境中非常有用但你要确保CPU路径也能正常工作否则回退后直接报错。3. 批处理与调度让GPU不再空转的实战策略3.1 动态批处理的触发条件与超时权衡GPU最怕的就是空转。一个请求进来做一次前向传播然后等下一个请求GPU利用率可能只有5%。动态批处理的核心思想是攒几个请求一起算。但攒多久这就是超时窗口的权衡。窗口设得太短比如1毫秒那和单请求没区别设得太长比如50毫秒延迟就上去了用户体验变差。我的经验值是这样的对于在线服务超时窗口设为5到10毫秒对于离线批处理可以设为100毫秒甚至更长。但这不是固定的你要根据请求到达率动态调整。如果每秒请求数QPS很高比如1000那5毫秒内能攒到5个请求批大小为5GPU利用率就能到60%以上。如果QPS只有10那5毫秒内可能一个请求都没有这时候超时窗口应该适当延长比如20毫秒保证至少能凑到2个请求。实现上我用一个后台线程维护请求队列主线程收到请求后放入队列并等待条件变量。后台线程每隔一个超时周期检查队列如果非空就取出所有请求组成一个批次调用模型推理然后把结果分发给对应的等待者。这里的关键是条件变量的通知机制不能用time.sleep轮询那样CPU占用率会很高。import threading import time from collections import deque class BatchScheduler: def __init__(self, max_batch_size8, timeout_ms10): self.queue deque() self.max_batch_size max_batch_size self.timeout timeout_ms / 1000.0 self.lock threading.Lock() self.cond threading.Condition(self.lock) self.results {} self.running True self.worker threading.Thread(targetself._loop, daemonTrue) self.worker.start() def submit(self, request_id, input_data): with self.cond: self.queue.append((request_id, input_data)) self.cond.notify() while request_id not in self.results: self.cond.wait() return self.results.pop(request_id) def _loop(self): while self.running: with self.cond: if not self.queue: self.cond.wait(timeoutself.timeout) if not self.queue: continue batch [] while self.queue and len(batch) self.max_batch_size: batch.append(self.queue.popleft()) # 在锁外执行推理避免阻塞提交 inputs [item[1] for item in batch] outputs self._infer(inputs) with self.cond: for (req_id, _), out in zip(batch, outputs): self.results[req_id] out self.cond.notify_all()这段代码里_infer是实际调用模型的方法它应该在锁外执行否则提交请求的线程会被阻塞。另外self.results字典在请求量极大时会成为内存瓶颈生产环境应该用更高效的数据结构比如concurrent.futures.Future。3.2 批大小对显存和延迟的非线性影响很多人以为批大小翻倍显存占用也翻倍延迟也翻倍。实际上不是这样的。显存占用大致是线性的因为中间激活值随批大小线性增长。但延迟不是线性的因为GPU的并行计算单元在批大小较小时利用率不足批大小增大到一定程度后计算时间增长缓慢。我实测过一个ResNet-50模型批大小1时延迟8毫秒批大小8时延迟12毫秒批大小16时延迟18毫秒。也就是说批大小从1到8吞吐量提升了5倍多但延迟只增加了50%。但这里有个临界点。当批大小超过某个值后显存不够了或者计算单元饱和了延迟会急剧上升。这个临界点取决于模型大小和GPU型号。我的做法是写一个简单的压测脚本从批大小1开始每次翻倍记录延迟和显存占用画出曲线找到拐点。通常拐点在GPU显存的70%到80%利用率处。批大小延迟(ms)显存占用(MB)吞吐量(请求/秒)181200125291350222410165040081222506661618345088832355850914从表里可以看出批大小16到32吞吐量几乎没提升但延迟翻倍。所以最优批大小在16左右。这个数据因模型和硬件而异但方法论是通用的。3.3 请求优先级与超时丢弃机制生产环境里请求不是平等的。有些是实时交互请求用户等着看结果延迟要求高有些是后台分析请求可以等几秒钟。如果一视同仁实时请求会被后台请求拖累。我的做法是在请求头里加一个优先级字段调度器维护两个队列高优先级和低优先级。每次组批时先从高优先级队列取取完了再用低优先级填充剩余位置。超时丢弃也很重要。如果一个请求在队列里等了超过阈值比如200毫秒还没被处理就应该直接返回超时错误而不是继续等。因为用户可能已经重试了你再算出来也没用反而浪费GPU资源。实现上每个请求入队时记录时间戳调度器组批前检查时间戳过期的直接丢弃并通知等待者。注意超时丢弃的阈值要大于批处理超时窗口否则请求还没等到组批就被丢了。一般设为批处理窗口的5到10倍。4. 显存管理与量化在有限资源下榨取性能4.1 显存碎片化与内存池设计GPU显存和CPU内存一样反复分配释放会产生碎片。尤其是变长输入场景每次请求的输入尺寸不同如果每次都cudaMalloc和cudaFree很快显存就碎片化了最后明明有足够的总空闲显存却分配不出一块连续的大内存。解决方案是内存池启动时一次性分配一大块显存然后自己管理分配和释放。我的内存池实现很简单按2的幂次方分级比如1KB、2KB、4KB直到256MB。每个级别维护一个空闲链表。分配时向上取整到最近的级别从对应链表取一块释放时归还到对应链表。如果某个级别空了从更高级别切分一块下来。这个策略在请求尺寸分布比较集中的场景下非常高效碎片率可以控制在5%以内。但内存池有个缺点启动时就要确定总大小。设得太小高峰期不够用设得太大浪费显存。我的经验是先跑一轮压测记录峰值显存占用然后内存池大小设为峰值的1.2倍。另外要留出至少20%的显存给模型权重和CUDA上下文否则会OOM。4.2 量化不是银弹INT8精度的实际损失评估量化是显存优化的利器FP32转INT8显存直接降到四分之一推理速度也能提升2到3倍。但量化会损失精度而且损失程度因模型而异。我见过一个文本分类模型INT8量化后准确率只掉了0.1%完全可用也见过一个目标检测模型量化后小目标的召回率掉了15%直接不可用。所以量化之前一定要做精度评估。步骤是准备一个验证集至少1000个样本用FP32模型跑一遍记录每个样本的输出用INT8模型跑一遍对比输出差异。对于分类模型看Top-1准确率变化对于检测模型看mAP变化对于生成模型看BLEU或者ROUGE变化。如果掉点超过可接受范围就要考虑混合量化只量化对精度不敏感的层比如卷积层保留全连接层为FP32。ONNX Runtime提供了动态量化和静态量化两种模式。动态量化不需要校准数据直接对权重做量化激活值在推理时动态计算量化参数。静态量化需要校准数据提前计算好激活值的量化参数精度通常更好。我的建议是如果手头有代表性数据优先用静态量化如果没有动态量化也能用但要做好精度下降的心理准备。from onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_inputmodel_fp32.onnx, model_outputmodel_int8.onnx, weight_typeQuantType.QUInt8, optimize_modelTrue )QuantType.QUInt8是无符号8位整数量化适合ReLU激活的模型。如果是LeakyReLU或者GELU用QInt8有符号量化更合适。optimize_modelTrue会在量化前先做图优化通常能提升量化后的精度。4.3 显存不足时的降级策略即使做了量化和内存池高峰期仍然可能显存不足。这时候不能直接崩掉要有降级策略。我的降级顺序是这样的第一级降低批大小从16降到8显存占用减半吞吐量下降但服务不中断第二级切换到CPU推理延迟大幅上升但至少能返回结果第三级返回503错误让上游重试或降级。实现上用一个显存监控线程每隔100毫秒查询一次torch.cuda.memory_allocated()或者ONNX Runtime的显存统计。如果超过阈值触发降级信号。降级信号通过一个全局状态变量传递给调度器调度器动态调整批大小上限。这个机制在Kubernetes环境下尤其重要因为容器有显存限制超了直接被杀。提示降级策略要提前演练不能等线上出问题了才第一次跑。我一般会在预发环境用压力测试工具模拟显存打满验证降级链路是否通畅。5. 线上排障实录一次推理延迟毛刺的完整排查链路5.1 问题现象与初步定位那是一个周五晚上监控系统突然报警推理服务P99延迟从15毫秒飙升到200毫秒但QPS没有明显变化。更奇怪的是P50延迟还是正常的15毫秒只有P99异常。这说明大部分请求正常少数请求特别慢。我第一反应是某个请求的输入尺寸特别大导致计算时间过长。但查了日志输入尺寸分布和平时一样。接着我怀疑是GPU降频。用nvidia-smi -q -d PERFORMANCE查看发现GPU时钟确实从1.8GHz降到了1.2GHz。但为什么降频温度正常功耗也正常。继续查发现nvidia-smi显示有一个进程占用了大量显存但不是我们的推理进程。原来是一台共享GPU的机器上另一个团队跑了一个训练任务把显存吃掉了大半导致我们的推理进程频繁触发显存回收进而引发延迟毛刺。5.2 排查过程中的三个关键命令第一个命令是nvidia-smi --query-compute-appspid,used_memory --formatcsv列出所有占用GPU的进程和显存。这个命令比直接nvidia-smi更清晰因为它只显示计算进程不显示图形进程。第二个命令是py-spy dump --pid 推理进程PID这是Python的采样分析工具能直接打印出当前所有线程的调用栈。我用它发现推理主线程卡在了cudaMemcpy上说明数据在CPU和GPU之间搬运时被阻塞了。第三个命令是nsys profile --statstrue -o profile_result python infer.py这是NVIDIA的系统级性能分析工具能生成时间线显示每个CUDA核函数的执行时间和内存拷贝时间。从时间线上看正常情况下cudaMemcpy只占5%的时间出问题时占到了60%。5.3 根因分析与修复方案根因很明确共享GPU环境下另一个进程的显存占用导致我们的进程在分配显存时触发了同步等待。CUDA的显存分配是同步操作当显存不足时驱动会尝试回收其他进程的缓存这个回收过程会阻塞当前进程。修复方案有三个第一申请独占GPU不让其他进程共享第二设置CUDA_VISIBLE_DEVICES隔离第三在代码里预分配所有需要的显存避免运行时分配。我选了第三个方案因为最可控。具体做法是在服务启动时用torch.cuda.memory_reserved()或者ONNX Runtime的arena_extend_strategy参数一次性预留足够显存。ONNX Runtime的arena_extend_strategy设为kSameAsRequested表示按需扩展但扩展后不释放这样后续请求就不会触发新的分配。options ort.SessionOptions() options.enable_cpu_mem_arena False options.add_session_config_entry(session.use_env_allocators, 1) options.add_session_config_entry(session.arena_extend_strategy, kSameAsRequested)修复后P99延迟回落到18毫秒虽然比P50的15毫秒略高但已经可以接受。这次排查让我深刻体会到AI工程不只是模型和算法更是对系统资源的精细管理。5.4 从这次故障中提炼的检查清单后来我把这次排查经验整理成了一个检查清单每次上线新服务前过一遍GPU是否独占用nvidia-smi确认没有其他计算进程。显存是否预分配启动后观察nvidia-smi的显存占用是否稳定。是否有降级策略模拟显存打满验证降级链路。监控是否覆盖P99只看平均值会漏掉毛刺。日志是否记录输入尺寸方便定位大请求。是否有性能基线每次变更后对比基线发现退化。这个清单帮我避免了好几次潜在故障。比如有一次新来的同事在服务里加了一个日志打印把每个请求的完整输入都打出来了导致磁盘IO飙升推理延迟跟着上涨。用清单里的“性能基线”一对比立刻发现了问题。6. 从手搓到生产还需要补上的几块拼图6.1 健康检查与优雅退出手搓的服务往往忽略健康检查和优雅退出。健康检查不是简单的返回200而是要检查模型是否加载成功、GPU是否可用、显存是否充足。我的做法是暴露一个/health端点里面依次检查模型session是否非空、cudaGetDeviceCount是否大于0、显存空闲是否大于阈值。只有全部通过才返回200否则返回503。优雅退出是指收到SIGTERM信号后不再接受新请求等待正在处理的请求完成然后释放资源再退出。这个在Kubernetes滚动更新时特别重要否则正在处理的请求会被直接杀掉用户看到502错误。实现上用一个全局的shutting_down标志主循环检查这个标志如果为真就关闭监听套接字然后等待所有工作线程结束。6.2 日志与指标暴露日志要结构化用JSON格式包含时间戳、请求ID、输入尺寸、批大小、推理耗时、显存占用。这些字段在排查问题时缺一不可。指标暴露用Prometheus格式暴露inference_latency_seconds直方图、batch_size直方图、gpu_memory_used_bytes仪表盘。有了这些指标你就能在Grafana上画出延迟分布和显存趋势提前发现异常。我见过很多团队只打日志不暴露指标结果每次排查都要grep日志效率极低。指标是聚合的能一眼看出趋势日志是离散的用于定位具体请求。两者互补缺一不可。6.3 版本管理与回滚模型文件要版本化每次更新模型都要记录版本号、训练数据哈希、评估指标。服务启动时加载指定版本的模型并在日志里打印版本信息。如果新模型上线后指标恶化要能一键回滚到旧版本。我的做法是把模型文件放在对象存储里用版本号做路径服务启动时从配置中心读取当前版本号然后下载对应模型。回滚就是改配置中心的版本号重启服务。这个机制看起来简单但很多团队直到出了事故才想起来做。我经历过一次模型更新后准确率暴跌因为没有版本管理花了两个小时才找到旧模型文件。从那以后版本管理成了我的必选项。6.4 压测与容量规划上线前必须压测。压测不是简单地用ab或者wrk打流量而是要模拟真实请求分布输入尺寸有长有短请求到达有突发有平稳。我用Locust写压测脚本定义多个任务每个任务有不同的输入尺寸和权重。压测时逐步增加并发用户数观察延迟和显存的变化找到服务能承受的最大QPS。容量规划就是根据压测结果计算需要多少台机器。比如单机最大QPS是100预计峰值QPS是800那至少需要8台机器再留20%余量就是10台。这个计算要定期回顾因为模型更新后性能会变化。注意压测环境要和生产环境硬件一致否则数据没有参考价值。我见过在CPU机器上压测然后上线到GPU机器结果完全对不上。7. 一些让我少走弯路的个人习惯我刚开始做AI工程的时候总想一步到位把服务写得尽善尽美。结果往往是过度设计代码复杂到自己也看不懂。后来我养成了一个习惯先写一个最笨的版本能跑通就行然后在此基础上逐步优化。比如批处理第一版就是批大小1跑通了再加动态批处理量化第一版就是FP32跑通了再试INT8。这样每一步都有基线出了问题也知道是哪个改动引入的。另一个习惯是记录“失败日志”。每次遇到坑解决之后花五分钟写下来现象是什么、排查过程、根因、修复方案。这些记录后来成了我自己的知识库下次遇到类似问题直接搜关键词就能找到答案。我建议你也这么做不用写得多正式几句话就行关键是坚持。还有一个反直觉的经验不要过早优化。我见过一个团队服务还没上线就开始搞模型蒸馏、算子融合、多流并行结果上线后发现QPS只有个位数根本用不上这些优化。先让服务跑起来有了真实流量再根据瓶颈优化。瓶颈可能在CPU预处理可能在网络IO也可能在GPU计算不跑起来你永远不知道。最后保持对底层的好奇心。当你用nvidia-smi看到GPU利用率只有30%的时候不要满足于“能用就行”去查查为什么只有30%是批大小不够还是数据加载拖了后腿还是核函数本身效率低。每一次深挖都会让你对AI工程的理解更深一层。这个领域变化很快但底层原理变化很慢把底层吃透了上层的新框架新工具你都能快速上手。
返回列表