1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件调几个API然后跑通一个Demo就觉得自己已经掌握了。我刚开始也是这么想的直到有一次线上推理服务在高峰期直接雪崩日志里全是显存溢出和请求超时我才意识到——那些封装好的高级接口在关键时刻根本救不了你。你连模型为什么吃这么多显存、请求为什么排队、算子为什么慢都说不清楚怎么排查ai-engineering-from-scratch这个标题核心不是教你从零训练一个GPT而是教你从零搭建一套能跑、能调、能扛的AI工程链路。它解决的是“只会调包、不懂底层”的普遍痛点适合那些已经会用PyTorch或TensorFlow写模型但一遇到部署、优化、并发就抓瞎的开发者。说白了就是从“能跑通”到“跑得稳、跑得快、跑得省”之间的那段路。我见过太多团队模型指标刷得很漂亮一上线就翻车。问题往往不在算法本身而在工程细节数据加载成了瓶颈、推理没有做批处理、显存碎片化严重、服务没有做限流和降级。这些坑你调包是永远调不出来的。所以这篇内容我会按照一个真实项目的推进节奏把从环境搭建到性能调优的完整链路拆开讲每一步都告诉你为什么这么做以及我踩过哪些坑。2. 环境搭建别让CUDA版本成为你的第一个噩梦2.1 驱动、运行时与框架的三角关系AI工程的第一步永远是环境。很多人觉得装个PyTorch有什么难的pip install一行命令的事。但实际情况是你的显卡驱动版本、CUDA运行时版本、PyTorch编译时链接的CUDA版本这三者必须匹配。我见过最离谱的案例是驱动是470CUDA运行时是11.4但PyTorch装的是cu118的包结果就是torch.cuda.is_available()返回True但一跑矩阵乘法就报“no kernel image is available”。这个错误信息极其误导它让你以为是显卡坏了其实是版本不匹配。正确的做法是先确定你的驱动版本然后查NVIDIA官方的兼容性表确定驱动支持的最高CUDA运行时版本再根据这个版本去选PyTorch的安装命令。比如驱动是515最高支持CUDA 11.7那你就应该装cu117的PyTorch而不是盲目装最新的cu121。这里有个小技巧用nvidia-smi看驱动版本用nvcc --version看CUDA编译器版本用torch.version.cuda看PyTorch链接的CUDA版本三个命令的输出必须能对上。注意不要用conda install pytorch这种命令conda的依赖解析经常会把CUDA运行时降级到一个很老的版本导致你明明装了新驱动却只能用老特性。我推荐用pip配合官方提供的index-url来装版本控制更精确。2.2 虚拟环境与依赖锁定的实操细节虚拟环境这件事很多人用venv或者conda create就完事了但AI项目的依赖树极其复杂torch、transformers、accelerate、datasets这些库之间经常有版本冲突。我的习惯是用conda创建基础环境但所有Python包都用pip装并且装完之后立刻用pip freeze requirements.txt锁定版本。注意pip freeze会把所有间接依赖都写进去这其实是好事因为AI项目的间接依赖经常出问题。还有一个坑是transformers库的版本和tokenizers库的版本是强绑定的你单独升级其中一个另一个就会报错。我一般会在requirements.txt里把这两个库的版本都写死比如transformers4.36.2和tokenizers0.15.0。另外如果你要用flash-attention这种需要编译的库一定要先装好CUDA toolkit和ninja否则编译过程会卡住而且报错信息很不友好。2.3 容器化从“在我机器上能跑”到“在哪都能跑”如果你只是本地开发虚拟环境够了。但一旦要部署到服务器或者集群容器化是绕不过去的。我推荐用NVIDIA的官方基础镜像比如nvcr.io/nvidia/pytorch:23.10-py3这个镜像里已经配好了驱动兼容的CUDA和cuDNN你只需要在上面装你的业务依赖。但要注意这个镜像的Python版本可能和你本地不一样所以requirements.txt里的包版本要重新验证一遍。构建镜像的时候有一个细节很容易被忽略pip install的时候要加--no-cache-dir否则镜像体积会膨胀好几个G。另外把requirements.txt的复制和安装放在Dockerfile的前面这样只要依赖没变构建缓存就能复用每次改代码只需要重新构建最后一层速度会快很多。我实测下来一个包含PyTorch和transformers的镜像优化前构建要15分钟优化后只要2分钟。3. 数据管道模型还没开始训瓶颈就已经出现了3.1 DataLoader的num_workers到底设多少DataLoader的num_workers参数很多人随手设成4或者8觉得越大越好。但实际情况是这个值设得不对不仅不会加速反而会拖慢训练。原因在于每个worker都是一个独立的进程它们需要把数据从磁盘读到内存再做预处理然后通过共享内存传给主进程。如果num_workers太大进程间切换的开销会超过并行读取的收益如果太小GPU就会等数据利用率上不去。我的经验是先看你的CPU核心数num_workers一般设成CPU核心数的50%到75%。比如16核的机器设8到12比较合适。但这不是绝对的还要看你的数据预处理复杂度。如果预处理很轻比如只是归一化可以设大一点如果预处理很重比如要做数据增强、分词那每个worker的耗时很长设太大反而会导致内存爆炸。我一般会先用num_workers0跑一个epoch看GPU利用率如果低于70%再逐步增加num_workers直到GPU利用率稳定在90%以上。还有一个隐藏的坑是pin_memoryTrue这个选项。它会把数据锁在页锁定内存里加速CPU到GPU的传输。但如果你用的是多进程DataLoader并且pin_memoryTrue那么每个worker都会尝试锁内存可能导致系统内存不足。我一般会在num_workers0的时候开pin_memory但会监控系统的可用内存如果低于20%就关掉它。3.2 数据预处理的流水线设计数据预处理最忌讳的是“一次性全做完”。我见过有人把整个数据集读进内存做完所有增强再开始训练。这样做有两个问题一是内存扛不住二是每个epoch的数据都是一样的没有随机性。正确的做法是把预处理拆成“轻量级在线处理”和“重量级离线处理”两部分。轻量级的操作比如归一化、张量转换、简单的随机裁剪放在Dataset的__getitem__里做这样每个epoch都能有不同的随机结果。重量级的操作比如分词、特征提取、图像解码如果耗时很长可以提前离线做好存成内存映射文件或者LMDB数据库训练时直接读取。我实测过把BERT分词从在线改成离线训练速度提升了3倍多因为分词是CPU密集型的在线做会严重拖慢数据加载。提示如果你用HuggingFace的datasets库它自带内存映射功能load_from_disk之后的数据集不会全部加载到内存而是按需读取。这个特性在处理大规模数据时非常有用但要注意内存映射文件的读取速度受磁盘IO限制最好放在SSD上。3.3 数据并行与分布式采样的坑当你用DistributedDataParallel做多卡训练时数据采样必须用DistributedSampler否则每张卡都会看到相同的数据训练就失去了并行的意义。但DistributedSampler有一个坑它默认会把数据集补齐到能被world_size整除的长度补齐的部分是重复的样本。如果你没设置drop_lastTrue这些重复样本会影响梯度更新的准确性。另外DistributedSampler的shuffle参数在每个epoch开始时需要手动调用set_epoch(epoch)否则随机种子不会变每个epoch的数据顺序都一样。这个细节很容易被忽略但影响很大尤其是当你的数据集本身有序的时候模型会学到错误的顺序信息。我一般会在训练循环的开头加上train_sampler.set_epoch(epoch)确保每个epoch的采样都是随机的。4. 模型训练从能收敛到收敛得快、收敛得稳4.1 混合精度训练省显存不是唯一目的混合精度训练AMP现在已经是标配了但很多人只知道它能省显存不知道它还能加速。原理是FP16的矩阵乘法在Tensor Core上的吞吐量是FP32的8倍所以计算密集型的层比如卷积、全连接用FP16跑速度提升非常明显。但要注意不是所有操作都适合FP16比如softmax、layer norm、loss计算这些对数值范围敏感的操作必须用FP32否则会溢出或者下溢。PyTorch的torch.cuda.amp提供了自动混合精度你只需要用autocast上下文管理器包住前向传播用GradScaler包住反向传播和优化器更新。但这里有一个坑GradScaler会在检测到梯度溢出时自动缩小loss scale如果溢出的次数太多scale会变得很小导致梯度更新几乎停滞。我一般会监控scaler.get_scale()的值如果它持续下降说明模型里有数值不稳定的操作需要检查是不是某个层的初始化有问题或者学习率设得太大了。还有一个细节是用AMP的时候BatchNorm层最好保持FP32因为它的running mean和running var对精度很敏感。PyTorch的autocast默认会把BatchNorm排除在FP16之外但如果你自己写了自定义的norm层就要手动加torch.cuda.amp.custom_fwd(cast_inputstorch.float32)装饰器。4.2 梯度累积与学习率调度的配合梯度累积是模拟大batch的常用技巧但很多人只改了累积步数忘了同步调整学习率。假设你把batch size从32扩大到128累积4步那学习率也应该相应放大但放大多少是有讲究的。线性缩放规则learning rate乘以累积步数在SGD上效果不错但在Adam上往往会导致训练不稳定因为Adam的自适应学习率本身就会根据梯度的大小调整步长。我的做法是先用线性缩放规则设一个初始学习率然后跑几百步看loss曲线。如果loss震荡得很厉害就把学习率降一半如果loss下降得太慢就稍微加一点。另外warmup步数也要跟着调整累积步数越大warmup应该越长因为等效batch size变大了模型需要更多步来稳定统计量。我一般会把warmup步数设成总步数的5%到10%具体看数据集大小。注意梯度累积的时候loss要除以累积步数否则梯度会放大。但如果你用的是GradScaler除法要在scaler.scale(loss)之前做否则scale会乱掉。这个顺序很容易搞反我踩过好几次坑。4.3 检查点保存与恢复的工程细节训练大模型的时候检查点保存是个容易被忽视但极其重要的环节。我见过有人训练了三天结果因为磁盘满了最后一个检查点没保存下来前功尽弃。所以检查点保存必须做几件事第一保存到临时目录写完再原子性地重命名到目标路径防止写入过程中断导致文件损坏第二保存的时候要包含优化器状态、学习率调度器状态、当前的epoch和step否则恢复训练的时候会从头开始第三定期清理旧的检查点但至少保留最近三个防止最新的检查点损坏。还有一个坑是如果你用了DistributedDataParallel保存检查点的时候只需要在主进程rank 0保存其他进程等待。但恢复的时候所有进程都要加载同一个检查点并且要确保加载后的模型参数在各个进程之间是一致的。我一般会用torch.save保存model.module.state_dict()而不是model.state_dict()因为后者会带上module.前缀加载的时候容易出错。5. 推理部署训练只是开始上线才是大考5.1 模型导出ONNX还是TorchScript模型训练完之后下一步就是导出成推理引擎能用的格式。常见的选择有ONNX和TorchScript。ONNX的优点是跨框架你可以在TensorRT、OpenVINO、ONNX Runtime上跑缺点是有些动态控制流比如if-else、循环不支持导出的时候会报错。TorchScript的优点是和PyTorch无缝集成支持动态控制流缺点是只能在PyTorch生态里跑跨框架能力弱。我的建议是如果你的模型结构比较规整比如ResNet、BERT优先用ONNX因为后续可以用TensorRT做量化加速性能提升非常明显。如果你的模型有复杂的控制流比如Beam Search、动态解码那就用TorchScript或者把控制流部分留在Python里只把计算密集的部分导出。我实测过一个BERT-base模型ONNX Runtime比原生PyTorch推理快1.5倍TensorRT FP16比ONNX Runtime再快2倍。导出ONNX的时候有一个坑是dynamic_axes的设置。如果你的模型支持变长输入比如不同长度的文本必须把序列长度维度设成动态的否则导出的模型只能处理固定长度。但设成动态之后TensorRT的优化会受限因为很多kernel fusion需要静态shape。我一般会导出两个版本一个静态shape的用于固定长度场景一个动态shape的用于变长场景根据实际请求选择。5.2 批处理与动态批处理的实现推理服务的吞吐量很大程度上取决于批处理策略。最简单的做法是固定batch size攒够一批就推理。但这样有两个问题一是延迟高如果请求稀疏用户要等很久才能攒够一批二是资源浪费如果请求突然增多固定batch size会导致排队。动态批处理Dynamic Batching解决了这个问题它设置一个最大batch size和一个最大等待时间在这个时间窗口内有多少请求就处理多少攒够最大batch size就立即处理。实现动态批处理可以用NVIDIA Triton Inference Server它内置了动态批处理功能你只需要在配置里设max_batch_size和max_queue_delay_microseconds。但Triton的配置比较复杂如果你不想引入这么重的依赖也可以自己用Python的asyncio或者queue实现。我一般会用一个后台线程做批处理调度主线程收到请求后把输入放进队列后台线程每隔几毫秒检查一次队列如果队列长度达到阈值或者等待时间超时就取出当前所有请求拼成一个batch调用模型推理然后把结果分发给各个请求。提示动态批处理的最大batch size不是越大越好因为显存占用是随batch size线性增长的。我一般会先用torch.cuda.max_memory_allocated()测一下单个样本的显存占用然后根据可用显存算出最大batch size再留20%的余量给中间激活值。5.3 服务限流与降级策略线上服务最怕的就是突发流量。如果没有限流所有请求都涌进来显存瞬间打满服务直接挂掉。限流的策略有很多我常用的是令牌桶算法以固定的速率往桶里放令牌每个请求消耗一个令牌桶空了就拒绝请求或者排队。令牌桶的好处是允许一定的突发流量只要桶里有足够的令牌。但限流只是第一道防线如果流量持续超过服务能力还需要降级。降级的策略包括降低推理精度FP16换成INT8、缩小模型用蒸馏后的小模型、减少批处理等待时间牺牲吞吐保延迟。我一般会在服务里加一个监控模块实时统计QPS、延迟、显存占用当这些指标超过阈值时自动触发降级。比如当显存占用超过90%时自动把最大batch size减半当P99延迟超过500ms时自动关闭动态批处理改成单条推理。还有一个细节是降级之后要有恢复机制。不能一直降级否则服务质量永远上不去。我一般会设置一个冷却时间比如降级后5分钟内不再触发降级5分钟后如果指标恢复正常就逐步恢复原来的配置。这个恢复过程要慢比如每次只把batch size增加10%观察一段时间再继续增加防止一下子恢复导致再次雪崩。6. 性能调优那些文档里不会写的实测经验6.1 显存碎片化与内存池配置PyTorch的显存分配器默认用的是缓存分配器它会缓存已经分配的显存块避免频繁向CUDA申请和释放。但这个缓存机制有一个问题如果模型训练过程中有不同大小的显存申请缓存块会变得碎片化导致明明有足够的空闲显存却分配不出一块连续的大显存。这个现象在变长输入的场景下特别明显因为每个batch的序列长度不一样中间激活值的大小也不一样。解决碎片化的方法有两个一是设置PYTORCH_CUDA_ALLOC_CONF环境变量把max_split_size_mb设成一个合适的值比如128这样分配器会尽量把大的空闲块保留下来而不是切碎。二是用torch.cuda.memory_summary()定期检查显存碎片率如果碎片率超过30%就重启服务。我实测过设置max_split_size_mb128之后显存碎片率从40%降到了10%以下服务稳定性提升了很多。还有一个坑是如果你用了torch.cuda.empty_cache()它会释放所有未使用的缓存显存但这会导致下一次分配显存时重新向CUDA申请速度很慢。所以empty_cache()不能频繁调用我一般只在OOM之后调用一次然后重新加载模型。6.2 算子融合与图优化PyTorch 2.0引入了torch.compile它可以把模型的计算图编译成优化的kernel自动做算子融合。我实测下来torch.compile在Transformer类模型上的加速效果非常明显推理速度能提升30%到50%。但torch.compile有一个坑它第一次运行的时候会花很长时间做编译如果服务是冷启动的第一个请求的延迟会非常高。所以我一般会在服务启动之后先用一些样例数据跑几次预热让torch.compile完成编译然后再开始接收真实请求。另外torch.compile的mode参数有default、reduce-overhead、max-autotune三个选项。default是默认的优化级别reduce-overhead会减少Python解释器的开销适合小模型max-autotune会花更多时间搜索最优的kernel配置适合大模型。我一般会在离线测试的时候用max-autotune找到最优配置后再用default模式上线因为max-autotune的编译时间太长了。6.3 监控指标与告警阈值设定AI服务的监控不能只看CPU和内存还要看GPU利用率、显存占用、推理延迟、批处理大小、队列长度。我一般会用Prometheus采集这些指标用Grafana做可视化。关键的告警阈值包括GPU利用率持续低于30%说明有瓶颈、显存占用超过90%有OOM风险、P99延迟超过1秒用户体验差、队列长度超过100服务过载。但阈值不是拍脑袋定的要根据实际压测结果来。我一般会先用Locust或者wrk做压力测试逐步增加QPS观察各个指标的变化找到服务开始劣化的拐点然后把告警阈值设在拐点的80%处。比如压测发现QPS到500的时候延迟开始飙升那告警阈值就设在400。这样可以在服务真正过载之前就发出告警留出处理时间。还有一个经验是监控指标要分维度。比如延迟要分P50、P90、P99不能只看平均值因为平均值会被大量快请求拉低掩盖了慢请求的问题。显存要分已分配、已缓存、空闲不能只看总占用因为缓存显存是可以释放的已分配显存才是真正被模型占用的。这些细节文档里不会写但线上排查问题的时候非常有用。7. 我踩过的那些坑真实案例复盘7.1 一次OOM引发的连锁反应有一次我们的推理服务在高峰期突然OOM然后Kubernetes自动重启了Pod。但重启之后服务并没有恢复而是陷入了“启动-预热-OOM-重启”的循环。排查了半天发现原因是预热阶段用的样例数据比真实请求长很多导致预热时的显存占用比正常推理高出一大截。Pod重启后预热阶段就OOM了根本没机会处理真实请求。这个坑的教训是预热数据必须和真实请求的分布一致。我后来改成了从线上流量里采样一批真实请求做预热并且把预热时的最大batch size设成正常值的一半确保预热不会OOM。另外Kubernetes的重启策略也要调整不能无限重启我设了restartPolicy: OnFailure和backoffLimit: 3重启3次还失败就停止避免无限循环消耗资源。7.2 动态批处理导致的延迟毛刺还有一次我们上线了动态批处理吞吐量确实上去了但P99延迟出现了周期性的毛刺。查了很久发现原因是动态批处理的等待时间设成了10毫秒但有些请求的处理时间特别长比如输入特别长导致这一批的处理时间超过了10毫秒下一批请求就要等更久形成了正反馈延迟越来越高。解决办法是给动态批处理加一个“最大处理时间”限制。如果一批请求的处理时间超过了阈值比如50毫秒就把这一批拆成两批先返回已经处理完的部分剩下的放到下一批。这样虽然会增加一些批处理次数但能避免延迟毛刺。另外对于特别长的请求可以单独走一个“慢通道”不参与动态批处理直接单条推理避免拖累其他请求。7.3 版本升级导致的精度下降最隐蔽的一个坑是我们把PyTorch从1.12升级到2.0之后模型的推理精度下降了0.5个百分点。这个下降很小一开始没注意到后来做A/B测试才发现。排查后发现PyTorch 2.0的torch.compile默认开启了一些算子融合其中有一个融合操作在数值上做了近似导致精度损失。这个坑的教训是版本升级之后一定要做精度回归测试。我后来在CI流程里加了一个步骤每次升级框架版本都用固定的测试集跑一遍推理对比输出和之前版本的差异如果差异超过阈值就阻止升级。另外如果精度要求极高可以在torch.compile里关掉有损的融合操作比如设置modedefault并且禁用fuse_ops虽然会损失一些性能但能保证精度。8. 写给想从零开始做AI工程的人如果你看到这里说明你对AI工程是真的有兴趣。我想说的是从零开始做AI工程最难的不是写代码而是建立一套系统性的思维方式。你要时刻问自己这个操作的瓶颈在哪里是计算、内存、IO还是网络这个参数为什么设这个值有没有理论依据或者实测数据支撑这个方案在流量翻倍的时候还能不能扛住我刚开始做AI工程的时候也是到处找“最佳实践”但后来发现所谓的“最佳实践”都是有前提的。别人的batch size设128是因为他们的显存有80G别人的学习率设1e-4是因为他们的数据集有100万条。你直接抄过来大概率会翻车。所以我建议你养成一个习惯每做一个决策都记录下当时的上下文和理由过一段时间回头看你会发现很多决策其实可以做得更好。最后分享一个我常用的调试技巧当你遇到一个性能问题不要急着改代码先用py-spy或者torch.profiler做一次性能剖析找到真正的瓶颈。我见过太多人凭直觉优化结果优化了不重要的部分真正的瓶颈还在那里。数据不会骗人让数据告诉你该优化哪里。
从零手搓AI工程：避开调包陷阱，掌握底层部署与性能调优
1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件调几个API然后跑通一个Demo就觉得自己已经掌握了。我刚开始也是这么想的直到有一次线上推理服务在高峰期直接雪崩日志里全是显存溢出和请求超时我才意识到——那些封装好的高级接口在关键时刻根本救不了你。你连模型为什么吃这么多显存、请求为什么排队、算子为什么慢都说不清楚怎么排查ai-engineering-from-scratch这个标题核心不是教你从零训练一个GPT而是教你从零搭建一套能跑、能调、能扛的AI工程链路。它解决的是“只会调包、不懂底层”的普遍痛点适合那些已经会用PyTorch或TensorFlow写模型但一遇到部署、优化、并发就抓瞎的开发者。说白了就是从“能跑通”到“跑得稳、跑得快、跑得省”之间的那段路。我见过太多团队模型指标刷得很漂亮一上线就翻车。问题往往不在算法本身而在工程细节数据加载成了瓶颈、推理没有做批处理、显存碎片化严重、服务没有做限流和降级。这些坑你调包是永远调不出来的。所以这篇内容我会按照一个真实项目的推进节奏把从环境搭建到性能调优的完整链路拆开讲每一步都告诉你为什么这么做以及我踩过哪些坑。2. 环境搭建别让CUDA版本成为你的第一个噩梦2.1 驱动、运行时与框架的三角关系AI工程的第一步永远是环境。很多人觉得装个PyTorch有什么难的pip install一行命令的事。但实际情况是你的显卡驱动版本、CUDA运行时版本、PyTorch编译时链接的CUDA版本这三者必须匹配。我见过最离谱的案例是驱动是470CUDA运行时是11.4但PyTorch装的是cu118的包结果就是torch.cuda.is_available()返回True但一跑矩阵乘法就报“no kernel image is available”。这个错误信息极其误导它让你以为是显卡坏了其实是版本不匹配。正确的做法是先确定你的驱动版本然后查NVIDIA官方的兼容性表确定驱动支持的最高CUDA运行时版本再根据这个版本去选PyTorch的安装命令。比如驱动是515最高支持CUDA 11.7那你就应该装cu117的PyTorch而不是盲目装最新的cu121。这里有个小技巧用nvidia-smi看驱动版本用nvcc --version看CUDA编译器版本用torch.version.cuda看PyTorch链接的CUDA版本三个命令的输出必须能对上。注意不要用conda install pytorch这种命令conda的依赖解析经常会把CUDA运行时降级到一个很老的版本导致你明明装了新驱动却只能用老特性。我推荐用pip配合官方提供的index-url来装版本控制更精确。2.2 虚拟环境与依赖锁定的实操细节虚拟环境这件事很多人用venv或者conda create就完事了但AI项目的依赖树极其复杂torch、transformers、accelerate、datasets这些库之间经常有版本冲突。我的习惯是用conda创建基础环境但所有Python包都用pip装并且装完之后立刻用pip freeze requirements.txt锁定版本。注意pip freeze会把所有间接依赖都写进去这其实是好事因为AI项目的间接依赖经常出问题。还有一个坑是transformers库的版本和tokenizers库的版本是强绑定的你单独升级其中一个另一个就会报错。我一般会在requirements.txt里把这两个库的版本都写死比如transformers4.36.2和tokenizers0.15.0。另外如果你要用flash-attention这种需要编译的库一定要先装好CUDA toolkit和ninja否则编译过程会卡住而且报错信息很不友好。2.3 容器化从“在我机器上能跑”到“在哪都能跑”如果你只是本地开发虚拟环境够了。但一旦要部署到服务器或者集群容器化是绕不过去的。我推荐用NVIDIA的官方基础镜像比如nvcr.io/nvidia/pytorch:23.10-py3这个镜像里已经配好了驱动兼容的CUDA和cuDNN你只需要在上面装你的业务依赖。但要注意这个镜像的Python版本可能和你本地不一样所以requirements.txt里的包版本要重新验证一遍。构建镜像的时候有一个细节很容易被忽略pip install的时候要加--no-cache-dir否则镜像体积会膨胀好几个G。另外把requirements.txt的复制和安装放在Dockerfile的前面这样只要依赖没变构建缓存就能复用每次改代码只需要重新构建最后一层速度会快很多。我实测下来一个包含PyTorch和transformers的镜像优化前构建要15分钟优化后只要2分钟。3. 数据管道模型还没开始训瓶颈就已经出现了3.1 DataLoader的num_workers到底设多少DataLoader的num_workers参数很多人随手设成4或者8觉得越大越好。但实际情况是这个值设得不对不仅不会加速反而会拖慢训练。原因在于每个worker都是一个独立的进程它们需要把数据从磁盘读到内存再做预处理然后通过共享内存传给主进程。如果num_workers太大进程间切换的开销会超过并行读取的收益如果太小GPU就会等数据利用率上不去。我的经验是先看你的CPU核心数num_workers一般设成CPU核心数的50%到75%。比如16核的机器设8到12比较合适。但这不是绝对的还要看你的数据预处理复杂度。如果预处理很轻比如只是归一化可以设大一点如果预处理很重比如要做数据增强、分词那每个worker的耗时很长设太大反而会导致内存爆炸。我一般会先用num_workers0跑一个epoch看GPU利用率如果低于70%再逐步增加num_workers直到GPU利用率稳定在90%以上。还有一个隐藏的坑是pin_memoryTrue这个选项。它会把数据锁在页锁定内存里加速CPU到GPU的传输。但如果你用的是多进程DataLoader并且pin_memoryTrue那么每个worker都会尝试锁内存可能导致系统内存不足。我一般会在num_workers0的时候开pin_memory但会监控系统的可用内存如果低于20%就关掉它。3.2 数据预处理的流水线设计数据预处理最忌讳的是“一次性全做完”。我见过有人把整个数据集读进内存做完所有增强再开始训练。这样做有两个问题一是内存扛不住二是每个epoch的数据都是一样的没有随机性。正确的做法是把预处理拆成“轻量级在线处理”和“重量级离线处理”两部分。轻量级的操作比如归一化、张量转换、简单的随机裁剪放在Dataset的__getitem__里做这样每个epoch都能有不同的随机结果。重量级的操作比如分词、特征提取、图像解码如果耗时很长可以提前离线做好存成内存映射文件或者LMDB数据库训练时直接读取。我实测过把BERT分词从在线改成离线训练速度提升了3倍多因为分词是CPU密集型的在线做会严重拖慢数据加载。提示如果你用HuggingFace的datasets库它自带内存映射功能load_from_disk之后的数据集不会全部加载到内存而是按需读取。这个特性在处理大规模数据时非常有用但要注意内存映射文件的读取速度受磁盘IO限制最好放在SSD上。3.3 数据并行与分布式采样的坑当你用DistributedDataParallel做多卡训练时数据采样必须用DistributedSampler否则每张卡都会看到相同的数据训练就失去了并行的意义。但DistributedSampler有一个坑它默认会把数据集补齐到能被world_size整除的长度补齐的部分是重复的样本。如果你没设置drop_lastTrue这些重复样本会影响梯度更新的准确性。另外DistributedSampler的shuffle参数在每个epoch开始时需要手动调用set_epoch(epoch)否则随机种子不会变每个epoch的数据顺序都一样。这个细节很容易被忽略但影响很大尤其是当你的数据集本身有序的时候模型会学到错误的顺序信息。我一般会在训练循环的开头加上train_sampler.set_epoch(epoch)确保每个epoch的采样都是随机的。4. 模型训练从能收敛到收敛得快、收敛得稳4.1 混合精度训练省显存不是唯一目的混合精度训练AMP现在已经是标配了但很多人只知道它能省显存不知道它还能加速。原理是FP16的矩阵乘法在Tensor Core上的吞吐量是FP32的8倍所以计算密集型的层比如卷积、全连接用FP16跑速度提升非常明显。但要注意不是所有操作都适合FP16比如softmax、layer norm、loss计算这些对数值范围敏感的操作必须用FP32否则会溢出或者下溢。PyTorch的torch.cuda.amp提供了自动混合精度你只需要用autocast上下文管理器包住前向传播用GradScaler包住反向传播和优化器更新。但这里有一个坑GradScaler会在检测到梯度溢出时自动缩小loss scale如果溢出的次数太多scale会变得很小导致梯度更新几乎停滞。我一般会监控scaler.get_scale()的值如果它持续下降说明模型里有数值不稳定的操作需要检查是不是某个层的初始化有问题或者学习率设得太大了。还有一个细节是用AMP的时候BatchNorm层最好保持FP32因为它的running mean和running var对精度很敏感。PyTorch的autocast默认会把BatchNorm排除在FP16之外但如果你自己写了自定义的norm层就要手动加torch.cuda.amp.custom_fwd(cast_inputstorch.float32)装饰器。4.2 梯度累积与学习率调度的配合梯度累积是模拟大batch的常用技巧但很多人只改了累积步数忘了同步调整学习率。假设你把batch size从32扩大到128累积4步那学习率也应该相应放大但放大多少是有讲究的。线性缩放规则learning rate乘以累积步数在SGD上效果不错但在Adam上往往会导致训练不稳定因为Adam的自适应学习率本身就会根据梯度的大小调整步长。我的做法是先用线性缩放规则设一个初始学习率然后跑几百步看loss曲线。如果loss震荡得很厉害就把学习率降一半如果loss下降得太慢就稍微加一点。另外warmup步数也要跟着调整累积步数越大warmup应该越长因为等效batch size变大了模型需要更多步来稳定统计量。我一般会把warmup步数设成总步数的5%到10%具体看数据集大小。注意梯度累积的时候loss要除以累积步数否则梯度会放大。但如果你用的是GradScaler除法要在scaler.scale(loss)之前做否则scale会乱掉。这个顺序很容易搞反我踩过好几次坑。4.3 检查点保存与恢复的工程细节训练大模型的时候检查点保存是个容易被忽视但极其重要的环节。我见过有人训练了三天结果因为磁盘满了最后一个检查点没保存下来前功尽弃。所以检查点保存必须做几件事第一保存到临时目录写完再原子性地重命名到目标路径防止写入过程中断导致文件损坏第二保存的时候要包含优化器状态、学习率调度器状态、当前的epoch和step否则恢复训练的时候会从头开始第三定期清理旧的检查点但至少保留最近三个防止最新的检查点损坏。还有一个坑是如果你用了DistributedDataParallel保存检查点的时候只需要在主进程rank 0保存其他进程等待。但恢复的时候所有进程都要加载同一个检查点并且要确保加载后的模型参数在各个进程之间是一致的。我一般会用torch.save保存model.module.state_dict()而不是model.state_dict()因为后者会带上module.前缀加载的时候容易出错。5. 推理部署训练只是开始上线才是大考5.1 模型导出ONNX还是TorchScript模型训练完之后下一步就是导出成推理引擎能用的格式。常见的选择有ONNX和TorchScript。ONNX的优点是跨框架你可以在TensorRT、OpenVINO、ONNX Runtime上跑缺点是有些动态控制流比如if-else、循环不支持导出的时候会报错。TorchScript的优点是和PyTorch无缝集成支持动态控制流缺点是只能在PyTorch生态里跑跨框架能力弱。我的建议是如果你的模型结构比较规整比如ResNet、BERT优先用ONNX因为后续可以用TensorRT做量化加速性能提升非常明显。如果你的模型有复杂的控制流比如Beam Search、动态解码那就用TorchScript或者把控制流部分留在Python里只把计算密集的部分导出。我实测过一个BERT-base模型ONNX Runtime比原生PyTorch推理快1.5倍TensorRT FP16比ONNX Runtime再快2倍。导出ONNX的时候有一个坑是dynamic_axes的设置。如果你的模型支持变长输入比如不同长度的文本必须把序列长度维度设成动态的否则导出的模型只能处理固定长度。但设成动态之后TensorRT的优化会受限因为很多kernel fusion需要静态shape。我一般会导出两个版本一个静态shape的用于固定长度场景一个动态shape的用于变长场景根据实际请求选择。5.2 批处理与动态批处理的实现推理服务的吞吐量很大程度上取决于批处理策略。最简单的做法是固定batch size攒够一批就推理。但这样有两个问题一是延迟高如果请求稀疏用户要等很久才能攒够一批二是资源浪费如果请求突然增多固定batch size会导致排队。动态批处理Dynamic Batching解决了这个问题它设置一个最大batch size和一个最大等待时间在这个时间窗口内有多少请求就处理多少攒够最大batch size就立即处理。实现动态批处理可以用NVIDIA Triton Inference Server它内置了动态批处理功能你只需要在配置里设max_batch_size和max_queue_delay_microseconds。但Triton的配置比较复杂如果你不想引入这么重的依赖也可以自己用Python的asyncio或者queue实现。我一般会用一个后台线程做批处理调度主线程收到请求后把输入放进队列后台线程每隔几毫秒检查一次队列如果队列长度达到阈值或者等待时间超时就取出当前所有请求拼成一个batch调用模型推理然后把结果分发给各个请求。提示动态批处理的最大batch size不是越大越好因为显存占用是随batch size线性增长的。我一般会先用torch.cuda.max_memory_allocated()测一下单个样本的显存占用然后根据可用显存算出最大batch size再留20%的余量给中间激活值。5.3 服务限流与降级策略线上服务最怕的就是突发流量。如果没有限流所有请求都涌进来显存瞬间打满服务直接挂掉。限流的策略有很多我常用的是令牌桶算法以固定的速率往桶里放令牌每个请求消耗一个令牌桶空了就拒绝请求或者排队。令牌桶的好处是允许一定的突发流量只要桶里有足够的令牌。但限流只是第一道防线如果流量持续超过服务能力还需要降级。降级的策略包括降低推理精度FP16换成INT8、缩小模型用蒸馏后的小模型、减少批处理等待时间牺牲吞吐保延迟。我一般会在服务里加一个监控模块实时统计QPS、延迟、显存占用当这些指标超过阈值时自动触发降级。比如当显存占用超过90%时自动把最大batch size减半当P99延迟超过500ms时自动关闭动态批处理改成单条推理。还有一个细节是降级之后要有恢复机制。不能一直降级否则服务质量永远上不去。我一般会设置一个冷却时间比如降级后5分钟内不再触发降级5分钟后如果指标恢复正常就逐步恢复原来的配置。这个恢复过程要慢比如每次只把batch size增加10%观察一段时间再继续增加防止一下子恢复导致再次雪崩。6. 性能调优那些文档里不会写的实测经验6.1 显存碎片化与内存池配置PyTorch的显存分配器默认用的是缓存分配器它会缓存已经分配的显存块避免频繁向CUDA申请和释放。但这个缓存机制有一个问题如果模型训练过程中有不同大小的显存申请缓存块会变得碎片化导致明明有足够的空闲显存却分配不出一块连续的大显存。这个现象在变长输入的场景下特别明显因为每个batch的序列长度不一样中间激活值的大小也不一样。解决碎片化的方法有两个一是设置PYTORCH_CUDA_ALLOC_CONF环境变量把max_split_size_mb设成一个合适的值比如128这样分配器会尽量把大的空闲块保留下来而不是切碎。二是用torch.cuda.memory_summary()定期检查显存碎片率如果碎片率超过30%就重启服务。我实测过设置max_split_size_mb128之后显存碎片率从40%降到了10%以下服务稳定性提升了很多。还有一个坑是如果你用了torch.cuda.empty_cache()它会释放所有未使用的缓存显存但这会导致下一次分配显存时重新向CUDA申请速度很慢。所以empty_cache()不能频繁调用我一般只在OOM之后调用一次然后重新加载模型。6.2 算子融合与图优化PyTorch 2.0引入了torch.compile它可以把模型的计算图编译成优化的kernel自动做算子融合。我实测下来torch.compile在Transformer类模型上的加速效果非常明显推理速度能提升30%到50%。但torch.compile有一个坑它第一次运行的时候会花很长时间做编译如果服务是冷启动的第一个请求的延迟会非常高。所以我一般会在服务启动之后先用一些样例数据跑几次预热让torch.compile完成编译然后再开始接收真实请求。另外torch.compile的mode参数有default、reduce-overhead、max-autotune三个选项。default是默认的优化级别reduce-overhead会减少Python解释器的开销适合小模型max-autotune会花更多时间搜索最优的kernel配置适合大模型。我一般会在离线测试的时候用max-autotune找到最优配置后再用default模式上线因为max-autotune的编译时间太长了。6.3 监控指标与告警阈值设定AI服务的监控不能只看CPU和内存还要看GPU利用率、显存占用、推理延迟、批处理大小、队列长度。我一般会用Prometheus采集这些指标用Grafana做可视化。关键的告警阈值包括GPU利用率持续低于30%说明有瓶颈、显存占用超过90%有OOM风险、P99延迟超过1秒用户体验差、队列长度超过100服务过载。但阈值不是拍脑袋定的要根据实际压测结果来。我一般会先用Locust或者wrk做压力测试逐步增加QPS观察各个指标的变化找到服务开始劣化的拐点然后把告警阈值设在拐点的80%处。比如压测发现QPS到500的时候延迟开始飙升那告警阈值就设在400。这样可以在服务真正过载之前就发出告警留出处理时间。还有一个经验是监控指标要分维度。比如延迟要分P50、P90、P99不能只看平均值因为平均值会被大量快请求拉低掩盖了慢请求的问题。显存要分已分配、已缓存、空闲不能只看总占用因为缓存显存是可以释放的已分配显存才是真正被模型占用的。这些细节文档里不会写但线上排查问题的时候非常有用。7. 我踩过的那些坑真实案例复盘7.1 一次OOM引发的连锁反应有一次我们的推理服务在高峰期突然OOM然后Kubernetes自动重启了Pod。但重启之后服务并没有恢复而是陷入了“启动-预热-OOM-重启”的循环。排查了半天发现原因是预热阶段用的样例数据比真实请求长很多导致预热时的显存占用比正常推理高出一大截。Pod重启后预热阶段就OOM了根本没机会处理真实请求。这个坑的教训是预热数据必须和真实请求的分布一致。我后来改成了从线上流量里采样一批真实请求做预热并且把预热时的最大batch size设成正常值的一半确保预热不会OOM。另外Kubernetes的重启策略也要调整不能无限重启我设了restartPolicy: OnFailure和backoffLimit: 3重启3次还失败就停止避免无限循环消耗资源。7.2 动态批处理导致的延迟毛刺还有一次我们上线了动态批处理吞吐量确实上去了但P99延迟出现了周期性的毛刺。查了很久发现原因是动态批处理的等待时间设成了10毫秒但有些请求的处理时间特别长比如输入特别长导致这一批的处理时间超过了10毫秒下一批请求就要等更久形成了正反馈延迟越来越高。解决办法是给动态批处理加一个“最大处理时间”限制。如果一批请求的处理时间超过了阈值比如50毫秒就把这一批拆成两批先返回已经处理完的部分剩下的放到下一批。这样虽然会增加一些批处理次数但能避免延迟毛刺。另外对于特别长的请求可以单独走一个“慢通道”不参与动态批处理直接单条推理避免拖累其他请求。7.3 版本升级导致的精度下降最隐蔽的一个坑是我们把PyTorch从1.12升级到2.0之后模型的推理精度下降了0.5个百分点。这个下降很小一开始没注意到后来做A/B测试才发现。排查后发现PyTorch 2.0的torch.compile默认开启了一些算子融合其中有一个融合操作在数值上做了近似导致精度损失。这个坑的教训是版本升级之后一定要做精度回归测试。我后来在CI流程里加了一个步骤每次升级框架版本都用固定的测试集跑一遍推理对比输出和之前版本的差异如果差异超过阈值就阻止升级。另外如果精度要求极高可以在torch.compile里关掉有损的融合操作比如设置modedefault并且禁用fuse_ops虽然会损失一些性能但能保证精度。8. 写给想从零开始做AI工程的人如果你看到这里说明你对AI工程是真的有兴趣。我想说的是从零开始做AI工程最难的不是写代码而是建立一套系统性的思维方式。你要时刻问自己这个操作的瓶颈在哪里是计算、内存、IO还是网络这个参数为什么设这个值有没有理论依据或者实测数据支撑这个方案在流量翻倍的时候还能不能扛住我刚开始做AI工程的时候也是到处找“最佳实践”但后来发现所谓的“最佳实践”都是有前提的。别人的batch size设128是因为他们的显存有80G别人的学习率设1e-4是因为他们的数据集有100万条。你直接抄过来大概率会翻车。所以我建议你养成一个习惯每做一个决策都记录下当时的上下文和理由过一段时间回头看你会发现很多决策其实可以做得更好。最后分享一个我常用的调试技巧当你遇到一个性能问题不要急着改代码先用py-spy或者torch.profiler做一次性能剖析找到真正的瓶颈。我见过太多人凭直觉优化结果优化了不重要的部分真正的瓶颈还在那里。数据不会骗人让数据告诉你该优化哪里。
