1. 从零手搓AI工程：为什么我不建议你直接调包
很多人一听到“AI工程”这四个字，第一反应就是打开某个云平台，拖几个组件，调几个API，然后跑通一个Demo，就觉得自己已经掌握了。我刚开始也是这么想的，直到有一次线上推理服务在高峰期直接雪崩，日志里全是显存溢出和请求超时，我才意识到——那些封装好的高级接口，在关键时刻根本救不了你。你连模型为什么吃这么多显存、请求为什么排队、算子为什么慢，都说不清楚，怎么排查？
ai-engineering-from-scratch这个标题，核心不是教你从零训练一个GPT，而是教你从零搭建一套能跑、能调、能扛的AI工程链路。它解决的是“只会调包、不懂底层”的普遍痛点，适合那些已经会用PyTorch或TensorFlow写模型，但一遇到部署、优化、并发就抓瞎的开发者。说白了，就是从“能跑通”到“跑得稳、跑得快、跑得省”之间的那段路。
我见过太多团队，模型指标刷得很漂亮，一上线就翻车。问题往往不在算法本身，而在工程细节：数据加载成了瓶颈、推理没有做批处理、显存碎片化严重、服务没有做限流和降级。这些坑，你调包是永远调不出来的。所以这篇内容，我会按照一个真实项目的推进节奏，把从环境搭建到性能调优的完整链路拆开讲，每一步都告诉你为什么这么做，以及我踩过哪些坑。
2. 环境搭建：别让CUDA版本成为你的第一个噩梦
2.1 驱动、运行时与框架的三角关系
AI工程的第一步永远是环境。很多人觉得装个PyTorch有什么难的，pip install一行命令的事。但实际情况是，你的显卡驱动版本、CUDA运行时版本、PyTorch编译时链接的CUDA版本，这三者必须匹配。我见过最离谱的案例是，驱动是470，CUDA运行时是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里把这两个库的版本都写死，比如transformers==4.36.2和tokenizers==0.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_workers=0跑一个epoch，看GPU利用率，如果低于70%，再逐步增加num_workers，直到GPU利用率稳定在90%以上。
还有一个隐藏的坑是，pin_memory=True这个选项。它会把数据锁在页锁定内存里，加速CPU到GPU的传输。但如果你用的是多进程DataLoader，并且pin_memory=True，那么每个worker都会尝试锁内存，可能导致系统内存不足。我一般会在num_workers>0的时候开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_last=True，这些重复样本会影响梯度更新的准确性。
另外，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_inputs=torch.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_mb=128之后，显存碎片率从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里关掉有损的融合操作，比如设置mode="default"并且禁用fuse_ops，虽然会损失一些性能，但能保证精度。
8. 写给想从零开始做AI工程的人
如果你看到这里，说明你对AI工程是真的有兴趣。我想说的是，从零开始做AI工程，最难的不是写代码，而是建立一套系统性的思维方式。你要时刻问自己：这个操作的瓶颈在哪里？是计算、内存、IO还是网络？这个参数为什么设这个值？有没有理论依据或者实测数据支撑？这个方案在流量翻倍的时候还能不能扛住？
我刚开始做AI工程的时候，也是到处找“最佳实践”，但后来发现，所谓的“最佳实践”都是有前提的。别人的batch size设128，是因为他们的显存有80G；别人的学习率设1e-4，是因为他们的数据集有100万条。你直接抄过来，大概率会翻车。所以，我建议你养成一个习惯：每做一个决策，都记录下当时的上下文和理由，过一段时间回头看，你会发现很多决策其实可以做得更好。
最后分享一个我常用的调试技巧：当你遇到一个性能问题，不要急着改代码，先用py-spy或者torch.profiler做一次性能剖析，找到真正的瓶颈。我见过太多人凭直觉优化，结果优化了不重要的部分，真正的瓶颈还在那里。数据不会骗人，让数据告诉你该优化哪里。
