【AI专栏】图解深度学习-AI infra工程师必知必会 - 第 04 章：训练 vs 推理
博主介绍：程序喵大人
- 35 - 资深C/C++/Rust/Android/iOS客户端开发
- 10年大厂工作经验
- 嵌入式/人工智能/自动驾驶/音视频/游戏开发入门级选手
- 《C++20高级编程》《C++23高级编程》等多本书籍著译者
- 更多原创精品文章，首发gzh，见文末
- 👇👇记得订阅专栏，以防走丢👇👇
😉C++基础系列专栏
😃C语言基础系列专栏
🤣C++大佬养成攻略专栏
🤓C++训练营
👉🏻个人网站
咱们前面花了不少时间逛车间、认算子。到了这一章，风格变了，咱们站到了这块工程白板前。作为 AI Infra 工程师，你接到一个新模型任务的第一反应，绝对不应该是去抠里面的网络结构，而是要先问一个极其核心的问题：“这个任务是做训练（Training），还是做推理（Inference）？”
你看白板上这左右两边的对比。左边的推理，干脆利落，数据像流水一样顺着箭头从头走到尾就结束了，底下的显存账本也很清爽，只有两行。右边的训练呢？它不光要正着走一遍，还得反着再绕一圈形成个闭环，底下的显存账本直接被塞得满满当当。
这两者的显存需求随随便便就能差出 3 到 8 倍。弄混了这两个概念，你要么会浪费成百上千万的机器资源，要么就是一跑程序直接爆显存（OOM）。接下来，咱们就把这个“正着走”和“反着走”的过程拆开看，把这笔显存账算得明明白白的。
一、Forward Pass：从输入算到输出
不管训练还是推理，起手式都是一样的：正向计算（Forward Pass）。
你看这条蓝色的流水线，数据框从左边进入第一层，算出个中间结果，再送进第二层……就这么一路单向流动，直到最右边输出预测结果。这中间每次经过那些 Linear、Norm 节点，数据的 shape 和数值都在按我们上一章讲的规矩发生变化。
这里藏着一个极度关键的区别：每一个节点算完之后，产生的那个中间结果（就是绿色的 Activations 缓存），要不要留着？如果是推理，算完这层进入下一层，前面的中间结果就没用了，显存直接覆盖掉，非常省空间。但如果是训练，抱歉，这些绿色的缓存你全得原封不动地保存在显存里，因为等会儿“反着走”的时候要用到它们。这就埋下了训练极其吃显存的第一个伏笔。
二、Loss：模型答得有多差
数据正向流到了终点，拿到预测结果了。在训练模式下，咱们就得掏出这个“偏差测量仪”了。
我们把模型自己算出来的概率（比如这里的[0.05, 0.72, 0.08, 0.15]）拿过来，和真实正确的答案（比如[0, 1, 0, 0]）放在一起对比。这个仪表盘测量出来的差距，就是 Loss（损失函数）。
Loss 算出来其实就是一个小小的浮点数（比如这里的0.327）。在显存里，它小到完全可以忽略不计。但它是整个训练过程的终极指挥棒——我们后面搞的所有大动作，唯一的目标就是让曲线图上的这个数字越来越小，让模型猜得越来越准。
也是从这个小小的数字开始，训练专属的“反向逆流”就要启动了。
三、Backward Pass：梯度从输出回到参数
你看，白板上多出了一条红色的反向路径。这就是反向传播（Backward Pass）。
它从最右边的 Loss 开始，顺着刚才蓝色箭头走过的路，倒着退回去。每经过一个层，它就会掏出刚才正向计算时特意存下来的那些绿色缓存（Activations），然后用微积分里的链式法则，算出一个叫“梯度”（Gradient）的东西。
如果说 Forward 是在算“模型觉得答案是什么”，那 Backward 就是在算“每个参数该怎么调，才能让刚才的回答更准一点”。这条红色的逆流计算量非常大，大约是正向计算的 2 倍。
四、梯度是什么：和参数同形状的另一组数字
刚才提到了“梯度”，咱们把镜头拉近，看看它到底是个什么长相。
其实一点都不神秘。你看白板左边，蓝色的参数矩阵 $$是 是是768 \times 307大小，里面装满了参与计算的权重。右边红色的就是刚才反向算出来的梯度矩阵 大小，里面装满了参与计算的权重。右边红色的就是刚才反向算出来的梯度矩阵大小，里面装满了参与计算的权重。右边红色的就是刚才反向算出来的梯度矩阵\partial L / \partial $$。
梯度的形状，和参数的形状一模一样！ 也是768 × 307 768 \times 307768×307。参数矩阵里的每一个数字，在梯度矩阵里都有一个对应的兄弟，这个兄弟记录了“你应该往上调一点还是往下调一点”。
到了最后一步，我们就用当前的参数，减去学习率乘以梯度，完成参数更新（W = W − l r × g r a d i e n W = W - lr \times gradienW=W−lr×gradien）。这也意味着，在训练的时候，只要你有一个参数矩阵，你的显存里就必须再为它准备一块同等大小的空地，专门用来放它的梯度，显存账单这就翻倍了。
五、优化器状态：Adam 为什么“再吃一份显存”
算完梯度直接更新不就完事了吗？理论上是的，那种最朴素的方法叫 SGD（Stochastic Gradient Descent，随机梯度下降）。
SGD 就像一个没有任何记忆的急先锋，它只看眼前的这一步。如果当前位置的梯度指向南，它就往南迈一步；下一时刻的梯度指向北，它又立刻调头往北。它的更新公式非常直观：W = W − l r × g r a d i e n W = W - lr \times gradienW=W−lr×gradien。这种策略特别省钱，显存账单上除了模型本身的权重外，只需要额外存储一份同等大小的梯度。它保持极简的无状态特征，省去了存储历史记忆的开销，账单极为清爽。
不过这种无状态的更新策略在大模型极度高维、复杂的空间里比较吃力。它遇到一些平坦的马鞍形山谷时，会在两边的山壁上反复剧烈震荡，在真正的谷底方向上前进得极其缓慢。这使得它的收敛速度偏慢，训练耗时很长。因此在实际工程中训练大模型，大家几乎清一色用的是 Adam 或者 AdamW 优化器。
Adam 聪明就聪明在，它会根据历史经验“自适应”地调整步子大小。为了记住这些历史经验，它给每个参数都准备了两个小本本：一个是“一阶矩”（m，橙色），一个是“二阶矩”（v，紫色）。
看出问题了吗？这两个小本本，也是和参数形状一模一样的矩阵，而且通常还得用最高精度（FP32）来存！
太无奈了，为了用 Adam 加速收敛，咱们又得多掏出两份巨大的显存空间，这就是为什么大家天天都在抱怨训练太费显存的根源所在！
六、训练显存账本：四笔费用加起来
现在咱们可以来汇总这张训练专属的显存明细单了。你看这四根横条：
- Weights（权重）：蓝条。模型本身的大小，常驻显存。
- Activations（激活值）：绿条。正向传出来的中间结果，为了留给反向传播用，不得不全囤着。随着 batch 大小水涨船高。
- Gradients（梯度）：红条。和权重一样大的一份开销。
- Optimizer States（优化器状态）：橙条。Adam 的m mm和v vv，通常是整个账单里最粗、最占地方的一笔。
四笔账加起来，如果训一个 7B 的模型，可能轻轻松松就破了 100GB，单张卡根本塞不下。少算任何一笔，你的训练程序都会在半夜无情地崩溃。
七、推理显存账本：只做 forward，账本立刻瘦身
看完了高昂的训练账本，再转头看看推理，你会觉得特别宽慰。
只要确定了是“进入推理模式”，那个巨大的红箭头就发挥作用了：不需要反向传播，所以梯度（Gradients）划掉！不需要更新参数，所以优化器状态（Optimizer States）划掉！甚至连中间的 Activations，用完一层就可以立刻扔掉，显存占用大跳水。
所以你看右边这个瘦身的账本，基本上就只剩下模型本身的权重（Weights），再加上极少量的动态空间（如果是跑 LLM，这里主要是留给 KV cache 用的，咱们后面讲）。这就解释了为什么一个需要几十张卡才能训练出来的 7B 模型，最后只要一张普通点的显卡就能舒舒服服地跑推理了。
八、总结：infra 工程师为什么先问“训练还是推理”
最后看这幅决策全景图。
以后接到模型任务，在分叉路口，先认准方向。如果是走左边去搞推理，你脑子里想的就是怎么管理瘦身后的显存（比如优化 KV cache），怎么把吞吐量拉高，怎么降低延迟让用户等得少一点。
如果是走右边去搞训练，面对那四座大山（特别是庞大的优化器状态），你想的就得是怎么做分布式，怎么把梯度在多张卡之间切分和同步，怎么抠下每一点显存来塞进更大的 batch。
路选对了，账本才能算对，后面要用到的工程手段才对得上号。这就是咱们下一章要展开讲的重头戏——怎么在具体的显卡上精打细算地过日子。
码字不易，欢迎大家点赞，关注，评论，谢谢！