训练做到一半loss 曲线突然往上蹿接着直接变 NaN。查了一整天最后用梯度检查把问题定位在一个自定义算子的反向传播上——不是梯度爆炸不是学习率太大而是 (dX G W^T) 这个公式里的矩阵转置方向抄反了。在 LLM 训练这种动辄几千亿 FLOPs 的场景里一个 (dX) 的错误会让整个反向传播链从最后一层开始逐层污染前面所有层的梯度全部报废。这篇文章要讲的就是反向传播里的 (dX)它到底是什么、为什么是 (G W^T) 而不是 (W^T G)、它在计算图里怎么流转、以及从 FLOPs 角度看它到底烧了多少算力。这是“FLOPs 计算详解——LLM 训练基础知识”系列的一部分我们默认你已经知道全连接层前向 (Y XW b) 是什么但不要求你会手推反向传播。我会从最基础的矩阵约定开始一路推到底再给你一份能直接跑的 NumPy 验证代码。1. 先搞清楚 (dX) 到底在问什么1.1 一个典型的 LLM 训练排查现场我们当时训的是一个参数量 1.3B 左右的 Decoder-only 自回归模型。前向正常、loss 正常下降跑到第 3 万多步时开始出现偶发 NaN。一开始怀疑是学习率 warmup 策略问题调低了峰值 LRNaN 照旧又怀疑是数据 pipeline 里混入了异常样本清洗之后问题依旧。最后我在损失函数和最后一个全连接层之间插了一个梯度 hook把每一层的梯度形状和数值打印出来才发现问题根本不在 loss而在最后一个线性层反向传播计算 (dX) 时把公式写错了。这个教训让我意识到一件事现在用自动微分框架写模型我们几乎永远不需要手写 (dX)。但正因为如此一旦遇到需要自定义算子、改 kernel、重计算激活值或者手动排查梯度传导问题时绝大多数人会在 (dX) 的推导上卡住。它比 (dW) 和 (db) 更隐蔽——(dW) 错了你会马上看到形状对不上(dX) 错了形状却可以完全正确梯度检查之前根本发现不了。1.2 用一条“退货流水线”理解反向传播想象一个制造工厂原材料 (X) 进入第一道工序加工后变成 (Z_1)经过激活函数变成 (A_1)再经过第二道工序变成 (Y)最后质检部门根据 (Y) 和标准答案对比算出一个损失 (L)。前向是不折不扣的流水线。反向传播则是“退货追责”质检部门算出每一件成品 (Y) 的“责任度”也就是 (G \partial L / \partial Y)。现在要追责第一道工序的原材料 (X)就必须沿着流水线从后往前走。走到任意一道工序时你需要知道两件事这批货从后面传过来的责任度是多少以及当前工序对输入有多敏感。两者相乘就是传给前一道工序的责任度。在数学上这个“传给前一道工序的责任度”就是 (dX \partial L / \partial X)。它不叫“梯度消失时的梯度”也不叫“损失对激活值的梯度”它有一个更准确的名字上游梯度upstream gradient。它不关心权重怎么更新只负责把信息往上一层传。2. 符号约定为什么全篇坚持用行向量表示2.1 行向量约定 vs 列向量约定矩阵求导最烦人的不是链式法则本身而是符号约定不统一。PyTorch 等主流框架的惯例是输入 (X) 形状为 ((m, n))权重 (W) 形状为 ((n, k))线性层输出 (Y XW b) 形状为 ((m, k))。这里的 (m) 是 batch size(n) 是输入维度(k) 是输出维度。按照这种约定每个样本是行向量一次性处理一个 batch 就是 (m) 行。有些教材和个人笔记喜欢用列向量一个样本是 ((n, 1)) 的列向量权重变成 (W^T)公式变成 (Y W^T X b)。两种约定数学上完全等价但公式会多出很多转置推导时左右乘顺序特别容易错。我的建议是写代码、写笔记、做推导统一采用行向量约定。理由有二第一和 PyTorch / TensorFlow 的 API 行为直接对应nn.Linear内部就是 (y xW^T b)如果你的 (x) 是 ((m, n))那么weight是 ((k, n))本质上仍然是一个 (W) 满足 (Y XW^T)。我们推导时直接记成 (Y XW) 更干净。第二行向量约定下梯度 (dX) 的形状永远和 (X) 一致这是一个极好的 sanity check。如果你算出一个 (dX) 形状和 (X) 不同别犹豫肯定前面某步错了。顺带一提不要相信“矩阵导数就是雅可比矩阵”这种说法。(dX) 不是雅可比矩阵雅可比矩阵是 ((\partial Y / \partial X))形状 ((m, n, m, k)) 这种四维张量。损失 (L) 是一个标量(dX) 是标量对矩阵每个元素的偏导形状和 (X) 一样。这是新手最容易混淆的地方。2.2 层层拆解L、Y、G 三者的关系设最后一层线性层的输出为 (Y XW b)其中(X)当前层输入形状 ((m, n))(W)权重形状 ((n, k))(b)偏置形状 ((1, k))(Y)logits形状 ((m, k))损失函数 (L) 是一个标量通常由 (Y) 经过 softmax 交叉熵等得到。在反向传播的第一站我们已经拿到 (G \partial L / \partial Y)形状与 (Y) 相同即 ((m, k))。这个 (G) 就是“从损失函数传过来的已经算好的上游梯度”。对于最后那个线性层来说(G) 是已知量由损失函数直接给出。我们的目标是算出三个梯度(dW \partial L / \partial W)形状 ((n, k))用来更新当前层权重(db \partial L / \partial b)形状 ((1, k))用来更新偏置(dX \partial L / \partial X)形状 ((m, n))不是用来更新任何参数而是传给前一层注意第三点的措辞(dX) 的名字里带个“梯度”但它不是参数梯度。在 PyTorch 里如果X是requires_gradTrue的叶子张量X.grad会被写入如果X是中间激活值它只会作为 grad_output 传给上一层的反向函数。它存在的唯一目的就是让前一层的 (dW) 有东西可算。2.3 三个梯度各自的用途一张表说清楚梯度数学表达式形状用途(dW)(X^T G)((n, k))更新当前层权重(db)(\sum_i G_i)对 batch 求和((1, k))更新当前层偏置(dX)(G W^T)((m, n))传给前一层作为上游梯度供前一层计算 (dW)这张表值得刻在脑子里。尤其是 (dX) 那行它叫“输入梯度”但它服务的是“前一层的权重更新”。标题里那句“传给前一层来计算 dW 的上游梯度”说的就是这么回事。3. 最后一层 (dX G W^T) 的完整推导3.1 从标量链式法则扩展到矩阵先看最朴素的标量场景(y xw)损失 (L \ell(y))那么[ \frac{\partial L}{\partial x} \frac{\partial L}{\partial y} \cdot \frac{\partial y}{\partial x} g \cdot w ]你看标量情况下就是“上游梯度乘以权重”。现在把它推广到矩阵多了一个维度求和的问题[ Y_{i,p} \sum_{q1}^{n} X_{i,q} W_{q,p} b_p ]一次对单个元素求导[ \frac{\partial Y_{i,p}}{\partial X_{i,j}} W_{j,p} ]这是关键一步。因为 (Y_{i,p}) 只依赖 (X) 的第 (i) 行所有元素不依赖其他行的元素所以求导结果里没有对 (i) 的求和。现在用链式法则把损失对 (X_{i,j}) 的偏导写出来[ dX_{i,j} \frac{\partial L}{\partial X_{i,j}} \sum_{p1}^{k} \frac{\partial L}{\partial Y_{i,p}} \cdot \frac{\partial Y_{i,p}}{\partial X_{i,j}} \sum_{p1}^{k} G_{i,p} W_{j,p} ]最后这个求和式恰好就是两个矩阵相乘的公式如果令 (dX G W^T)那么[ (G W^T){i,j} \sum{p1}^{k} G_{i,p} (W^T){p,j} \sum{p1}^{k} G_{i,p} W_{j,p} ]完全一致。所以 (dX G W^T) 不是拍脑袋的结论而是一个逐元素求和后整理出来的矩阵形式。3.2 维度验证三行代码的 sanity check推导完必须验形状。已知 (G) 是 ((m, k))(W^T) 是 ((k, n))矩阵相乘得到 ((m, n))正好等于输入 (X) 的形状。这一步在任何全连接层上做验证import numpy as np m, n, k 8, 16, 32 rng np.random.default_rng(42) X rng.standard_normal((m, n)) W rng.standard_normal((n, k)) b rng.standard_normal((1, k)) G rng.standard_normal((m, k)) dX G W.T # 形状 (m, n) (8, 16) dW X.T G # 形状 (n, k) (16, 32) db G.sum(axis0, keepdimsTrue) # 形状 (1, k) (1, 32) assert dX.shape X.shape, fdX shape mismatch: {dX.shape} vs {X.shape} assert dW.shape W.shape, fdW shape mismatch: {dW.shape} vs {W.shape} assert db.shape b.shape, fdb shape mismatch: {db.shape} vs {b.shape}很多教程只给公式不验维度导致读者抄代码时把 (W^T) 抄成 (W)最后从不报错但梯度全错。形状断言是反向传播里成本最低的保险丝强烈建议在每一层的自定义反向函数里都加上。3.3 为什么是“右乘 (W^T)”而不是“左乘”有人可能问既然 (Y XW)那导数不应该是 (\partial Y / \partial X W^T) 吗链式法则乘起来怎么变成 (G W^T) 而不是 (W^T G)关键在矩阵乘法的行/列对应关系。从上一节的逐元素推导可以看得很清楚(Y) 的第 (p) 列只由 (W) 的第 (p) 列决定而 (X) 的第 (j) 个元素影响的是一整行 (Y)。展开后上游梯度 (G) 的第 (i) 行向量 (G_{i,:}) 要对 (W) 的每一列做内积恰好是右乘 (W^T)。反过来的 (W^T G) 形状是 ((n, m))根本不对。这里有一个记忆技巧前向是“输入从左边进”反向就是“梯度从右边出”。前向 (X) 左乘 (W)反向 (G) 右乘 (W^T)对称关系一目了然。同理(dW X^T G) 是另一个对称方向(X) 转置后左乘 (G)正好匹配 (W) 的形状 ((n, k))。如果换成列向量约定 (Y W_h^T X b)最终推导结果是 (dX W_h G)。公式形态完全不同但 FLOPs 完全一样。我在实操里见过有人把两种约定的推导混在一起最后写出来一个形状对不上但“看起来很像”的式子。没有哪一种约定更好但混用一定出 bug。3.4 一个重要细节为什么 (dX) 的外面看得到 (W^T)把 (dX G W^T) 展开写得更直观一些[ dX \underbrace{G}{(m,k)} \cdot \underbrace{W^T}{(k,n)} \sum_{p1}^{k} G_{:,p} \otimes W_{p,:} ]也就是说(dX) 是 (k) 个秩一矩阵的叠加。每一层输出的第 (p) 个神经元的上游梯度 (G_{:,p})会把权重矩阵第 (p) 行的信息“印”到输入梯度上。这意味着(dX) 里包含了当前层权重的全部信息前一层看到 (dX)不仅知道“我该怎么调整”还能反推“下一层大概是怎么处理我的输出的”。这就是为什么 (dX) 必须被完整地计算出来——它不只是传给前一层它本身就是前一层学习信号的载体。4. (dX) 在反向传播中的真实角色不只是“输入梯度”4.1 一条链上的信息接力假设模型只有两层(X \to W_1 \to Z_1 \to \sigma \to A_1 \to W_2 \to Y \to L)。反向传播的顺序是计算 (G \partial L / \partial Y)第二层线性层反推 (dW_2 A_1^T G)、(db_2)以及 (dA_1 G W_2^T)激活函数层反推 (dZ_1 dA_1 \circ \sigma(Z_1))逐元素相乘第一层线性层拿到 (dZ_1) 作为自己的上游梯度计算 (dW_1 X^T dZ_1)、(db_1)以及 (dX dZ_1 W_1^T)第 4 步里第一层算出的 (dX) 如果还要传得更深它就继续作为第 0 层Embedding 或上一层 Transformer Block的上游梯度。所以你看到 (dX) 这个名字在每一层反向代码里反复出现线性层的 (dX) 恰好是它上一层的上游梯度 G。同一块数据换个层看名字就变了。这种“换个视角身份就变了”的设计在代码里最容易造成混乱。我见过有人写自定义反向函数时把一个变量同时当作当前层的 (dX) 和上一层的 (G)结果不小心在前向里多乘了一个什么东西后梯度偏差被“传递性”放大很难查。建议在代码中显式命名grad_out不要在函数内部调用dX。4.2 残差连接处(dX) 必须走“加法分流”LLM 里的 Transformer Block 处处是残差连接残差连接在反向传播中会变成一个梯度求和节点[ Y X F(X) ]那么[ dX dY dF ]也就是说残差块的输入梯度等于两条路径的上游梯度之和。很多初学的人在写自定义残差算子时只算了 (F) 那条分支的梯度忘了把跳跃连接的 (dY) 原样加回去。单独看一条路径没错加总后损失函数关于输入的导数就是错的。在 1.3B 的模型里这种错误不会让形状报错只会让 loss 低效下降甚至在某几层出现梯度“神秘消失”排查起来非常恶心。需要注意上面这张图的思维模型是前向时主河道的水绕过了巨石形成两个分支反向时两个分支的回流必须重新合在一起。用数学写就是# 残差块反向 dX_main dY # 跳跃连接路径原样传递 dX_branch backward_of_F(dY) # 主干路径 dX dX_main dX_branch # 分叉合并有了这个理解后续看 FlashAttention 或各种融合 kernel 的反向实现会轻松很多——它们本质上都在算矩阵乘法的 (dX)、(dW)只是加上了 mask 和分块。4.3 (dX) 与激活函数梯度的配合激活函数层没有权重反向时只需要把输入梯度逐元素乘以激活函数的导数[ dZ dA \circ f(Z) ]这里 (dA) 正是上一层线性层传下来的 (dX)。以 ReLU 为例(f(Z) (Z 0))所以 (dZ dA \circ (Z 0))。ReLU 的梯度极其稀疏这也意味着越深、越大规模的模型梯度在传播过程中天然会丢失超过一半的信息。但注意即使激活函数的梯度为 0它也不能“跳过”上游梯度直接置 0必须是先乘上游梯度再逐元素过滤。顺序写反数值上不等于原结果虽然等于零的时候无差别但这会让你的实现很难推广到 Leaky ReLU、GELU 等非饱和激活。在 GELU 这类平滑激活上导数是连续函数乘法和过滤顺序的区别就更体现出来了。PyTorch 的F.gelu反向实现里有完整的 (dZ dA \cdot \Phi(Z) dA \cdot Z \cdot \phi(Z)) 这一套式子它们的上游输入就是更前一层的 (dX)。5. 用 NumPy 手工实现一次完整反向传播顺便算 FLOPs5.1 最小三层网络的前向与反向光看公式不够我们干脆从零手写一个两层 MLP 的梯度。这里不用任何自动微分框架保证每一步都透明。import numpy as np from numpy.random import default_rng rng default_rng(42) m, n, h, k 8, 16, 32, 4 X rng.standard_normal((m, n)) W1 rng.standard_normal((n, h)) b1 rng.standard_normal((1, h)) W2 rng.standard_normal((h, k)) b2 rng.standard_normal((1, k)) Y_true rng.standard_normal((m, k)) # ---------- 前向 ---------- def forward(X): Z1 X W1 b1 A1 np.maximum(0, Z1) # ReLU Y A1 W2 b2 return Z1, A1, Y Z1, A1, Y forward(X) # 使用简单 MSE 损失作为 L def loss(Y): return np.mean((Y - Y_true) ** 2) # ---------- 反向 ---------- G 2.0 * (Y - Y_true) / (m * k) # dL/dY形状 (m, k) # 第二层线性层 dW2 A1.T G # (h, k) db2 G.sum(axis0, keepdimsTrue) # (1, k) dA1 G W2.T # (m, h) —— 这就是第二层的 dX作为下一层上游梯度 # 激活层 dZ1 dA1 * (Z1 0) # (m, h) # 第一层线性层 dW1 X.T dZ1 # (n, h) db1 dZ1.sum(axis0, keepdimsTrue) # (1, h) dX dZ1 W1.T # (m, n) —— 这就是我们要的 dX如果还有更深的层它会继续往上送注意这里dA1和dX的关系dA1是第二层输出的“下一层上游梯度”dZ1是第一层激活后传出的梯度而真正的最外层dX是dZ1 W1.T。整个链条里上一层的dX就是下一层的G但变量名不同容易看晕。建议拿纸笔把每个变量的形状标一遍。5.2 数值梯度验证最靠谱的 debug 工具算法都实现了但怎么证明它是对的黄金标准是数值梯度把某个参数扰动一个极小量看损失函数的变化率和解析梯度比。只要两者误差在 (1 \times 10^{-6}) 量级以下基本可以确认反向传播公式正确。def numerical_grad(f, param, eps1e-6): g np.zeros_like(param) it np.nditer(param, flags[multi_index]) while not it.finished: idx it.multi_index old param[idx] param[idx] old eps fp f() param[idx] old - eps fm f() param[idx] old g[idx] (fp - fm) / (2 * eps) it.iternext() return g f lambda: loss(forward(X)[2]) # 对 W2 做验证 g_analytic dW2 g_numeric numerical_grad(f, W2) print(W2 解析梯度 max abs error:, np.abs(g_analytic - g_numeric).max()) print(W1 解析梯度 max abs error:, np.abs(dW1 - numerical_grad(f, W1)).max()) print(dX 解析梯度 max abs error:, np.abs(dX - numerical_grad(f, X)).max())实测下来用 (eps1e-6) 时上面三个误差通常在 (1 \times 10^{-10}) 量级。如果你的输出数值超过 (1 \times 10^{-5})不用怀疑反向里十有八九有 bug。这个方法排查了我在自定义算子中遇到的无数个问题比看代码找错高效得多。5.3 从 FLOPs 视角看反向传播它比前向贵一倍标题里挂着“FLOPs 计算详解”现在把前向和反向的开销摆到台面上。约定一次浮点操作FLOP为一次乘法或一次加法矩阵乘法 (C_{(m, k)} A_{(m, n)} B_{(n, k)}) 需要[ 2mnk ](m \times k) 个输出元素每个元素做 (n) 次乘法和 (n) 次加法对于一层全连接计算公式FLOPs前向(Y XW)(2mnk)反向计算 (dX)(G W^T)(2mnk)反向计算 (dW)(X^T G)(2mnk)反向计算 (db)(G) 按行求和(O(mk))可忽略合计反向(4mnk)所以一层线性层在完整训练前向 反向中的 FLOPs 是 (6mnk)其中反向占 (2/3)。这还不是全部——激活函数、Softmax、LayerNorm 的反向也要算 FLOPs只是通常远小于矩阵乘法。在 Transformer 类 LLM 里训练一次整体迭代的 FLOPs 中反向传播大约占三分之二。这也是混合精度训练、梯度检查点activation checkpointing这类技术要重点优化反向路径的原因。有一个容易漏算的地方如果你为了省显存启用了 activation checkpointing前向时不保存激活值反向时重新计算前向那么每个 checkpoint 区间会多付出一次前向的 (2mnk) FLOPs。这也是为什么很多人说“checkpointing 是用 FLOPs 换显存”的直观来源。在你做 FLOPs 预算和训练吞吐估算时别只算 forward pass否则会高估 MFU模型浮点利用率一倍左右。6. 常见问题与排查技巧实录6.1 问题速查表现象可能原因处理方式反向时报 shape mismatch(dX) 与 (X) 形状不一致左右乘反了用assert卡形状对照 3.2 节维度验证数值梯度误差大链式法则漏乘上游梯度用numerical_grad逐层定位loss 下降缓慢(dW) 和 (dX) 都算错但形状碰巧对数值梯度检查不要只看 loss 曲线残差网络出现梯度 NaN残差分支梯度求和漏项拆两条路径分别打印梯度显存爆掉前向缓存了全量激活值考虑 activation checkpointing接受额外 FLOPs混合精度下梯度溢出(dX) 数值范围过大/GPU 精度不足loss scaling、梯度裁剪、检查算子内部是否用 fp16 累加6.2 左乘右乘搞错的典型症状把 (dX G W^T) 写成 (dX W^T G)形状从 ((m, n)) 变成 ((n, m))反向时立即报错这类错误反而容易发现。更阴险的是写成 (dX G W)形状刚好是 ((m, n)) 没错前提是 (k n)即权重是方阵但数值完全错误。方阵情况下所有维度全部对得上只有数值对不上此时数值梯度几乎是唯一能快速定位的手段。我在处理某跨平台系统的序列模型时就遇到过一个方阵权重下 (dX) 少转置的 bug那个错误在 loss 上表现出来的只是收敛慢了一倍但因为形状全对找人 review 了两轮才通过梯度检查发现。6.3 忽略上游梯度只写了一半链式法则一个很容易犯的错是反向计算只写了“当前函数对输入的局部导数”忘了乘上游梯度。比如错误地把激活层反向写成 (dZ (Z 0)) 而不是 (dZ dA \circ (Z 0))。这在单层网络里只影响 loss 数值形状多层网络里会让梯度逐层放大或缩小产生“假梯度爆炸”。一个辅助记忆法任何反向函数的输入都包括 grad_output 和当前层的输入/权重。grad_output就是上游梯度是别人给你的 (dX)必须乘进你的输出里。如果你的反向函数里没有出现 grad_output那几乎一定漏乘了链式法则因子。6.4 广播场景的梯度求和偏置 (b) 在前向时是 ((1, k))会广播到 ((m, k))。反向时 (db) 必须对 batch 维度求和db G.sum(axis0, keepdimsTrue)如果直接写成db G形状对不上倒是会报错如果写成db G.sum(axis1)形状变成 ((m, 1))而偏置需要 ((1, k))又报错。最危险的是 (k1) 的特殊情况(G) 是 ((m, 1))两种求和方式形状都能对上但只有sum(axis0)是正确的。务必小心。6.5 残差连接分叉合并容易漏项前面 4.2 节已经强调过残差块的反向里 (dX dY dF)。实际操作时我建议在代码注释里写明“分叉合并”并且单独用一个变量保存跳跃路径梯度不要默认dX等于主干路径梯度之后再加。因为多分支求和是反向传播的通用规则前向有多少个分支汇合反向就有多少个分支求导后相加。Transformer Block 里有两处残差Attention 之前、FFN 之前。两处都要做分叉合并。很多 bug 都是从第二个残差没加跳跃梯度开始的因为形状完全一致loss 看起来还在降只是速度和质量差了不少。6.6 混合精度训练中 (dX) 的额外风险LLM 训练基本都用混合精度。这带来一个隐藏问题如果自定义反向 kernel 里用 fp16 做矩阵累加(dX) 的中间累加可能会溢出或下溢。PyTorch 在标准算子内部用 fp32 累加但如果你手写 CUDA kernel 或调用某些底层库函数很容易忽略这一点。排查这类问题最直接的方法是在梯度计算中临时把所有量转成 fp32 跑一次对比两者 loss 曲线。如果 fp32 正常而 fp16 不稳定几乎可以断定是累加精度问题。7. 最后再分享几点实操经验这一路踩过来我最深的体会是(dX) 是一个“为别人服务的梯度”。它不更新任何参数却被每一层当作上游信号。理解它的流转比背下公式重要得多。第一从零手推一遍最后一层 (dX G W^T)并用 NumPy 做数值梯度验证是值得花半小时做的基本功。不要以为会用自动微分框架就不需要理解它。等你需要融合算子、重计算、手写量化 kernel 的时候这些理解会成倍地回报你。当年我遇上某次梯度异常正是靠 gradcheck 一路从损失函数反查到最后一层线性层才发现自定义算子里把转置搞反了。第二给自定义反向实现加形状断言是非常划算的防御。三行assert就能拦住一大批“形状刚好对上但数值全错”的潜在问题。再加一个可选的梯度自检开关训练前对小模型跑一次 gradcheck成本极低却能避免数小时乃至数天的排错。第三计算 FLOPs 时要养成把 backward 算进去的习惯。训练一次完整迭代的 FLOPs前向只占 1/3。无论你是估算训练成本、设计优化策略还是评估硬件利用率把 forward-only 当作训练开销都会高估一倍以上。希望这篇拆解能帮你把 (dX) 这层窗户纸彻底捅破。如果你手头正在写自定义算子试着手推一遍自己的反向传播再用数值梯度验证一次——那个“形状和公式终于对上”的时刻会比照着别人的代码抄要踏实得多。
