1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件、调几个API然后跑通一个Demo就觉得自己已经入门了。我刚开始接触这个方向的时候也是这么想的直到有一次线上推理服务在高峰期直接雪崩日志里全是显存溢出的报错我才意识到——只会调包的人根本不知道模型在底层到底经历了什么。ai-engineering-from-scratch这个项目标题本身就说明了一件事它要做的不是教你调库而是带你从最原始的矩阵乘法开始一步步把AI工程里那些被封装得严严实实的环节拆开来看。这篇文章适合三类人第一类是有一定编程基础但没接触过AI系统底层的人第二类是做过模型训练但没碰过推理部署的人第三类是被各种框架的“黑盒”坑过、想搞清楚内部机制的人。我写这篇东西的出发点很简单市面上讲AI工程的文章要么停留在“pip install然后调用”的层面要么直接跳到分布式训练的源码分析中间那一大段“从单机到集群、从训练到推理、从原型到生产”的工程细节几乎没人系统讲。而恰恰是这段细节决定了你的AI系统能不能扛住真实流量。接下来的内容我会按照一个AI工程从零搭建的真实顺序来展开先搞清楚数据管道怎么设计再聊模型训练的基础设施然后是推理服务的性能优化最后是监控和迭代。每一块我都会给出具体的代码示例、参数计算过程以及我在实际项目中踩过的坑。2. 数据管道AI工程里最容易被低估的环节2.1 为什么数据加载会成为训练瓶颈大部分人第一次写训练脚本的时候都是直接把整个数据集读进内存然后开始for循环。小数据集上没问题一旦数据量上到几十GB你就会发现GPU利用率忽高忽低有时候甚至掉到10%以下。这不是模型的问题是数据管道拖了后腿。我做过一个实测在单卡V100上训练一个ResNet-50数据集是ImageNet的一个子集大约50万张图片。如果直接用Python的PIL逐张读取每个epoch要花将近40分钟GPU利用率平均只有35%。后来改成先用lmdb把图片序列化成二进制格式再用多进程预取每个epoch降到12分钟GPU利用率稳定在85%以上。这里的核心逻辑是GPU的计算速度远快于磁盘IO和CPU预处理的速度。如果你不让数据加载和模型计算并行起来GPU大部分时间都在等数据。解决办法就是预取prefetch和多进程加载。2.2 从零实现一个带预取的数据加载器很多人直接用PyTorch的DataLoader这没问题但你要知道它内部做了什么。下面是一个简化版的实现思路帮你理解预取的原理import threading import queue class PrefetchLoader: def __init__(self, dataset, batch_size, num_workers4, prefetch_factor2): self.dataset dataset self.batch_size batch_size self.num_workers num_workers self.prefetch_factor prefetch_factor self.queue queue.Queue(maxsizeprefetch_factor) self._stop_event threading.Event() def _worker(self): indices list(range(len(self.dataset))) batch [] for idx in indices: batch.append(self.dataset[idx]) if len(batch) self.batch_size: if not self._stop_event.is_set(): self.queue.put(batch) batch [] def __iter__(self): threads [] for _ in range(self.num_workers): t threading.Thread(targetself._worker) t.start() threads.append(t) while True: try: yield self.queue.get(timeout1) except queue.Empty: break self._stop_event.set() for t in threads: t.join()这段代码的关键在于queue.Queue(maxsizeprefetch_factor)。队列有容量限制当队列满的时候worker线程会自动阻塞这样就不会无限占用内存。prefetch_factor一般设成2到4太小了起不到预取效果太大了会浪费内存。注意多进程加载在Windows上会有额外的序列化开销因为Windows的进程启动方式是spawn而不是fork。如果你在Windows上开发建议把num_workers设小一点或者直接用Linux环境。2.3 数据版本管理别让脏数据毁掉你的实验我见过太多团队在数据管理上翻车。同一个模型今天跑出来准确率85%明天跑出来82%查了半天发现是数据清洗脚本被人改了一行。数据版本管理不是可选项是必选项。最土但最有效的办法是每次数据预处理完算一个MD5校验和存到数据库里。训练脚本启动的时候先校验不匹配就直接报错退出。进阶一点可以用DVCData Version Control它能把数据文件和Git提交关联起来回滚代码的时候数据也能跟着回滚。方案适用场景优点缺点MD5校验小团队、数据集固定实现简单、零依赖无法追溯变更历史DVC中等团队、频繁迭代与Git集成、支持远程存储学习成本较高自建元数据服务大团队、多项目共享灵活可控开发和维护成本高我的建议是团队人数少于5个、数据集不经常变MD5就够了超过5个人或者数据集每周都在更新直接上DVC别犹豫。3. 训练基础设施从单卡到多卡的工程细节3.1 混合精度训练到底省了什么混合精度训练Mixed Precision Training是现在训练大模型的标配但很多人只知道它能省显存不知道它为什么能省、省了多少、有什么代价。先说原理。默认情况下模型参数和梯度都是FP3232位浮点数每个数占4个字节。混合精度把前向传播和反向传播中的大部分计算换成FP1616位浮点数每个数占2个字节。显存占用直接减半计算速度也能提升因为现代GPU对FP16的吞吐量通常是FP32的2到8倍。但FP16有个致命问题它的表示范围比FP32小很多。FP16能表示的最小正数是约6e-8最大是65504。训练过程中梯度值经常小于6e-8直接变成0模型就学不动了。解决办法是损失缩放Loss Scaling把损失值放大一个倍数反向传播出来的梯度也跟着放大等更新参数的时候再缩回去。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in dataloader: optimizer.zero_grad() with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler会自动调整缩放倍数。如果发现梯度溢出了变成inf或nan它会把倍数调小如果连续几次没溢出就把倍数调大。这个动态调整的过程不需要你手动干预。实测数据在A100上训练BERT-baseFP32的显存占用约12GBFP16降到6.5GB左右每个step的时间从0.42秒降到0.28秒。但要注意混合精度对某些操作如softmax、layer norm仍然需要FP32所以显存不会精确减半。3.2 数据并行和模型并行的选择逻辑单卡装不下模型的时候你就得考虑多卡了。多卡方案主要有两种数据并行Data Parallelism和模型并行Model Parallelism。数据并行的逻辑很简单每张卡上放一份完整的模型副本把batch切成N份每张卡算一份然后同步梯度。PyTorch的DistributedDataParallelDDP就是干这个的。它的通信量是模型参数量的大小每步都要做一次all-reduce。模型并行则是把模型本身切开不同的层放在不同的卡上。通信量取决于切分点如果切得不好卡之间等待的时间会很长。选择逻辑是这样的模型能装进单卡但训练太慢 → 数据并行模型装不进单卡但能装进两卡 → 模型并行或流水线并行模型巨大百亿参数以上→ 混合并行数据模型流水线我个人的经验是能用数据并行就用数据并行因为它的实现最简单、通信模式最成熟。只有当单卡真的装不下模型的时候才去折腾模型并行。模型并行的调试成本非常高一个切分点没选好性能可能还不如单卡。3.3 梯度累积小显存跑大batch的妥协方案有时候你只有一张8GB的卡但论文里说batch size要设256才能复现结果。怎么办梯度累积Gradient Accumulation就是答案。它的思路是把一个大batch拆成几个小batch分别做前向和反向但不清空梯度等累积够了再更新一次参数。这样等效于大batch训练但显存占用只有小batch的水平。accumulation_steps 8 optimizer.zero_grad() for i, (data, target) in enumerate(dataloader): output model(data) loss criterion(output, target) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意loss要除以accumulation_steps否则梯度会累积成原来的N倍相当于学习率被放大了N倍。但梯度累积有个副作用BatchNorm层的行为会变。因为BatchNorm是在每个小batch上算均值和方差的累积之后等效的batch统计量和小batch一样而不是大batch的。如果模型里有BatchNorm要么换成GroupNorm或LayerNorm要么用同步BatchNormSyncBN。4. 推理服务把模型变成能扛流量的API4.1 推理和训练的本质区别训练的时候你关心的是吞吐量每秒处理多少样本和收敛速度。推理的时候你关心的是延迟单个请求多久返回和并发能力。这两个目标经常是矛盾的。训练可以用大batch把GPU喂饱推理不行因为用户等不了。一个请求过来你不可能等凑够32个再一起算。所以推理的优化方向是在保证延迟可接受的前提下尽可能提高吞吐。这就引出了几个关键技术动态批处理Dynamic Batching、模型量化Quantization、算子融合Operator Fusion。4.2 动态批处理的实现与调参动态批处理的逻辑是维护一个请求队列每隔几毫秒把队列里的请求打包成一个batch送进GPU计算算完再拆开返回。这样既不会让用户等太久又能利用GPU的并行能力。import time import threading class DynamicBatcher: def __init__(self, model, max_batch_size32, max_wait_ms10): self.model model self.max_batch_size max_batch_size self.max_wait_ms max_wait_ms self.queue [] self.lock threading.Lock() def add_request(self, input_data): with self.lock: self.queue.append(input_data) if len(self.queue) self.max_batch_size: return self._process() time.sleep(self.max_wait_ms / 1000.0) with self.lock: if self.queue: return self._process() return None def _process(self): batch self.queue[:self.max_batch_size] self.queue self.queue[self.max_batch_size:] return self.model(batch)max_wait_ms这个参数很关键。设太小了batch凑不大GPU利用率低设太大了用户延迟高。我的经验值是在线服务设5到10毫秒离线批处理设50到100毫秒。具体还要看你的SLA要求。4.3 模型量化用精度换速度的账怎么算模型量化是把FP32的权重和激活值转换成INT8这样显存占用降到四分之一计算速度提升2到4倍。但精度会掉掉多少取决于模型和量化方法。最简单的量化是训练后量化Post-Training Quantization, PTQ拿一批校准数据跑一遍统计每层激活值的分布算出量化参数scale和zero_point然后直接转换。不需要重新训练几分钟就能搞定。import torch.quantization as tq model.eval() model.qconfig tq.get_default_qconfig(fbgemm) model_prepared tq.prepare(model) # 用校准数据跑几轮 for data in calibration_loader: model_prepared(data) model_quantized tq.convert(model_prepared)实测下来ResNet-50在ImageNet上FP32的准确率是76.1%PTQ量化后是75.8%掉了0.3个百分点但推理速度提升了2.8倍。这个 trade-off 在大多数场景下是划算的。但如果你的模型对精度极其敏感比如医疗影像诊断PTQ可能不够需要用量化感知训练Quantization-Aware Training, QAT在训练过程中模拟量化的误差让模型自己去适应。QAT的精度损失通常能控制在0.1个百分点以内但需要重新训练成本高很多。注意量化后的模型不是所有硬件都支持。INT8的矩阵乘法需要特定的指令集如x86的VNNI、ARM的dot product老旧的CPU上可能反而更慢。部署前一定要在目标硬件上实测。4.4 算子融合减少内存访问的隐形优化算子融合是把多个连续的小算子合并成一个大的算子减少中间结果的读写。比如Conv BatchNorm ReLU这三个操作如果不融合需要把Conv的输出写回显存再读出来做BatchNorm再写回去做ReLU。融合之后数据在寄存器或共享内存里就完成了所有计算只读写一次。TensorRT和TVM这类推理引擎会自动做算子融合但你需要知道它的原理才能在遇到性能问题时知道从哪里下手。我遇到过一个case模型在TensorRT上推理速度比PyTorch还慢查了半天发现是某个自定义算子不支持融合导致整个图被切成了两段中间插了一次显存拷贝。后来把那个算子用原生操作重写了一遍速度直接翻了3倍。5. 监控与迭代上线只是开始5.1 推理服务的核心监控指标模型上线之后你需要盯着几个关键指标P50/P95/P99延迟平均延迟没有意义要看长尾。P99延迟超过1秒用户就会明显感觉到卡顿。GPU利用率持续低于50%说明资源浪费持续高于90%说明快到瓶颈了。显存占用要留20%的余量否则突发流量一来就OOM。请求队列长度队列持续增长说明处理速度跟不上到达速度需要扩容或优化。这些指标用Prometheus Grafana就能搞定。我建议在服务启动的时候就埋好点别等出了问题再加那时候已经晚了。5.2 模型漂移的检测与应对模型不是上线就一劳永逸的。数据分布会变用户行为会变模型的效果会慢慢下降。这就是模型漂移Model Drift。检测漂移最简单的方法是监控预测结果的分布。比如一个分类模型如果之前预测结果中正样本占30%最近一周变成了60%那大概率是数据分布变了。更严谨的方法是用KL散度或PSIPopulation Stability Index来量化训练集和线上数据的分布差异。import numpy as np from scipy.stats import entropy def calculate_psi(expected, actual, buckets10): breakpoints np.linspace(0, 100, buckets 1) expected_percents np.histogram(expected, breakpoints)[0] / len(expected) actual_percents np.histogram(actual, breakpoints)[0] / len(actual) def sub_psi(e_perc, a_perc): if a_perc 0: a_perc 0.0001 if e_perc 0: e_perc 0.0001 return (e_perc - a_perc) * np.log(e_perc / a_perc) psi_value sum(sub_psi(expected_percents[i], actual_percents[i]) for i in range(len(expected_percents))) return psi_valuePSI小于0.1说明分布稳定0.1到0.25之间需要警惕超过0.25就说明分布显著变化了需要考虑重新训练。5.3 持续训练的数据闭环怎么建模型迭代的核心是数据闭环线上服务的请求和预测结果要存下来人工标注或用户反馈要关联上去然后定期用新数据重新训练。这个闭环听起来简单做起来坑很多。第一个坑是存储成本每天几百万请求原始输入加上中间结果存储费用很快就上去了。我的做法是只存采样后的数据比如每天随机采10%同时把所有异常请求延迟超高、预测置信度极低全量存下来。第二个坑是标注质量用户点击不等于满意用户不点击不等于不满意。隐式反馈的噪声很大需要设计合理的标注策略。我通常会用A/B测试来验证新模型的效果而不是直接看离线指标。第三个坑是训练和推理的一致性训练时的特征处理逻辑和推理时必须完全一致否则会出现训练时AUC 0.9、上线后AUC 0.6的惨剧。解决办法是把特征处理逻辑封装成独立的模块训练和推理共用同一份代码。6. 一些让我少走弯路的实操心得6.1 环境隔离不是可选项我见过太多“在我机器上能跑”的悲剧。AI工程涉及Python、CUDA、cuDNN、各种框架的版本依赖关系极其复杂。我的做法是每个项目一个Docker镜像基础镜像固定CUDA和Python版本所有依赖写进requirements.txt并锁定版本号。FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 RUN apt-get update apt-get install -y python3.10 python3-pip COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txtrequirements.txt里不要写torch1.12这种模糊版本要写torch1.13.1。差一个小版本API可能就变了。6.2 日志要打够但别打太多推理服务的日志我一般分三级ERROR级别记录所有异常和堆栈INFO级别记录每个请求的ID、输入摘要、输出摘要、延迟DEBUG级别记录完整的输入输出只在排查问题时临时开启。日志量太大会拖慢服务也会让存储成本失控。我的经验是INFO级别的日志量控制在请求量的1%以内DEBUG级别默认关闭。6.3 性能优化要先测量再动手很多人一遇到性能问题就开始瞎猜是不是模型太大了是不是batch太小了是不是GPU不够快正确的做法是先profile找到真正的瓶颈。PyTorch自带的torch.profiler就很好用with torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], scheduletorch.profiler.schedule(wait1, warmup1, active3), on_trace_readytorch.profiler.tensorboard_trace_handler(./log) ) as prof: for step, data in enumerate(dataloader): if step 5: break model(data) prof.step()跑完之后在TensorBoard里看哪个算子耗时最长、哪次显存拷贝最频繁一目了然。我做过一次优化profile之后发现瓶颈不在模型计算而在数据从CPU传到GPU的那一步后来用了pin_memory和non_blockingTrue速度直接提升了40%。6.4 别忽视CPU端的预处理GPU再快如果CPU预处理跟不上整体速度还是上不去。图片解码、归一化、数据增强这些操作如果全用Python做单核性能根本不够。我的做法是能用GPU做的预处理就放到GPU上做比如归一化和简单的几何变换必须用CPU做的就用多进程并行并且尽量用C实现的库如OpenCV、Pillow-SIMD而不是纯Python。还有一个容易被忽视的点数据从CPU内存拷贝到GPU显存的时间。如果预处理后的数据在CPU上是非连续的比如经过多次切片和拼接拷贝效率会很低。用np.ascontiguousarray()整理一下内存布局有时候能省不少时间。6.5 模型版本管理要像代码一样严格每次训练的模型都要有唯一的版本号关联到具体的代码提交、数据版本、超参数配置。我见过有人用model_final_v2_real_final.pth这种命名过了一周自己都不知道哪个是哪个。我的做法是用MLflow或者Weights Biases这类工具自动记录每次实验的参数、指标和产物。如果不想引入外部工具至少要在模型文件旁边存一个JSON写清楚训练配置和评估结果。{ model_version: 20240115-001, git_commit: a3f8c2d, data_version: imagenet-subset-v3, hyperparameters: { lr: 0.001, batch_size: 64, epochs: 50 }, metrics: { val_accuracy: 0.842, val_loss: 0.412 } }这样出了问题能快速回滚也能追溯是哪个改动导致了效果下降。6.6 压测要模拟真实流量模式上线前的压测不能只用均匀的请求速率。真实流量是有波峰的早上9点、中午12点、晚上8点请求量可能是平时的3到5倍。压测的时候要模拟这种突发流量看看服务能不能扛住。我一般用Locust或者wrk做压测设置一个阶梯式的请求速率从每秒10个请求开始每30秒翻一倍直到服务出现超时或错误。记录下每个速率下的P99延迟和错误率找到服务的拐点。注意压测环境要和生产环境尽量一致包括GPU型号、网络带宽、依赖服务的响应时间。在开发机上压测出来的数据没有参考价值。7. 写在最后从零构建的意义回到ai-engineering-from-scratch这个主题。我之所以强调“从零”不是因为调包不好而是因为只有你自己动手实现过一遍数据加载、梯度累积、动态批处理这些环节你才能在出问题的时候快速定位在需要优化的时候知道从哪里下手。调包能让你快速跑通Demo但从Demo到生产之间隔着数据管道、训练基础设施、推理优化、监控迭代这一整套工程体系。这套体系里的每一个环节都有它存在的理由和背后的权衡。理解这些理由和权衡才是一个AI工程师真正值钱的地方。我个人的体会是每次我深入一个“黑盒”的内部把它拆开来看清楚后面再遇到类似的问题解决速度就会快很多。因为你知道它为什么这样设计也就知道了它的边界在哪里什么时候会出问题出了问题该往哪个方向查。这个方向的内容还有很多可以展开的比如分布式训练中的通信优化、推理服务的自动扩缩容、多模型共享GPU的资源调度等等。每一个话题都值得单独写一篇。后面有机会我再逐个拆解。 从零手搓AI工程：数据管道、训练基础设施与推理服务优化实战
1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件、调几个API然后跑通一个Demo就觉得自己已经入门了。我刚开始接触这个方向的时候也是这么想的直到有一次线上推理服务在高峰期直接雪崩日志里全是显存溢出的报错我才意识到——只会调包的人根本不知道模型在底层到底经历了什么。ai-engineering-from-scratch这个项目标题本身就说明了一件事它要做的不是教你调库而是带你从最原始的矩阵乘法开始一步步把AI工程里那些被封装得严严实实的环节拆开来看。这篇文章适合三类人第一类是有一定编程基础但没接触过AI系统底层的人第二类是做过模型训练但没碰过推理部署的人第三类是被各种框架的“黑盒”坑过、想搞清楚内部机制的人。我写这篇东西的出发点很简单市面上讲AI工程的文章要么停留在“pip install然后调用”的层面要么直接跳到分布式训练的源码分析中间那一大段“从单机到集群、从训练到推理、从原型到生产”的工程细节几乎没人系统讲。而恰恰是这段细节决定了你的AI系统能不能扛住真实流量。接下来的内容我会按照一个AI工程从零搭建的真实顺序来展开先搞清楚数据管道怎么设计再聊模型训练的基础设施然后是推理服务的性能优化最后是监控和迭代。每一块我都会给出具体的代码示例、参数计算过程以及我在实际项目中踩过的坑。2. 数据管道AI工程里最容易被低估的环节2.1 为什么数据加载会成为训练瓶颈大部分人第一次写训练脚本的时候都是直接把整个数据集读进内存然后开始for循环。小数据集上没问题一旦数据量上到几十GB你就会发现GPU利用率忽高忽低有时候甚至掉到10%以下。这不是模型的问题是数据管道拖了后腿。我做过一个实测在单卡V100上训练一个ResNet-50数据集是ImageNet的一个子集大约50万张图片。如果直接用Python的PIL逐张读取每个epoch要花将近40分钟GPU利用率平均只有35%。后来改成先用lmdb把图片序列化成二进制格式再用多进程预取每个epoch降到12分钟GPU利用率稳定在85%以上。这里的核心逻辑是GPU的计算速度远快于磁盘IO和CPU预处理的速度。如果你不让数据加载和模型计算并行起来GPU大部分时间都在等数据。解决办法就是预取prefetch和多进程加载。2.2 从零实现一个带预取的数据加载器很多人直接用PyTorch的DataLoader这没问题但你要知道它内部做了什么。下面是一个简化版的实现思路帮你理解预取的原理import threading import queue class PrefetchLoader: def __init__(self, dataset, batch_size, num_workers4, prefetch_factor2): self.dataset dataset self.batch_size batch_size self.num_workers num_workers self.prefetch_factor prefetch_factor self.queue queue.Queue(maxsizeprefetch_factor) self._stop_event threading.Event() def _worker(self): indices list(range(len(self.dataset))) batch [] for idx in indices: batch.append(self.dataset[idx]) if len(batch) self.batch_size: if not self._stop_event.is_set(): self.queue.put(batch) batch [] def __iter__(self): threads [] for _ in range(self.num_workers): t threading.Thread(targetself._worker) t.start() threads.append(t) while True: try: yield self.queue.get(timeout1) except queue.Empty: break self._stop_event.set() for t in threads: t.join()这段代码的关键在于queue.Queue(maxsizeprefetch_factor)。队列有容量限制当队列满的时候worker线程会自动阻塞这样就不会无限占用内存。prefetch_factor一般设成2到4太小了起不到预取效果太大了会浪费内存。注意多进程加载在Windows上会有额外的序列化开销因为Windows的进程启动方式是spawn而不是fork。如果你在Windows上开发建议把num_workers设小一点或者直接用Linux环境。2.3 数据版本管理别让脏数据毁掉你的实验我见过太多团队在数据管理上翻车。同一个模型今天跑出来准确率85%明天跑出来82%查了半天发现是数据清洗脚本被人改了一行。数据版本管理不是可选项是必选项。最土但最有效的办法是每次数据预处理完算一个MD5校验和存到数据库里。训练脚本启动的时候先校验不匹配就直接报错退出。进阶一点可以用DVCData Version Control它能把数据文件和Git提交关联起来回滚代码的时候数据也能跟着回滚。方案适用场景优点缺点MD5校验小团队、数据集固定实现简单、零依赖无法追溯变更历史DVC中等团队、频繁迭代与Git集成、支持远程存储学习成本较高自建元数据服务大团队、多项目共享灵活可控开发和维护成本高我的建议是团队人数少于5个、数据集不经常变MD5就够了超过5个人或者数据集每周都在更新直接上DVC别犹豫。3. 训练基础设施从单卡到多卡的工程细节3.1 混合精度训练到底省了什么混合精度训练Mixed Precision Training是现在训练大模型的标配但很多人只知道它能省显存不知道它为什么能省、省了多少、有什么代价。先说原理。默认情况下模型参数和梯度都是FP3232位浮点数每个数占4个字节。混合精度把前向传播和反向传播中的大部分计算换成FP1616位浮点数每个数占2个字节。显存占用直接减半计算速度也能提升因为现代GPU对FP16的吞吐量通常是FP32的2到8倍。但FP16有个致命问题它的表示范围比FP32小很多。FP16能表示的最小正数是约6e-8最大是65504。训练过程中梯度值经常小于6e-8直接变成0模型就学不动了。解决办法是损失缩放Loss Scaling把损失值放大一个倍数反向传播出来的梯度也跟着放大等更新参数的时候再缩回去。from torch.cuda.amp import autocast, GradScaler scaler GradScaler() for data, target in dataloader: optimizer.zero_grad() with autocast(): output model(data) loss criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()GradScaler会自动调整缩放倍数。如果发现梯度溢出了变成inf或nan它会把倍数调小如果连续几次没溢出就把倍数调大。这个动态调整的过程不需要你手动干预。实测数据在A100上训练BERT-baseFP32的显存占用约12GBFP16降到6.5GB左右每个step的时间从0.42秒降到0.28秒。但要注意混合精度对某些操作如softmax、layer norm仍然需要FP32所以显存不会精确减半。3.2 数据并行和模型并行的选择逻辑单卡装不下模型的时候你就得考虑多卡了。多卡方案主要有两种数据并行Data Parallelism和模型并行Model Parallelism。数据并行的逻辑很简单每张卡上放一份完整的模型副本把batch切成N份每张卡算一份然后同步梯度。PyTorch的DistributedDataParallelDDP就是干这个的。它的通信量是模型参数量的大小每步都要做一次all-reduce。模型并行则是把模型本身切开不同的层放在不同的卡上。通信量取决于切分点如果切得不好卡之间等待的时间会很长。选择逻辑是这样的模型能装进单卡但训练太慢 → 数据并行模型装不进单卡但能装进两卡 → 模型并行或流水线并行模型巨大百亿参数以上→ 混合并行数据模型流水线我个人的经验是能用数据并行就用数据并行因为它的实现最简单、通信模式最成熟。只有当单卡真的装不下模型的时候才去折腾模型并行。模型并行的调试成本非常高一个切分点没选好性能可能还不如单卡。3.3 梯度累积小显存跑大batch的妥协方案有时候你只有一张8GB的卡但论文里说batch size要设256才能复现结果。怎么办梯度累积Gradient Accumulation就是答案。它的思路是把一个大batch拆成几个小batch分别做前向和反向但不清空梯度等累积够了再更新一次参数。这样等效于大batch训练但显存占用只有小batch的水平。accumulation_steps 8 optimizer.zero_grad() for i, (data, target) in enumerate(dataloader): output model(data) loss criterion(output, target) / accumulation_steps loss.backward() if (i 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()注意loss要除以accumulation_steps否则梯度会累积成原来的N倍相当于学习率被放大了N倍。但梯度累积有个副作用BatchNorm层的行为会变。因为BatchNorm是在每个小batch上算均值和方差的累积之后等效的batch统计量和小batch一样而不是大batch的。如果模型里有BatchNorm要么换成GroupNorm或LayerNorm要么用同步BatchNormSyncBN。4. 推理服务把模型变成能扛流量的API4.1 推理和训练的本质区别训练的时候你关心的是吞吐量每秒处理多少样本和收敛速度。推理的时候你关心的是延迟单个请求多久返回和并发能力。这两个目标经常是矛盾的。训练可以用大batch把GPU喂饱推理不行因为用户等不了。一个请求过来你不可能等凑够32个再一起算。所以推理的优化方向是在保证延迟可接受的前提下尽可能提高吞吐。这就引出了几个关键技术动态批处理Dynamic Batching、模型量化Quantization、算子融合Operator Fusion。4.2 动态批处理的实现与调参动态批处理的逻辑是维护一个请求队列每隔几毫秒把队列里的请求打包成一个batch送进GPU计算算完再拆开返回。这样既不会让用户等太久又能利用GPU的并行能力。import time import threading class DynamicBatcher: def __init__(self, model, max_batch_size32, max_wait_ms10): self.model model self.max_batch_size max_batch_size self.max_wait_ms max_wait_ms self.queue [] self.lock threading.Lock() def add_request(self, input_data): with self.lock: self.queue.append(input_data) if len(self.queue) self.max_batch_size: return self._process() time.sleep(self.max_wait_ms / 1000.0) with self.lock: if self.queue: return self._process() return None def _process(self): batch self.queue[:self.max_batch_size] self.queue self.queue[self.max_batch_size:] return self.model(batch)max_wait_ms这个参数很关键。设太小了batch凑不大GPU利用率低设太大了用户延迟高。我的经验值是在线服务设5到10毫秒离线批处理设50到100毫秒。具体还要看你的SLA要求。4.3 模型量化用精度换速度的账怎么算模型量化是把FP32的权重和激活值转换成INT8这样显存占用降到四分之一计算速度提升2到4倍。但精度会掉掉多少取决于模型和量化方法。最简单的量化是训练后量化Post-Training Quantization, PTQ拿一批校准数据跑一遍统计每层激活值的分布算出量化参数scale和zero_point然后直接转换。不需要重新训练几分钟就能搞定。import torch.quantization as tq model.eval() model.qconfig tq.get_default_qconfig(fbgemm) model_prepared tq.prepare(model) # 用校准数据跑几轮 for data in calibration_loader: model_prepared(data) model_quantized tq.convert(model_prepared)实测下来ResNet-50在ImageNet上FP32的准确率是76.1%PTQ量化后是75.8%掉了0.3个百分点但推理速度提升了2.8倍。这个 trade-off 在大多数场景下是划算的。但如果你的模型对精度极其敏感比如医疗影像诊断PTQ可能不够需要用量化感知训练Quantization-Aware Training, QAT在训练过程中模拟量化的误差让模型自己去适应。QAT的精度损失通常能控制在0.1个百分点以内但需要重新训练成本高很多。注意量化后的模型不是所有硬件都支持。INT8的矩阵乘法需要特定的指令集如x86的VNNI、ARM的dot product老旧的CPU上可能反而更慢。部署前一定要在目标硬件上实测。4.4 算子融合减少内存访问的隐形优化算子融合是把多个连续的小算子合并成一个大的算子减少中间结果的读写。比如Conv BatchNorm ReLU这三个操作如果不融合需要把Conv的输出写回显存再读出来做BatchNorm再写回去做ReLU。融合之后数据在寄存器或共享内存里就完成了所有计算只读写一次。TensorRT和TVM这类推理引擎会自动做算子融合但你需要知道它的原理才能在遇到性能问题时知道从哪里下手。我遇到过一个case模型在TensorRT上推理速度比PyTorch还慢查了半天发现是某个自定义算子不支持融合导致整个图被切成了两段中间插了一次显存拷贝。后来把那个算子用原生操作重写了一遍速度直接翻了3倍。5. 监控与迭代上线只是开始5.1 推理服务的核心监控指标模型上线之后你需要盯着几个关键指标P50/P95/P99延迟平均延迟没有意义要看长尾。P99延迟超过1秒用户就会明显感觉到卡顿。GPU利用率持续低于50%说明资源浪费持续高于90%说明快到瓶颈了。显存占用要留20%的余量否则突发流量一来就OOM。请求队列长度队列持续增长说明处理速度跟不上到达速度需要扩容或优化。这些指标用Prometheus Grafana就能搞定。我建议在服务启动的时候就埋好点别等出了问题再加那时候已经晚了。5.2 模型漂移的检测与应对模型不是上线就一劳永逸的。数据分布会变用户行为会变模型的效果会慢慢下降。这就是模型漂移Model Drift。检测漂移最简单的方法是监控预测结果的分布。比如一个分类模型如果之前预测结果中正样本占30%最近一周变成了60%那大概率是数据分布变了。更严谨的方法是用KL散度或PSIPopulation Stability Index来量化训练集和线上数据的分布差异。import numpy as np from scipy.stats import entropy def calculate_psi(expected, actual, buckets10): breakpoints np.linspace(0, 100, buckets 1) expected_percents np.histogram(expected, breakpoints)[0] / len(expected) actual_percents np.histogram(actual, breakpoints)[0] / len(actual) def sub_psi(e_perc, a_perc): if a_perc 0: a_perc 0.0001 if e_perc 0: e_perc 0.0001 return (e_perc - a_perc) * np.log(e_perc / a_perc) psi_value sum(sub_psi(expected_percents[i], actual_percents[i]) for i in range(len(expected_percents))) return psi_valuePSI小于0.1说明分布稳定0.1到0.25之间需要警惕超过0.25就说明分布显著变化了需要考虑重新训练。5.3 持续训练的数据闭环怎么建模型迭代的核心是数据闭环线上服务的请求和预测结果要存下来人工标注或用户反馈要关联上去然后定期用新数据重新训练。这个闭环听起来简单做起来坑很多。第一个坑是存储成本每天几百万请求原始输入加上中间结果存储费用很快就上去了。我的做法是只存采样后的数据比如每天随机采10%同时把所有异常请求延迟超高、预测置信度极低全量存下来。第二个坑是标注质量用户点击不等于满意用户不点击不等于不满意。隐式反馈的噪声很大需要设计合理的标注策略。我通常会用A/B测试来验证新模型的效果而不是直接看离线指标。第三个坑是训练和推理的一致性训练时的特征处理逻辑和推理时必须完全一致否则会出现训练时AUC 0.9、上线后AUC 0.6的惨剧。解决办法是把特征处理逻辑封装成独立的模块训练和推理共用同一份代码。6. 一些让我少走弯路的实操心得6.1 环境隔离不是可选项我见过太多“在我机器上能跑”的悲剧。AI工程涉及Python、CUDA、cuDNN、各种框架的版本依赖关系极其复杂。我的做法是每个项目一个Docker镜像基础镜像固定CUDA和Python版本所有依赖写进requirements.txt并锁定版本号。FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 RUN apt-get update apt-get install -y python3.10 python3-pip COPY requirements.txt . RUN pip install --no-cache-dir -r requirements.txtrequirements.txt里不要写torch1.12这种模糊版本要写torch1.13.1。差一个小版本API可能就变了。6.2 日志要打够但别打太多推理服务的日志我一般分三级ERROR级别记录所有异常和堆栈INFO级别记录每个请求的ID、输入摘要、输出摘要、延迟DEBUG级别记录完整的输入输出只在排查问题时临时开启。日志量太大会拖慢服务也会让存储成本失控。我的经验是INFO级别的日志量控制在请求量的1%以内DEBUG级别默认关闭。6.3 性能优化要先测量再动手很多人一遇到性能问题就开始瞎猜是不是模型太大了是不是batch太小了是不是GPU不够快正确的做法是先profile找到真正的瓶颈。PyTorch自带的torch.profiler就很好用with torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA], scheduletorch.profiler.schedule(wait1, warmup1, active3), on_trace_readytorch.profiler.tensorboard_trace_handler(./log) ) as prof: for step, data in enumerate(dataloader): if step 5: break model(data) prof.step()跑完之后在TensorBoard里看哪个算子耗时最长、哪次显存拷贝最频繁一目了然。我做过一次优化profile之后发现瓶颈不在模型计算而在数据从CPU传到GPU的那一步后来用了pin_memory和non_blockingTrue速度直接提升了40%。6.4 别忽视CPU端的预处理GPU再快如果CPU预处理跟不上整体速度还是上不去。图片解码、归一化、数据增强这些操作如果全用Python做单核性能根本不够。我的做法是能用GPU做的预处理就放到GPU上做比如归一化和简单的几何变换必须用CPU做的就用多进程并行并且尽量用C实现的库如OpenCV、Pillow-SIMD而不是纯Python。还有一个容易被忽视的点数据从CPU内存拷贝到GPU显存的时间。如果预处理后的数据在CPU上是非连续的比如经过多次切片和拼接拷贝效率会很低。用np.ascontiguousarray()整理一下内存布局有时候能省不少时间。6.5 模型版本管理要像代码一样严格每次训练的模型都要有唯一的版本号关联到具体的代码提交、数据版本、超参数配置。我见过有人用model_final_v2_real_final.pth这种命名过了一周自己都不知道哪个是哪个。我的做法是用MLflow或者Weights Biases这类工具自动记录每次实验的参数、指标和产物。如果不想引入外部工具至少要在模型文件旁边存一个JSON写清楚训练配置和评估结果。{ model_version: 20240115-001, git_commit: a3f8c2d, data_version: imagenet-subset-v3, hyperparameters: { lr: 0.001, batch_size: 64, epochs: 50 }, metrics: { val_accuracy: 0.842, val_loss: 0.412 } }这样出了问题能快速回滚也能追溯是哪个改动导致了效果下降。6.6 压测要模拟真实流量模式上线前的压测不能只用均匀的请求速率。真实流量是有波峰的早上9点、中午12点、晚上8点请求量可能是平时的3到5倍。压测的时候要模拟这种突发流量看看服务能不能扛住。我一般用Locust或者wrk做压测设置一个阶梯式的请求速率从每秒10个请求开始每30秒翻一倍直到服务出现超时或错误。记录下每个速率下的P99延迟和错误率找到服务的拐点。注意压测环境要和生产环境尽量一致包括GPU型号、网络带宽、依赖服务的响应时间。在开发机上压测出来的数据没有参考价值。7. 写在最后从零构建的意义回到ai-engineering-from-scratch这个主题。我之所以强调“从零”不是因为调包不好而是因为只有你自己动手实现过一遍数据加载、梯度累积、动态批处理这些环节你才能在出问题的时候快速定位在需要优化的时候知道从哪里下手。调包能让你快速跑通Demo但从Demo到生产之间隔着数据管道、训练基础设施、推理优化、监控迭代这一整套工程体系。这套体系里的每一个环节都有它存在的理由和背后的权衡。理解这些理由和权衡才是一个AI工程师真正值钱的地方。我个人的体会是每次我深入一个“黑盒”的内部把它拆开来看清楚后面再遇到类似的问题解决速度就会快很多。因为你知道它为什么这样设计也就知道了它的边界在哪里什么时候会出问题出了问题该往哪个方向查。这个方向的内容还有很多可以展开的比如分布式训练中的通信优化、推理服务的自动扩缩容、多模型共享GPU的资源调度等等。每一个话题都值得单独写一篇。后面有机会我再逐个拆解。