从矩阵反向传播到 DP、TP 与 CP：Transformer 训练如何拆开计算

分布式训练容易被缩写切碎：先学反向传播和 Adam，接着遇到 DDP、AllReduce，再往后是 TP、CP、ZeRO。每个概念单独看都有一组公式，但真正困难的是判断：这些 GPU 到底在拆什么，通信回来时为什么有时取平均、有时只求和？

答案可以从一个线性层推出。batch 内不同样本的 loss 可以相加，所以它们对共享参数的梯度也可以相加；矩阵乘法被切成多个分片后，各路径对输入的梯度贡献也要相加；Adam 的状态又是逐参数递推的，所以状态的存放位置可以拆开而不改变公式。

本文沿这条数学主线，从 C=AB 的反向传播走到 Transformer 的 DP、TP、CP 与 ZeRO。重点不是记通信算子的名字，而是建立一个判断方法：先找被拆的维度，再判断每张 GPU 持有的是独立样本、同一运算的部分结果，还是某段模型状态。

本文承接站内前文《从 QKV 到 KV Cache：一步步理解 Transformer 训练与大模型推理》。前文关注一个 Transformer block 如何训练和推理；本文只展开梯度与并行边界。

Transformer 训练张量沿 batch、sequence、hidden 维拆成 DP、CP、TP，并把参数梯度和 Adam 状态按 ZeRO 分片 Fig. 四种“拆分”不能混为一谈：DP 拆独立数据，CP 拆同一序列的 token，TP 拆同一次矩阵运算，ZeRO 拆模型状态的所有权；只有 DP 维增加彼此独立的数据份数。

1 从一个 batch 的 next-token loss 开始
2 矩阵反向传播：转置和求和从哪里来
3 batch 梯度如何汇总为一次参数更新
4 Adam 为什么保存两份历史状态
5 DP：为什么不同 GPU 的梯度可以平均
6 TP：为什么反向传播经常做 SUM 而不是平均
7 CP：沿上下文拆分后 Attention 怎么恢复
8 128K 训练为什么可以配置 1M 推理
9 ZeRO：为什么优化器状态也能沿 DP 拆分
10 DP、TP、CP 如何组合
11 通信成本与部署边界
12 最终工作模型
参考资料

1 从一个 batch 的 next-token loss 开始

GPT 类因果语言模型不会只用一行序列预测最后一个 token。假设输入 batch 是：

X = [
  [t0, t1, t2],
  [t3, t4, t5],
]

训练标签是输入整体左移一位：

Y = [
  [t1, t2, t3],
  [t4, t5, t6],
]

第一行同时产生三项监督：

t0       -> t1
t0,t1    -> t2
t0,t1,t2 -> t3

模型输出 logits 的 shape 是 [B,T,V]，标签是 [B,T]。causal mask 保证位置 t 只能读取自己和过去；虽然一整段序列并行送入 GPU，它仍然没有看到未来答案。

训练 loss 通常对所有有效 token 取平均：

L=\frac{1}{N_{\mathrm{valid}}}\sum_i \ell(\mathrm{token}_i)

这一行定义非常重要。后面 DP 为什么取平均、padding token 为什么需要正确计数，都取决于全局 loss 究竟按什么分母归一化。

实际数据加载通常也不会构造高度重叠的滑动窗口：

[t0,t1,t2]
[t1,t2,t3]
[t2,t3,t4]

这种写法合法，却重复计算了大量相同前缀。预训练一般从连续 token 流切出不同片段，再让每个片段内所有位置参与 next-token loss。

2 矩阵反向传播：转置和求和从哪里来

考虑最常见的线性层：

C=AB,\qquad A\in\mathbb{R}^{m\times n},\quad B\in\mathbb{R}^{n\times p},\quad C\in\mathbb{R}^{m\times p}

假设上游已经给出：

G_C=\frac{\partial L}{\partial C}

反向传播公式是：

\frac{\partial L}{\partial A}=G_CB^\mathsf{T},\qquad \frac{\partial L}{\partial B}=A^\mathsf{T}G_C

也就是 $G_A=G_CB^\mathsf{T}$ 、 $G_B=A^\mathsf{T}G_C$ 。

这两个转置不是为了背 shape 而人为加上的。把矩阵乘法写到元素级：

C_{ij}=\sum_k A_{ik}B_{kj}

固定 $A_{pq}$ 。它只影响 $C$ 的第 $p$ 行，并且：

\frac{\partial C_{pj}}{\partial A_{pq}}=B_{qj}

$A_{pq}$ 会经过所有 $C_{pj}$ 影响最终 loss，多变量链式法则要求把这些路径相加：

\begin{aligned} \frac{\partial L}{\partial A_{pq}} &=\sum_j \frac{\partial L}{\partial C_{pj}} \frac{\partial C_{pj}}{\partial A_{pq}} \\ &=\sum_j (G_C)_{pj}B_{qj} \end{aligned}

最后一行恰好是 $(G_CB^\mathsf{T})_{pq}$ 。同理可以得到：

\frac{\partial L}{\partial B_{pq}} =\sum_i A_{ip}(G_C)_{iq} =(A^\mathsf{T}G_C)_{pq}

这里最值得保留的不是公式本身，而是链式法则的工作方式：

一条路径：沿途局部导数相乘
多条路径：所有路径贡献相加

这个“路径贡献相加”会在 batch、DP 和 TP 中反复出现，只是路径代表的对象不同。

3 batch 梯度如何汇总为一次参数更新

用一个两样本线性层把前向和反向完整走一遍：

X = [[1,2],
     [3,4]]

W = [[ 2],
     [-1]]

Y = [[1],
     [3]]

前向传播：

Y_hat = X W = [[0],
               [2]]

error = Y_hat - Y = [[-1],
                     [-1]]

使用 batch 平均的平方损失：

L=\frac{1}{2B}\sum_i(\hat Y_i-Y_i)^2

B=2，所以输出梯度是：

G_{\hat Y}=\frac{\hat Y-Y}{B} =\begin{bmatrix}-0.5\\-0.5\end{bmatrix}

权重梯度：

\begin{aligned} G_W &=X^\mathsf{T}G_{\hat Y} \\ &=\begin{bmatrix}1&3\\2&4\end{bmatrix} \begin{bmatrix}-0.5\\-0.5\end{bmatrix} \\ &=\begin{bmatrix}-2\\-3\end{bmatrix} \end{aligned}

展开后可以看到每个样本各贡献一项：

sample 0: [[1],[2]] × (-0.5) = [[-0.5],[-1.0]]
sample 1: [[3],[4]] × (-0.5) = [[-1.5],[-2.0]]

sum: [[-2],[-3]]

因此 $X^\mathsf{T}G_{\hat Y}$ 已经在矩阵乘法里完成了 batch 维的梯度汇总。优化器不会为 batch 中每一行分别更新参数，而是等整个 batch 得到一个 $G_W$ 后，只执行一次 step()。

如果 SGD 学习率是 0.1：

W_{\mathrm{new}}=W-0.1G_W =\begin{bmatrix}2.2\\-0.7\end{bmatrix}

4 Adam 为什么保存两份历史状态

SGD 对所有参数使用同一个学习率：

\theta_t=\theta_{t-1}-\alpha g_t

考虑一个曲率差异很大的损失：

L(x,y)=\frac{1}{2}(100x^2+y^2),\qquad \nabla L=\begin{bmatrix}100x&y\end{bmatrix}

从 [x,y]=[1,1] 出发，梯度是 [100,1]。如果 alpha=0.1，SGD 第一步变成：

\begin{bmatrix}1&1\end{bmatrix} -0.1\begin{bmatrix}100&1\end{bmatrix} =\begin{bmatrix}-9&0.9\end{bmatrix}

y 的步长合适，x 却直接越过最低点并发散。若把学习率降到足以稳定 x，y 又会非常慢。问题不是梯度算错，而是两个参数的梯度尺度差了 100 倍，却共享同一个步长。

Adam 为每个参数保存两份历史：

\begin{aligned} m_t&=\beta_1m_{t-1}+(1-\beta_1)g_t \\ v_t&=\beta_2v_{t-1}+(1-\beta_2)g_t^2 \end{aligned}

m 记录带符号的平均梯度，回答“最近总体往哪边走”；v 记录梯度平方的平均，回答“这个参数的梯度通常有多大”。常见设置是 beta1=0.9、beta2=0.999，本文手算采用 beta2=0.99 以便看清系数职责。

初始 m_0=v_0=0 会把前几步的滑动平均拉向 0。展开一阶矩：

m_t=(1-\beta_1)\sum_{k=1}^{t}\beta_1^{t-k}g_k

当前所有梯度权重之和只有 $1-\beta_1^t$ ，所以要重新归一化：

\hat m_t=\frac{m_t}{1-\beta_1^t},\qquad \hat v_t=\frac{v_t}{1-\beta_2^t}

这就是 bias correction。它不是对数据分布做校正，而是在补偿“滑动平均从全零状态启动”造成的系统性偏小。

最终更新为：

\theta_t=\theta_{t-1} -\alpha\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}

第一步的梯度 [100,1] 经修正后：

\hat m=\begin{bmatrix}100&1\end{bmatrix},\qquad \hat v=\begin{bmatrix}10000&1\end{bmatrix}

\frac{\hat m}{\sqrt{\hat v}} =\begin{bmatrix}1&1\end{bmatrix}

Adam 在第一步几乎只保留梯度符号：对梯度典型值大的 x 自动使用更小的有效学习率，对梯度小的 y 使用更大的有效学习率。后续步骤中，m 还会让方向反复变化的梯度彼此抵消，从而抑制震荡。

这说明 Adam 通常更容易适应不同参数的梯度尺度，不代表它在所有任务上都优于精心调参的 SGD。它解决的是优化路径和调参敏感度，不是无条件改善最终泛化。

5 DP：为什么不同 GPU 的梯度可以平均

假设有 N 张 GPU，每张卡处理 b 个样本，全局 batch 是 $B=Nb$ 。全局平均 loss 可以拆成：

\begin{aligned} L_{\mathrm{global}} &=\frac{1}{Nb}\sum_{r=1}^{N}\sum_{i=1}^{b}\ell_{r,i} \\ &=\frac{1}{N}\sum_{r=1}^{N}L_r \end{aligned}

求导的线性性质给出：

\begin{aligned} g_{\mathrm{global}} &=\nabla L_{\mathrm{global}} \\ &=\frac{1}{N}\sum_{r=1}^{N}\nabla L_r \\ &=\frac{1}{N}\sum_{r=1}^{N}g_r \end{aligned}

因此 DP 的正确性不是“不同梯度平均后大概更稳定”，而是：在所有 rank 使用同一份参数、每卡 local batch 相等、loss 使用相同平均规则时，local gradient 的平均就是全局 batch 的精确梯度。

回到上一节的两个样本，可以把它们拆到两张卡：

rank 0: sample [1,2] -> local dW = [[-1],[-2]]
rank 1: sample [3,4] -> local dW = [[-3],[-4]]

AllReduce mean:
([[ -1],[-2]] + [[-3],[-4]]) / 2
= [[-2],[-3]]

结果与单卡完整 batch 完全一致。验证脚本保存在 examples/transformer-training-math/training_math.py，不依赖第三方库：

python3 examples/transformer-training-math/training_math.py

$纯 Python 验证完整 batch 梯度等于 DP 平均梯度、TP 分片输入梯度之和等于完整矩阵反传，并验证 Adam 首步尺度归一化$ Fig. 两个 PASS 对应两种不同的等价关系：DP 对独立样本梯度取平均，TP 对同一输入经过不同参数分片的梯度贡献求和；二者不能仅凭都用了 collective 就混成一种操作。

DDP 的实际执行顺序是：

每个 rank 持有相同参数 W
  -> 各自处理不同 local batch
  -> backward 产生 local gradients
  -> gradient buckets 执行 AllReduce
  -> 每个 rank 得到相同 global gradient
  -> 每个 rank 本地执行相同 optimizer.step()

只要参数、全局梯度和优化器状态原本一致，相同的确定性更新会让各 rank 参数继续一致。Adam 必须在梯度同步之后执行；先让各卡用 local gradient 更新 Adam，再平均参数，通常不等价，因为平方和除法不是线性运算：

\left(\frac{g_0+g_1}{2}\right)^2 \ne\frac{g_0^2+g_1^2}{2}

DP 还有三个边界：

每卡有效样本或有效 token 数不同时，应按数量加权，不能平均各卡 local mean。
BatchNorm、全局对比学习等跨样本算子需要额外同步，前向本身不再是完全独立的样本分支。
global batch 变大后，梯度虽然仍然正确，但优化轨迹会改变，学习率和训练步数可能需要调整。

6 TP：为什么反向传播经常做 SUM 而不是平均

TP 不拆独立样本，而是让多张 GPU 合作完成同一次矩阵乘法。考虑：

y=xW,\qquad x=\begin{bmatrix}1&2\end{bmatrix},\qquad W=\begin{bmatrix}1&3\\2&4\end{bmatrix},\qquad y=\begin{bmatrix}5&11\end{bmatrix}

将 W 按列拆到两张卡：

rank 0: W0 = [[1],[2]] -> y0 = x W0 = [5]
rank 1: W1 = [[3],[4]] -> y1 = x W1 = [11]

y = concat(y0,y1)

假设上游梯度：

G_y=\begin{bmatrix}1&2\end{bmatrix}

输出也是按列分片的，所以每张卡拿到对应部分，独立计算自己的参数梯度：

G_{W_0}=x^\mathsf{T}\cdot1 =\begin{bmatrix}1\\2\end{bmatrix},\qquad G_{W_1}=x^\mathsf{T}\cdot2 =\begin{bmatrix}2\\4\end{bmatrix}

这些参数梯度属于不同权重分片，不应该彼此平均。每张 GPU 保留并更新自己的分片即可。

但输入 x 同时经过了两条计算路径：

G_x^{(0)}=1\cdot W_0^\mathsf{T} =\begin{bmatrix}1&2\end{bmatrix},\qquad G_x^{(1)}=2\cdot W_1^\mathsf{T} =\begin{bmatrix}6&8\end{bmatrix}

上标 $(0)$ 、 $(1)$ 表示两个 TP rank 对同一个输入梯度的局部贡献，并不是两个不同输入 $x_0,x_1$ 。

根据多变量链式法则，完整输入梯度必须汇总所有路径：

G_x=G_x^{(0)}+G_x^{(1)} =\begin{bmatrix}7&10\end{bmatrix}

直接使用完整矩阵计算也是：

G_x=G_yW^\mathsf{T} =\begin{bmatrix}7&10\end{bmatrix}

因此 Column Parallel 线性层在反向传播中常需要对 G_x 做 TP SUM。这里不能除以 TP size：每张卡拿到的是完整答案的一部分，而不是对同一个答案的重复估计。

若把 W 按行拆分，则输入也按特征维拆开：

x=\begin{bmatrix}x_0&x_1\end{bmatrix},\qquad W=\begin{bmatrix}W_0\\W_1\end{bmatrix},\qquad y=x_0W_0+x_1W_1

Row Parallel 在前向传播中需要把局部输出求和；反向传播得到的 $G_{x_0}$ 、 $G_{x_1}$ 则天然保持分片。Transformer 通常把两种线性层配对：

MLP:
Column Parallel W_up
  -> hidden 保持分片
  -> Row Parallel W_down
  -> SUM partial outputs

Attention:
Column Parallel QKV（不同 rank 负责不同 heads）
  -> local attention
  -> Row Parallel output projection
  -> SUM partial outputs

TP 的参数梯度在本 TP rank 上只对应一个参数分片。如果同时启用 DP，保存相同分片的 DP ranks 还要再对这份参数梯度做 DP AllReduce。

7 CP：沿上下文拆分后 Attention 怎么恢复

Context Parallel 沿 [B,T,D] 的 T 维拆同一条长序列：

sequence = [t0,t1,t2,t3,t4,t5,t6,t7]

CP rank 0: [t0,t1,t2,t3]
CP rank 1: [t4,t5,t6,t7]

LayerNorm 和 MLP 都逐 token 工作，各 rank 可以直接处理本地 token。Attention 不行，因为后半段 query 需要读取前半段的 K/V：

Q1 for t4..t7 needs K/V from t0..t7 under causal mask

最直接的实现是 AllGather K/V，让每张卡拿到完整上下文，但这会重新占用完整 K/V 显存。Ring Attention 类实现让 K/V block 在 CP ranks 之间循环：

Q 留在本地
K/V block 沿 ring 传递
每收到一块，就更新本地 query 的 online softmax 状态

Softmax 不能对每个 block 分别归一化后直接相加，因为分母必须覆盖所有可见 key。实现需要维护每个 query 的运行最大值、指数和与加权输出，最终结果才等价于一次完整 attention。

反向传播时，远程 K/V 被本地 query 使用过，因此产生的 dK/dV 贡献要沿通信路径送回并累加到 K/V 所属 rank。CP 和 TP 一样是在恢复同一次模型计算的完整依赖，不是在平均不同训练样本。

CP 的主要收益是让每张卡只保存约 T/CP 的 token 激活。它不改变 global batch：同一 CP group 的 GPU 在合作处理同一条序列。

8 128K 训练为什么可以配置 1M 推理

训练序列长度与推理引擎允许的最大上下文不是同一个概念。线性层参数依赖 D，不依赖 T，因此模型权重本身不会因为输入从 128K 变成 1M 而 shape 不匹配。真正的边界在位置编码、Attention 实现、KV cache 容量与模型是否学会利用远距离信息。

使用 RoPE 的模型可以公式化地产生训练范围外的位置，但直接外推通常会退化。常见做法是 RoPE scaling，将更长的位置映射到较熟悉的旋转频率范围。128K 扩到 1M，长度比例约为：

\frac{1{,}000{,}000}{128{,}000}\approx7.8

简单线性缩放可以理解为把位置 p 映射成 p/7.8；实际的 YaRN、LongRoPE 等方法会区别处理不同频率，避免局部位置分辨率被统一压缩。

技术上能接收 1M，不代表模型能可靠利用 1M：

max configured context != effective context

训练只见过 128K 依赖时，1M 范围内可能出现检索下降、lost in the middle 或跨段推理失败。是否真正具备 1M 能力，需要多位置、多目标检索与长程推理评测，不能只看服务配置。

不直接全量训练 1M 的主要原因是成本。Dense Attention 从 128K 增长到 1M，长度增加约 7.8 倍，理论计算量按 $T^2$ 增长约：

7.8^2\approx61

FlashAttention 避免显式保存完整 [T,T] 矩阵，却不会消除 dense attention 的平方计算。推理只做前向、batch 可以更小，还能使用 chunked prefill、KV cache 量化和 CP；因此“少量 1M 推理”通常比“持续 1M 训练”更容易承担。

9 ZeRO：为什么优化器状态也能沿 DP 拆分

普通 DDP 会在每个 rank 复制完整参数、梯度和 Adam 状态。Adam 对每个参数元素独立更新：

\begin{aligned} m_i&\leftarrow\beta_1m_i+(1-\beta_1)g_i \\ v_i&\leftarrow\beta_2v_i+(1-\beta_2)g_i^2 \\ w_i&\leftarrow\operatorname{Adam}(w_i,g_i,m_i,v_i) \end{aligned}

更新 w_i 只需要对应的 g_i,m_i,v_i，不需要另一个参数的状态。因此一个 DP group 可以指定：

rank 0 owns: w0,w1 and their m/v
rank 1 owns: w2,w3 and their m/v

只要 owner 得到正确的全局梯度分片，它计算出的新状态与单卡 Adam 完全相同。其他 rank 没有必要保存重复的 m_i,v_i。随后同步更新后的参数，下一轮前向仍使用一致模型。

这就是 ZeRO 各阶段的递进关系：

方式	参数	梯度	优化器状态
DDP	复制	复制	复制
ZeRO-1	复制	复制	分片
ZeRO-2	复制	分片	分片
ZeRO-3 / FSDP	分片	分片	分片

ZeRO-2 可以把 AllReduce 分解成：

AllReduce = ReduceScatter + AllGather

ReduceScatter 后，每张卡已经拿到自己负责参数的全局梯度分片，可以直接更新本地 optimizer states，不需要重新 AllGather 完整梯度。更新完成后再 AllGather 参数。

ZeRO-3 连参数也不常驻完整副本。某一层计算前 AllGather 该层参数，完成计算后释放非本地分片；反向传播对梯度 ReduceScatter，owner 更新自己的参数和状态。

所以 ZeRO 并没有把 Adam 数学公式拆成近似算法，它拆的是状态所有权。全局梯度裁剪、LAMB 层级范数、Shampoo 预条件器等跨参数操作仍需要额外通信，不能仅靠逐参数独立性完成。

10 DP、TP、CP 如何组合

假设 8 张 GPU，配置为 $TP=2$ 、 $CP=2$ 、 $DP=2$ ，因此：

N_{\mathrm{GPU}}=TP\times CP\times DP=8

其中 TP×CP=4 张 GPU 合作完成一个模型副本：TP 拆 hidden/heads，CP 拆 sequence。另一个 4 卡组处理不同数据，形成第二个 DP replica。

以只有 TP 和 DP 的 4 卡配置为例：

                 TP shard 0      TP shard 1
DP replica 0       rank 0          rank 1
DP replica 1       rank 2          rank 3

通信组是：

TP groups: [0,1], [2,3]
DP groups: [0,2], [1,3]

rank 0 与 rank 2 保存相同参数分片、处理不同数据，因此对该分片的参数梯度做 DP AllReduce。rank 0 与 rank 1 保存不同参数分片、处理相同数据，因此在层内交换局部输出或输入梯度。

global batch 只乘 DP，不乘 TP 或 CP：

B_{\mathrm{global}} =B_{\mathrm{micro}}\times DP\times N_{\mathrm{accumulation}}

若再加入 Pipeline Parallel，常见设备数关系是：

\mathrm{world\ size}=DP\times TP\times CP\times PP

PP 拆的是 Transformer layers，各 stage 之间传激活和激活梯度，同样不增加独立数据副本数。

11 通信成本与部署边界

TP 的通信通常比 DP 更难隐藏。TP 在每层前向和反向内部交换激活或局部结果，collective 没完成时，下一段依赖计算往往无法继续。DP 的参数梯度每 step 同步一次，框架可以按 bucket 在 backward 过程中提前启动异步 AllReduce，与前面层的梯度计算重叠。

这不意味着 TP 的总字节数永远大于 DP。DP 要同步整个参数或参数分片的梯度，模型很大时同样昂贵。更准确的比较是：

维度	DP	TP	CP
主要通信对象	参数梯度	激活、局部输出、输入梯度	K/V block 与其梯度
通信频率	每参数每 step 一次	几乎每层多次	几乎每个 attention layer
关键路径敏感度	较容易与 backward 重叠	高	高
常见放置	可跨节点扩展	优先节点内 NVLink/NVSwitch	优先高速互联

因此常见部署原则是：使用满足显存和单卡 GEMM 规模要求的最小 TP，在高速节点内完成 TP/CP，再沿节点扩展 DP。TP 增大后每卡矩阵乘法变小，通信占比上升；它不是越大越好。

12 最终工作模型

整套机制可以压缩成四个判断。

第一，先看全局目标如何定义：

global mean loss
-> gradient is the weighted mean of sample/token gradients
-> DP averages independent data contributions

第二，沿计算图判断一个变量有几条下游路径：

C = A B
-> dA = dC B^T
-> dB = A^T dC
-> path contributions are summed

TP 将同一矩阵运算拆成多条设备路径，所以 partial output 或 dX 通常求和，不能像 DP 那样随手除以 world size。

第三，区分张量维度与状态所有权：

DP   splits B: independent samples
CP   splits T: tokens of one context
TP   splits D: hidden/heads/parameter dimensions
ZeRO splits W, dW, m, v ownership

第四，优化器始终消费同步后的正确梯度：

backward
-> TP/CP restore complete model-path contributions
-> DP forms the global data gradient
-> ZeRO owner updates its state shard
-> parameters are made available for the next forward

分布式训练并没有改变反向传播的数学目标；它做的是把原本单设备上的求和、矩阵乘法和状态存储分配给不同 GPU，再通过通信恢复同一个结果。 判断一种并行是否正确，最终都要回到三件事：拆的是哪个维度、局部值代表完整结果还是部分贡献、通信后是否恢复了未拆分计算的 loss 与梯度。