1. 从零搭建AI工程能力为什么我劝你别急着调库这两年AI应用开发的门槛肉眼可见地在降低随便拉个框架、调几个API一个能跑通的demo就出来了。但我自己带过几支团队、也面试过不少号称“做过AI项目”的候选人之后发现一个很普遍的现象很多人能跑通demo却说不清楚一次推理请求背后到底发生了什么模型加载为什么慢、显存为什么爆、batch size调大之后延迟为什么非线性增长这些问题一出来就卡壳。ai-engineering-from-scratch这个标题在我看来切中的正是这个痛点。它不是让你从零训练一个大模型而是让你从零把AI工程这条链路搭起来——从环境、依赖、模型加载、推理服务、性能观测到部署上线每一环都自己动手过一遍。适合谁看我觉得有三类人一是刚转行做AI应用、只会调库的开发者二是想搞清楚推理服务底层原理的后端工程师三是需要自己搭一套可控推理环境、不想被各种托管平台绑死的技术负责人。这篇文章我会按我实际搭一套本地推理服务的顺序来讲把每一步的选型理由、参数计算、踩过的坑都摊开说。你不需要有很深的机器学习背景但最好懂一点Python和Linux基础剩下的我们边做边补。2. 整体思路与方案选型为什么这么搭2.1 先想清楚“从零”到底指什么很多人一看到“from scratch”就以为是手写矩阵乘法、自己实现反向传播。那是研究向的从零不是工程向的从零。工程向的从零指的是你不依赖某个高度封装的托管服务自己把模型跑起来、把服务暴露出去、把性能盯住。模型权重可以用开源的推理框架可以用成熟的但整条链路的组装、配置、调优必须是你自己掌控的。这个定位很重要因为它直接决定了你的技术选型。如果你目标是理解Transformer的每一个数学细节那你该去看论文和教学实现如果你目标是能独立交付一个可用的AI服务那你要关注的是推理引擎、服务框架、资源管理和观测体系。ai-engineering-from-scratch我倾向于把它归到后者。2.2 推理引擎怎么选别一上来就上最重的推理引擎这块市面上常见的有几类原生PyTorch直接推理、ONNX Runtime、TensorRT、以及各种针对大模型优化的推理框架。我的建议是分阶段来不要一上来就追求极致性能。原生PyTorch推理的好处是零转换成本模型拿来就能跑调试也方便缺点是性能一般、显存占用高。ONNX Runtime在中小模型上性价比很高转换一次之后跨平台部署也方便。TensorRT性能最强但转换和调优成本高模型结构一变就得重来。大模型场景下专门优化的推理框架在吞吐和显存上优势明显但配置复杂度也上去了。我一般的做法是先用原生PyTorch把链路跑通确认输入输出、预处理后处理逻辑都对然后根据实际性能瓶颈决定要不要上ONNX或更专门的引擎。这个顺序能帮你避免一个常见错误——花两天调TensorRT最后发现瓶颈其实在数据预处理上。2.3 服务框架FastAPI够用但要知道它的边界服务层我基本都用FastAPI原因很实际异步支持好、类型校验省心、自动生成接口文档、生态成熟。对于大多数中小规模的推理服务FastAPI完全够用。但你要清楚它的边界。FastAPI本身是异步框架可推理计算是同步阻塞的如果你直接在async函数里跑推理会把事件循环堵死。正确做法是把推理放到线程池或独立的进程里用run_in_executor或者干脆用同步路由。这个坑我见过太多人踩表现就是并发一上来延迟飙升但GPU利用率却不高。2.4 观测体系没有它你就是盲人摸象一套没有观测的AI服务等于闭着眼睛开车。我最少会盯这几个指标请求延迟的P50/P95/P99、GPU利用率和显存占用、每秒处理请求数、以及队列积压情况。这些指标能帮你快速定位瓶颈到底在计算、在IO还是在调度。工具上Prometheus加Grafana是经典组合成本低、生态好。如果只是本地调试用简单的日志加定时打印也够。关键不是工具多高级而是你得有数据能回答“现在到底慢在哪”这个问题。3. 核心细节解析与实操要点3.1 环境隔离别让依赖地狱毁了你的一天AI工程的依赖冲突是出了名的严重。PyTorch、CUDA、cuDNN、各种推理框架版本之间互相卡。我的铁律是一个项目一个独立环境绝不在系统Python里装东西。用conda还是venv涉及CUDA相关依赖时我更倾向conda因为它能帮你管理非Python的二进制依赖。纯Python依赖的话venv更轻。创建环境之后第一件事是锁定Python版本第二件事是先把PyTorch装对——注意要去官方看对应CUDA版本的安装命令别直接pip install torch那样装到的可能是CPU版本跑起来慢得让你怀疑人生。提示装完PyTorch后立刻用torch.cuda.is_available()验证返回False就说明CUDA没配好别急着往下走。3.2 模型加载显存是怎么被吃掉的模型加载这一步很多人只知道“加载慢”但不知道慢在哪、显存花在哪。我拆一下模型权重本身占一部分显存这部分是固定的推理过程中的中间激活值占一部分这部分跟batch size和序列长度强相关还有框架自身的一些开销。以FP16精度为例一个7B参数的模型权重约占14GB显存7B × 2字节。如果你用FP32直接翻倍到28GB。这就是为什么量化这么重要——INT8量化能把权重压到约7GBINT4能压到约3.5GB代价是精度会有一定损失。加载慢的问题一部分是磁盘IO一部分是权重初始化。如果模型放在机械硬盘上加载时间会明显拉长换SSD能改善不少。另外可以用内存映射的方式加载减少一次性内存占用。3.3 预处理与后处理最容易被忽视的性能杀手我见过太多案例GPU利用率上不去最后查出来瓶颈在CPU上的图像解码或文本分词。预处理和后处理跑在CPU上如果实现得不够高效就会成为整个链路的短板。文本场景下分词器的选择影响很大。有些分词器实现是纯Python的速度慢换成基于Rust或C实现的版本吞吐能提升好几倍。图像场景下解码和resize尽量用硬件加速或优化过的库别用最朴素的PIL逐张处理。还有一个细节预处理能不能批量化。单条处理时CPU开销被放大批量处理能摊薄。但批量又不能太大否则延迟上去了。这个平衡点要靠实测找。3.4 批处理策略吞吐和延迟的跷跷板批处理是提升吞吐最直接的手段但它和延迟是一对矛盾。batch size越大单位时间处理的请求越多但单个请求要等更久才能凑够一批。静态批处理实现简单固定一个batch size攒够就推理。缺点是请求不均匀时要么等太久要么批次装不满浪费算力。动态批处理更聪明设定一个最大等待时间超时就用当前攒到的请求推理能在延迟和吞吐之间取得更好的平衡。我一般会先测出不同batch size下的延迟曲线找到那个“再往上加延迟就明显恶化”的拐点把最大batch size设在那附近再配合一个合理的超时时间。4. 实操过程与核心环节实现4.1 环境搭建的完整步骤先把基础环境搭起来。我以Linux环境为例Windows的话建议用WSL2原生Windows在AI工程上坑太多。第一步确认显卡驱动和CUDA版本。用nvidia-smi看驱动支持的CUDA版本这个决定了你能装哪个版本的PyTorch。假设显示支持到12.1那你就装对应CUDA 12.1的PyTorch。第二步创建conda环境conda create -n ai-eng python3.10 -y conda activate ai-eng第三步装PyTorch。去官网查对应命令比如pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121第四步装服务框架和观测工具pip install fastapi uvicorn prometheus-client第五步验证import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))这三行输出正常环境就算搭好了。如果cuda.is_available()是False八成是PyTorch版本和CUDA版本对不上回去重装。4.2 一个最小可用的推理服务先写一个最朴素的版本把链路跑通。核心是把模型加载和服务启动分开模型只加载一次服务常驻。import torch from fastapi import FastAPI from pydantic import BaseModel from transformers import AutoTokenizer, AutoModelForCausalLM app FastAPI() model_name your-model-path tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained( model_name, torch_dtypetorch.float16, device_mapauto ) model.eval() class Request(BaseModel): text: str max_new_tokens: int 128 app.post(/generate) def generate(req: Request): inputs tokenizer(req.text, return_tensorspt).to(model.device) with torch.no_grad(): outputs model.generate( **inputs, max_new_tokensreq.max_new_tokens, do_sampleFalse ) result tokenizer.decode(outputs[0], skip_special_tokensTrue) return {result: result}启动命令uvicorn main:app --host 0.0.0.0 --port 8000这个版本能跑但有几个问题路由是同步的FastAPI会把它放到线程池里执行不会堵事件循环这点是对的但每次请求都重新做分词和生成没有批处理吞吐上不去也没有任何观测。我们一步步优化。4.3 加入动态批处理动态批处理的核心是维护一个请求队列后台有个worker不断从队列取请求攒到一定数量或超时后一起推理再把结果分发回去。import asyncio import time from collections import deque class BatchProcessor: def __init__(self, model, tokenizer, max_batch8, max_wait0.05): self.model model self.tokenizer tokenizer self.max_batch max_batch self.max_wait max_wait self.queue deque() self.lock asyncio.Lock() async def submit(self, text, max_new_tokens): future asyncio.get_event_loop().create_future() async with self.lock: self.queue.append((text, max_new_tokens, future)) return await future async def worker(self): while True: if not self.queue: await asyncio.sleep(0.001) continue batch [] start time.time() while self.queue and len(batch) self.max_batch: batch.append(self.queue.popleft()) if time.time() - start self.max_wait: break await self.process(batch) async def process(self, batch): texts [item[0] for item in batch] max_tokens max(item[1] for item in batch) inputs self.tokenizer( texts, return_tensorspt, paddingTrue, truncationTrue ).to(self.model.device) with torch.no_grad(): outputs self.model.generate( **inputs, max_new_tokensmax_tokens, do_sampleFalse ) for i, item in enumerate(batch): result self.tokenizer.decode(outputs[i], skip_special_tokensTrue) item[2].set_result(result)这里有几个关键点。max_batch和max_wait是两个核心参数前者决定吞吐上限后者决定延迟上限。padding是必须的因为一个batch里的序列长度不一样要补齐到相同长度。但padding会浪费一些计算所以序列长度差异特别大的时候可以考虑按长度分桶。4.4 参数计算batch size和显存的关系显存占用可以粗略估算。以推理为例主要分三块模型权重、激活值、以及框架开销。模型权重参数量 × 精度字节数。7B模型FP16约14GB。激活值跟batch size、序列长度、隐藏层维度相关。粗略估算每层激活值约等于 batch_size × seq_len × hidden_dim × 精度字节数 × 一个系数。这个系数跟具体结构有关实测更准。假设你的卡有24GB显存模型权重占14GB剩10GB给激活值和框架开销。框架开销留2GB剩8GB给激活值。如果单条序列激活值占0.5GB那batch size最多到16左右。但这只是估算实际要留余量因为显存碎片和峰值占用可能更高。我的经验是先按估算值的一半设跑起来看实际显存占用再逐步往上加加到接近上限但留10%余量为止。4.5 观测接入用prometheus-client暴露指标在推理前后打点from prometheus_client import Histogram, Counter, Gauge REQUEST_LATENCY Histogram(inference_latency_seconds, 推理延迟) REQUEST_COUNT Counter(inference_requests_total, 请求总数) BATCH_SIZE Histogram(inference_batch_size, 批大小分布) QUEUE_LENGTH Gauge(inference_queue_length, 队列长度)在process函数里记录延迟和批大小在submit时更新队列长度。然后起一个/metrics端点暴露出去Prometheus定时抓取Grafana画图。这套东西搭起来之后你就能看到延迟分布、批大小分布、队列积压情况。哪个环节有问题一眼就能看出来。5. 常见问题与排查技巧实录5.1 显存溢出OOM怎么排查OOM是最常见的问题。排查思路是先确认是加载时OOM还是推理时OOM。加载时OOM说明模型权重就超了得换量化版本或换卡。推理时OOM说明激活值超了得降batch size或降序列长度。一个容易被忽视的点是显存碎片。长时间运行的服务反复分配释放显存可能产生碎片导致明明总显存够用却分配不出来。解决办法是设置PYTORCH_CUDA_ALLOC_CONF环境变量开启更高效的内存分配策略。提示export PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128能缓解碎片问题但会略微增加内存占用。5.2 延迟忽高忽低延迟不稳定通常是几个原因批处理等待时间波动、GPU被其他进程抢占、或者请求长度差异大导致计算量波动。先看批大小分布如果批大小波动很大说明请求不均匀可以适当调大max_wait让批次更稳定。再看GPU是否有其他进程用nvidia-smi确认。如果是请求长度问题考虑按长度分桶把长度相近的请求放一批。5.3 吞吐上不去但GPU利用率低这个现象说明瓶颈不在GPU在CPU或IO。常见原因预处理太慢、分词器是纯Python实现、数据在CPU和GPU之间频繁拷贝。排查方法是用profiler看时间花在哪。PyTorch自带的profiler能给出每个操作的耗时。如果发现大量时间花在数据搬运上考虑用pin_memory和异步拷贝。如果花在分词上换更快的分词器实现。5.4 常见问题速查表问题现象可能原因排查方向解决思路加载时OOM模型权重超显存看模型参数量和精度用量化版本或换卡推理时OOM激活值超显存看batch size和序列长度降batch或降长度延迟波动大批处理不稳定看批大小分布调max_wait或分桶吞吐低GPU闲CPU瓶颈用profiler定位优化预处理或换实现并发上不去事件循环阻塞看路由是否同步用线程池或同步路由显存碎片反复分配释放看长时间运行表现设置分配策略环境变量5.5 几个我踩过的坑第一个坑是分词器的padding方向。有些模型是左padding有些是右padding搞错了生成结果会完全不对。这个一定要看模型文档确认。第二个坑是torch.no_grad()忘了加。推理时如果没加PyTorch会构建计算图显存占用暴涨速度也慢。这个错误很隐蔽因为结果是对的只是慢和费显存。第三个坑是模型eval模式忘了切。有些层在训练和推理时行为不同比如dropout和batchnorm。忘了model.eval()会导致结果不稳定。第四个坑是并发测试时用了同步客户端。测服务并发能力时客户端本身如果是同步阻塞的测出来的并发数是不准的。要用异步客户端或者多进程压测。6. 从能跑到好用进一步优化的方向6.1 量化精度和资源的权衡量化是降低显存、提升速度的有效手段。INT8量化通常精度损失很小INT4损失大一些但显存占用大幅下降。量化的方式有训练后量化和量化感知训练前者简单后者精度更好但需要训练。实操上如果只是推理训练后量化基本够用。要注意的是不是所有层都适合量化有些敏感层量化后精度掉得厉害需要混合精度策略。6.2 模型编译让计算图更高效PyTorch 2.0之后引入了编译功能能把模型的计算图优化减少Python开销和冗余计算。对于推理服务开启编译通常能带来明显的速度提升代价是首次编译需要一些时间。用法很简单model torch.compile(model)。但要注意不是所有模型都能顺利编译遇到不支持的算子会回退到eager模式。编译后的模型第一次推理会慢因为要编译之后才快。6.3 多卡与多实例单卡跑不下或者吞吐不够时考虑多卡。多卡有两种模式模型并行和数据并行。模型并行是把模型切到多张卡上适合单卡放不下的大模型数据并行是每张卡跑一份完整模型适合吞吐不够的场景。数据并行实现简单起多个服务实例前面挂个负载均衡就行。模型并行复杂一些需要框架支持。我一般优先考虑数据并行因为简单可靠扩展性也好。6.4 服务治理健康检查与优雅退出生产环境的服务健康检查和优雅退出是必须的。健康检查让负载均衡知道实例是否可用优雅退出保证正在处理的请求能完成再关闭。FastAPI里加健康检查很简单一个返回200的端点就行。优雅退出需要监听关闭信号停止接收新请求等队列清空后再退出。这个逻辑要自己写框架不会自动帮你做。7. 我个人的一些体会搭这套东西的过程中我最大的感受是AI工程的难点往往不在AI本身而在工程。模型是现成的框架是成熟的真正花时间的是环境配置、性能调优、问题排查这些“脏活累活”。但恰恰是这些活决定了一个AI服务能不能真正上线、能不能扛住流量。另一个体会是观测的重要性怎么强调都不过分。我早期做服务的时候出了问题全靠猜效率极低。后来把观测体系搭起来很多问题几分钟就能定位。数据不会骗人有数据就有方向。最后分享一个小技巧每次改动只改一个变量改完立刻测。AI工程涉及的因素太多一次改好几个地方出了问题根本不知道是哪个改动导致的。控制变量法虽然笨但最有效。
1. 从零搭建AI工程能力为什么我劝你别急着调库这两年AI应用开发的门槛肉眼可见地在降低随便拉个框架、调几个API一个能跑通的demo就出来了。但我自己带过几支团队、也面试过不少号称“做过AI项目”的候选人之后发现一个很普遍的现象很多人能跑通demo却说不清楚一次推理请求背后到底发生了什么模型加载为什么慢、显存为什么爆、batch size调大之后延迟为什么非线性增长这些问题一出来就卡壳。ai-engineering-from-scratch这个标题在我看来切中的正是这个痛点。它不是让你从零训练一个大模型而是让你从零把AI工程这条链路搭起来——从环境、依赖、模型加载、推理服务、性能观测到部署上线每一环都自己动手过一遍。适合谁看我觉得有三类人一是刚转行做AI应用、只会调库的开发者二是想搞清楚推理服务底层原理的后端工程师三是需要自己搭一套可控推理环境、不想被各种托管平台绑死的技术负责人。这篇文章我会按我实际搭一套本地推理服务的顺序来讲把每一步的选型理由、参数计算、踩过的坑都摊开说。你不需要有很深的机器学习背景但最好懂一点Python和Linux基础剩下的我们边做边补。2. 整体思路与方案选型为什么这么搭2.1 先想清楚“从零”到底指什么很多人一看到“from scratch”就以为是手写矩阵乘法、自己实现反向传播。那是研究向的从零不是工程向的从零。工程向的从零指的是你不依赖某个高度封装的托管服务自己把模型跑起来、把服务暴露出去、把性能盯住。模型权重可以用开源的推理框架可以用成熟的但整条链路的组装、配置、调优必须是你自己掌控的。这个定位很重要因为它直接决定了你的技术选型。如果你目标是理解Transformer的每一个数学细节那你该去看论文和教学实现如果你目标是能独立交付一个可用的AI服务那你要关注的是推理引擎、服务框架、资源管理和观测体系。ai-engineering-from-scratch我倾向于把它归到后者。2.2 推理引擎怎么选别一上来就上最重的推理引擎这块市面上常见的有几类原生PyTorch直接推理、ONNX Runtime、TensorRT、以及各种针对大模型优化的推理框架。我的建议是分阶段来不要一上来就追求极致性能。原生PyTorch推理的好处是零转换成本模型拿来就能跑调试也方便缺点是性能一般、显存占用高。ONNX Runtime在中小模型上性价比很高转换一次之后跨平台部署也方便。TensorRT性能最强但转换和调优成本高模型结构一变就得重来。大模型场景下专门优化的推理框架在吞吐和显存上优势明显但配置复杂度也上去了。我一般的做法是先用原生PyTorch把链路跑通确认输入输出、预处理后处理逻辑都对然后根据实际性能瓶颈决定要不要上ONNX或更专门的引擎。这个顺序能帮你避免一个常见错误——花两天调TensorRT最后发现瓶颈其实在数据预处理上。2.3 服务框架FastAPI够用但要知道它的边界服务层我基本都用FastAPI原因很实际异步支持好、类型校验省心、自动生成接口文档、生态成熟。对于大多数中小规模的推理服务FastAPI完全够用。但你要清楚它的边界。FastAPI本身是异步框架可推理计算是同步阻塞的如果你直接在async函数里跑推理会把事件循环堵死。正确做法是把推理放到线程池或独立的进程里用run_in_executor或者干脆用同步路由。这个坑我见过太多人踩表现就是并发一上来延迟飙升但GPU利用率却不高。2.4 观测体系没有它你就是盲人摸象一套没有观测的AI服务等于闭着眼睛开车。我最少会盯这几个指标请求延迟的P50/P95/P99、GPU利用率和显存占用、每秒处理请求数、以及队列积压情况。这些指标能帮你快速定位瓶颈到底在计算、在IO还是在调度。工具上Prometheus加Grafana是经典组合成本低、生态好。如果只是本地调试用简单的日志加定时打印也够。关键不是工具多高级而是你得有数据能回答“现在到底慢在哪”这个问题。3. 核心细节解析与实操要点3.1 环境隔离别让依赖地狱毁了你的一天AI工程的依赖冲突是出了名的严重。PyTorch、CUDA、cuDNN、各种推理框架版本之间互相卡。我的铁律是一个项目一个独立环境绝不在系统Python里装东西。用conda还是venv涉及CUDA相关依赖时我更倾向conda因为它能帮你管理非Python的二进制依赖。纯Python依赖的话venv更轻。创建环境之后第一件事是锁定Python版本第二件事是先把PyTorch装对——注意要去官方看对应CUDA版本的安装命令别直接pip install torch那样装到的可能是CPU版本跑起来慢得让你怀疑人生。提示装完PyTorch后立刻用torch.cuda.is_available()验证返回False就说明CUDA没配好别急着往下走。3.2 模型加载显存是怎么被吃掉的模型加载这一步很多人只知道“加载慢”但不知道慢在哪、显存花在哪。我拆一下模型权重本身占一部分显存这部分是固定的推理过程中的中间激活值占一部分这部分跟batch size和序列长度强相关还有框架自身的一些开销。以FP16精度为例一个7B参数的模型权重约占14GB显存7B × 2字节。如果你用FP32直接翻倍到28GB。这就是为什么量化这么重要——INT8量化能把权重压到约7GBINT4能压到约3.5GB代价是精度会有一定损失。加载慢的问题一部分是磁盘IO一部分是权重初始化。如果模型放在机械硬盘上加载时间会明显拉长换SSD能改善不少。另外可以用内存映射的方式加载减少一次性内存占用。3.3 预处理与后处理最容易被忽视的性能杀手我见过太多案例GPU利用率上不去最后查出来瓶颈在CPU上的图像解码或文本分词。预处理和后处理跑在CPU上如果实现得不够高效就会成为整个链路的短板。文本场景下分词器的选择影响很大。有些分词器实现是纯Python的速度慢换成基于Rust或C实现的版本吞吐能提升好几倍。图像场景下解码和resize尽量用硬件加速或优化过的库别用最朴素的PIL逐张处理。还有一个细节预处理能不能批量化。单条处理时CPU开销被放大批量处理能摊薄。但批量又不能太大否则延迟上去了。这个平衡点要靠实测找。3.4 批处理策略吞吐和延迟的跷跷板批处理是提升吞吐最直接的手段但它和延迟是一对矛盾。batch size越大单位时间处理的请求越多但单个请求要等更久才能凑够一批。静态批处理实现简单固定一个batch size攒够就推理。缺点是请求不均匀时要么等太久要么批次装不满浪费算力。动态批处理更聪明设定一个最大等待时间超时就用当前攒到的请求推理能在延迟和吞吐之间取得更好的平衡。我一般会先测出不同batch size下的延迟曲线找到那个“再往上加延迟就明显恶化”的拐点把最大batch size设在那附近再配合一个合理的超时时间。4. 实操过程与核心环节实现4.1 环境搭建的完整步骤先把基础环境搭起来。我以Linux环境为例Windows的话建议用WSL2原生Windows在AI工程上坑太多。第一步确认显卡驱动和CUDA版本。用nvidia-smi看驱动支持的CUDA版本这个决定了你能装哪个版本的PyTorch。假设显示支持到12.1那你就装对应CUDA 12.1的PyTorch。第二步创建conda环境conda create -n ai-eng python3.10 -y conda activate ai-eng第三步装PyTorch。去官网查对应命令比如pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121第四步装服务框架和观测工具pip install fastapi uvicorn prometheus-client第五步验证import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))这三行输出正常环境就算搭好了。如果cuda.is_available()是False八成是PyTorch版本和CUDA版本对不上回去重装。4.2 一个最小可用的推理服务先写一个最朴素的版本把链路跑通。核心是把模型加载和服务启动分开模型只加载一次服务常驻。import torch from fastapi import FastAPI from pydantic import BaseModel from transformers import AutoTokenizer, AutoModelForCausalLM app FastAPI() model_name your-model-path tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained( model_name, torch_dtypetorch.float16, device_mapauto ) model.eval() class Request(BaseModel): text: str max_new_tokens: int 128 app.post(/generate) def generate(req: Request): inputs tokenizer(req.text, return_tensorspt).to(model.device) with torch.no_grad(): outputs model.generate( **inputs, max_new_tokensreq.max_new_tokens, do_sampleFalse ) result tokenizer.decode(outputs[0], skip_special_tokensTrue) return {result: result}启动命令uvicorn main:app --host 0.0.0.0 --port 8000这个版本能跑但有几个问题路由是同步的FastAPI会把它放到线程池里执行不会堵事件循环这点是对的但每次请求都重新做分词和生成没有批处理吞吐上不去也没有任何观测。我们一步步优化。4.3 加入动态批处理动态批处理的核心是维护一个请求队列后台有个worker不断从队列取请求攒到一定数量或超时后一起推理再把结果分发回去。import asyncio import time from collections import deque class BatchProcessor: def __init__(self, model, tokenizer, max_batch8, max_wait0.05): self.model model self.tokenizer tokenizer self.max_batch max_batch self.max_wait max_wait self.queue deque() self.lock asyncio.Lock() async def submit(self, text, max_new_tokens): future asyncio.get_event_loop().create_future() async with self.lock: self.queue.append((text, max_new_tokens, future)) return await future async def worker(self): while True: if not self.queue: await asyncio.sleep(0.001) continue batch [] start time.time() while self.queue and len(batch) self.max_batch: batch.append(self.queue.popleft()) if time.time() - start self.max_wait: break await self.process(batch) async def process(self, batch): texts [item[0] for item in batch] max_tokens max(item[1] for item in batch) inputs self.tokenizer( texts, return_tensorspt, paddingTrue, truncationTrue ).to(self.model.device) with torch.no_grad(): outputs self.model.generate( **inputs, max_new_tokensmax_tokens, do_sampleFalse ) for i, item in enumerate(batch): result self.tokenizer.decode(outputs[i], skip_special_tokensTrue) item[2].set_result(result)这里有几个关键点。max_batch和max_wait是两个核心参数前者决定吞吐上限后者决定延迟上限。padding是必须的因为一个batch里的序列长度不一样要补齐到相同长度。但padding会浪费一些计算所以序列长度差异特别大的时候可以考虑按长度分桶。4.4 参数计算batch size和显存的关系显存占用可以粗略估算。以推理为例主要分三块模型权重、激活值、以及框架开销。模型权重参数量 × 精度字节数。7B模型FP16约14GB。激活值跟batch size、序列长度、隐藏层维度相关。粗略估算每层激活值约等于 batch_size × seq_len × hidden_dim × 精度字节数 × 一个系数。这个系数跟具体结构有关实测更准。假设你的卡有24GB显存模型权重占14GB剩10GB给激活值和框架开销。框架开销留2GB剩8GB给激活值。如果单条序列激活值占0.5GB那batch size最多到16左右。但这只是估算实际要留余量因为显存碎片和峰值占用可能更高。我的经验是先按估算值的一半设跑起来看实际显存占用再逐步往上加加到接近上限但留10%余量为止。4.5 观测接入用prometheus-client暴露指标在推理前后打点from prometheus_client import Histogram, Counter, Gauge REQUEST_LATENCY Histogram(inference_latency_seconds, 推理延迟) REQUEST_COUNT Counter(inference_requests_total, 请求总数) BATCH_SIZE Histogram(inference_batch_size, 批大小分布) QUEUE_LENGTH Gauge(inference_queue_length, 队列长度)在process函数里记录延迟和批大小在submit时更新队列长度。然后起一个/metrics端点暴露出去Prometheus定时抓取Grafana画图。这套东西搭起来之后你就能看到延迟分布、批大小分布、队列积压情况。哪个环节有问题一眼就能看出来。5. 常见问题与排查技巧实录5.1 显存溢出OOM怎么排查OOM是最常见的问题。排查思路是先确认是加载时OOM还是推理时OOM。加载时OOM说明模型权重就超了得换量化版本或换卡。推理时OOM说明激活值超了得降batch size或降序列长度。一个容易被忽视的点是显存碎片。长时间运行的服务反复分配释放显存可能产生碎片导致明明总显存够用却分配不出来。解决办法是设置PYTORCH_CUDA_ALLOC_CONF环境变量开启更高效的内存分配策略。提示export PYTORCH_CUDA_ALLOC_CONFmax_split_size_mb:128能缓解碎片问题但会略微增加内存占用。5.2 延迟忽高忽低延迟不稳定通常是几个原因批处理等待时间波动、GPU被其他进程抢占、或者请求长度差异大导致计算量波动。先看批大小分布如果批大小波动很大说明请求不均匀可以适当调大max_wait让批次更稳定。再看GPU是否有其他进程用nvidia-smi确认。如果是请求长度问题考虑按长度分桶把长度相近的请求放一批。5.3 吞吐上不去但GPU利用率低这个现象说明瓶颈不在GPU在CPU或IO。常见原因预处理太慢、分词器是纯Python实现、数据在CPU和GPU之间频繁拷贝。排查方法是用profiler看时间花在哪。PyTorch自带的profiler能给出每个操作的耗时。如果发现大量时间花在数据搬运上考虑用pin_memory和异步拷贝。如果花在分词上换更快的分词器实现。5.4 常见问题速查表问题现象可能原因排查方向解决思路加载时OOM模型权重超显存看模型参数量和精度用量化版本或换卡推理时OOM激活值超显存看batch size和序列长度降batch或降长度延迟波动大批处理不稳定看批大小分布调max_wait或分桶吞吐低GPU闲CPU瓶颈用profiler定位优化预处理或换实现并发上不去事件循环阻塞看路由是否同步用线程池或同步路由显存碎片反复分配释放看长时间运行表现设置分配策略环境变量5.5 几个我踩过的坑第一个坑是分词器的padding方向。有些模型是左padding有些是右padding搞错了生成结果会完全不对。这个一定要看模型文档确认。第二个坑是torch.no_grad()忘了加。推理时如果没加PyTorch会构建计算图显存占用暴涨速度也慢。这个错误很隐蔽因为结果是对的只是慢和费显存。第三个坑是模型eval模式忘了切。有些层在训练和推理时行为不同比如dropout和batchnorm。忘了model.eval()会导致结果不稳定。第四个坑是并发测试时用了同步客户端。测服务并发能力时客户端本身如果是同步阻塞的测出来的并发数是不准的。要用异步客户端或者多进程压测。6. 从能跑到好用进一步优化的方向6.1 量化精度和资源的权衡量化是降低显存、提升速度的有效手段。INT8量化通常精度损失很小INT4损失大一些但显存占用大幅下降。量化的方式有训练后量化和量化感知训练前者简单后者精度更好但需要训练。实操上如果只是推理训练后量化基本够用。要注意的是不是所有层都适合量化有些敏感层量化后精度掉得厉害需要混合精度策略。6.2 模型编译让计算图更高效PyTorch 2.0之后引入了编译功能能把模型的计算图优化减少Python开销和冗余计算。对于推理服务开启编译通常能带来明显的速度提升代价是首次编译需要一些时间。用法很简单model torch.compile(model)。但要注意不是所有模型都能顺利编译遇到不支持的算子会回退到eager模式。编译后的模型第一次推理会慢因为要编译之后才快。6.3 多卡与多实例单卡跑不下或者吞吐不够时考虑多卡。多卡有两种模式模型并行和数据并行。模型并行是把模型切到多张卡上适合单卡放不下的大模型数据并行是每张卡跑一份完整模型适合吞吐不够的场景。数据并行实现简单起多个服务实例前面挂个负载均衡就行。模型并行复杂一些需要框架支持。我一般优先考虑数据并行因为简单可靠扩展性也好。6.4 服务治理健康检查与优雅退出生产环境的服务健康检查和优雅退出是必须的。健康检查让负载均衡知道实例是否可用优雅退出保证正在处理的请求能完成再关闭。FastAPI里加健康检查很简单一个返回200的端点就行。优雅退出需要监听关闭信号停止接收新请求等队列清空后再退出。这个逻辑要自己写框架不会自动帮你做。7. 我个人的一些体会搭这套东西的过程中我最大的感受是AI工程的难点往往不在AI本身而在工程。模型是现成的框架是成熟的真正花时间的是环境配置、性能调优、问题排查这些“脏活累活”。但恰恰是这些活决定了一个AI服务能不能真正上线、能不能扛住流量。另一个体会是观测的重要性怎么强调都不过分。我早期做服务的时候出了问题全靠猜效率极低。后来把观测体系搭起来很多问题几分钟就能定位。数据不会骗人有数据就有方向。最后分享一个小技巧每次改动只改一个变量改完立刻测。AI工程涉及的因素太多一次改好几个地方出了问题根本不知道是哪个改动导致的。控制变量法虽然笨但最有效。
