【AI专栏】图解深度学习-AI infra工程师必知必会 - 第 04 章：训练 vs 推理

博主介绍程序喵大人35 - 资深C/C/Rust/Android/iOS客户端开发10年大厂工作经验嵌入式/人工智能/自动驾驶/音视频/游戏开发入门级选手《C20高级编程》《C23高级编程》等多本书籍著译者更多原创精品文章首发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 gradienWW−lr×gradien。这也意味着在训练的时候只要你有一个参数矩阵你的显存里就必须再为它准备一块同等大小的空地专门用来放它的梯度显存账单这就翻倍了。五、优化器状态Adam 为什么“再吃一份显存”算完梯度直接更新不就完事了吗理论上是的那种最朴素的方法叫 SGDStochastic Gradient Descent随机梯度下降。SGD 就像一个没有任何记忆的急先锋它只看眼前的这一步。如果当前位置的梯度指向南它就往南迈一步下一时刻的梯度指向北它又立刻调头往北。它的更新公式非常直观W W − l r × g r a d i e n W W - lr \times gradienWW−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。路选对了账本才能算对后面要用到的工程手段才对得上号。这就是咱们下一章要展开讲的重头戏——怎么在具体的显卡上精打细算地过日子。码字不易欢迎大家点赞关注评论谢谢