@TOC

一、背景概述

随着大语言模型参数规模持续增长，单卡显存已无法满足训练需求。分布式训练成为突破单卡瓶颈的核心手段，其中张量并行（Tensor Parallelism, TP） 与全分片数据并行（Fully Sharded Data Parallel, FSDP） 是两种互补的关键技术：TP将单层权重切分到多卡以加速计算，FSDP将参数、梯度与优化器状态在数据并行维度分片以降低显存压力。将二者正交组合为2D并行，可在Atlas 800I A2 NPU平台上高效支撑大规模模型训练。

本文以一个Transformer FFN层和Attention层为例，在2D并行（FSDP + TP）环境下，逐步推导前向传播与反向传播的完整计算流程，涵盖维度变化、矩阵乘法分解、通信时机与类型，以及结果统筹机制。所有集合通信均基于昇腾HCCL通信库实现，计算算子运行于NPU之上。

二、物理环境与模型维度设定

2.1 物理环境

配置项	设定值
NPU总数	4块（Atlas 800I A2）
逻辑网格	2（TP）× 2（DP）
TP组	NPU0-NPU1（处理同一条数据）；NPU2-NPU3（处理另一条数据）
DP组	NPU0-NPU2（共享w1_A、w2_A分片）；NPU1-NPU3（共享w1_B、w2_B分片）

2.2 模型维度设定

参数	符号	数值
输入特征维度	dim	4
隐藏层维度	hidden_dim	8
权重矩阵（升维）	w1	形状 [4, 8]
权重矩阵（门控）	w3	形状 [4, 8]
权重矩阵（降维）	w2	形状 [8, 4]
输入Token向量	X	形状 [1, 4]
数据并行另一输入	X'	形状 [1, 4]（独立样本）

三、初始状态：物理存储布局

训练开始前，所有权重经过两层切分：先沿TP维度切分，再沿FSDP维度二次切分，最终分片分布在各个NPU的HBM中。

3.1 第一步：沿TP维度切分（逻辑切分）

权重	原始形状	TP切分后	切分方式
w1	[4, 8]	w1_A：[4, 4]（输出维0~3）；w1_B：[4, 4]（输出维4~7）	列并行（按列切）
w3	[4, 8]	w3_A：[4, 4]；w3_B：[4, 4]	列并行（按列切）
w2	[8, 4]	w2_A：[4, 4]（输入维0~3）；w2_B：[4, 4]（输入维4~7）	行并行（按行切）

3.2 第二步：沿FSDP维度二次切分（物理存储）

每个TP分片再沿DP维度均分（此处简化为按行均分）：

NPU	存储的权重分片	形状	说明
NPU0	w1_A前半 / w3_A前半 / w2_A前半	[2, 4] / [2, 4] / [2, 4]	TP左半 + DP前半
NPU1	w1_B前半 / w3_B前半 / w2_B前半	[2, 4] / [2, 4] / [2, 4]	TP右半 + DP前半
NPU2	w1_A后半 / w3_A后半 / w2_A后半	[2, 4] / [2, 4] / [2, 4]	TP左半 + DP后半
NPU3	w1_B后半 / w3_B后半 / w2_B后半	[2, 4] / [2, 4] / [2, 4]	TP右半 + DP后半

关键点：没有任何一块NPU持有完整的TP分片（如完整的w1_A），每个TP分片被DP组内的两块NPU各分一半。

四、前向传播：完整分步推导

4.1 阶段1：FSDP All-Gather 参数还原

当模型运行到该FFN层时，FSDP的pre-forward钩子触发，在DP维度上收集参数以恢复完整的TP分片。该通信基于HCCL的All-Gather原语实现。

NPU0的操作（与同DP组的NPU2通信）：

• 发起 All-Gather 通信，从NPU2获取w1_A缺失的另一半（形状[2, 4]）。
• 拼接后，NPU0获得完整的 w1_A（形状 [4, 4]）。
• 同理获得完整的 w3_A（[4, 4]）和 w2_A（[4, 4]）。

NPU1的操作（与同DP组的NPU3通信）：

• 通过All-Gather获得完整的 w1_B、w3_B、w2_B（均形状 [4, 4]）。

此时各NPU持有的完整计算权重：

NPU	持有的完整TP分片	形状
NPU0	w1_A、w3_A、w2_A	各 [4, 4]
NPU1	w1_B、w3_B、w2_B	各 [4, 4]
NPU2	w1_A、w3_A、w2_A	各 [4, 4]
NPU3	w1_B、w3_B、w2_B	各 [4, 4]

通信类型：All-Gather（广播收集），基于HCCL实现。
通信量：每个NPU接收约 3 × [2, 4] = 24 个浮点数（本例中）。

4.2 阶段2：列并行升维（计算w1和w3）

输入向量 X = [x0, x1, x2, x3]（形状 [1, 4]）被同时广播到TP组内的NPU0和NPU1。此阶段无需任何通信，所有计算在本地NPU上完成。

4.2.1 NPU0上的计算（负责隐藏层前半部分，维度0~3）

步骤1：计算h1_A

h1_A = X @ w1_A
     = [x0, x1, x2, x3] × [4×4矩阵]
     = [h1_0, h1_1, h1_2, h1_3]    形状: [1, 4]

对应完整h1（形状[1, 8]）的前4维。

步骤2：计算h3_A

h3_A = X @ w3_A
     = [h3_0, h3_1, h3_2, h3_3]    形状: [1, 4]

对应完整h3（形状[1, 8]）的前4维。

步骤3：SwiGLU激活

Z_A = SiLU(h1_A) * h3_A
    = [z0, z1, z2, z3]             形状: [1, 4]

对应完整中间结果Z（形状[1, 8]）的前半部分。

4.2.2 NPU1上的计算（负责隐藏层后半部分，维度4~7）

步骤1：计算h1_B

h1_B = X @ w1_B
     = [h1_4, h1_5, h1_6, h1_7]    形状: [1, 4]

对应完整h1的后4维。

步骤2：计算h3_B

h3_B = X @ w3_B
     = [h3_4, h3_5, h3_6, h3_7]    形状: [1, 4]

步骤3：SwiGLU激活

Z_B = SiLU(h1_B) * h3_B
    = [z4, z5, z6, z7]             形状: [1, 4]

对应完整中间结果Z的后半部分。

4.2.3 升维阶段结果统筹

完整的中间张量为：

Z = [Z_A | Z_B] = [z0, z1, z2, z3, z4, z5, z6, z7]    形状: [1, 8]

物理上，Z_A存放在NPU0，Z_B存放在NPU1。通过列并行的数学特性，无需通信即完成了输出的自然分割。

4.3 阶段3：行并行降维（计算w2）

现在需要计算 Y = Z @ w2，其中Z = [Z_A | Z_B]，w2 = [w2_A ; w2_B]。

根据分块矩阵乘法法则：

Y = [Z_A | Z_B] × [w2_A ; w2_B]
  = (Z_A × w2_A) + (Z_B × w2_B)

这是一个线性分配律的应用，使得两个NPU可以独立计算局部累加项，最后通过一次All-Reduce Sum合并。

4.3.1 NPU0上的计算

持有 Z_A（形状[1, 4]）和 w2_A（形状[4, 4]）：

Partial_0 = Z_A @ w2_A
          = [z0, z1, z2, z3] × [4×4矩阵]
          = [p0_0, p0_1, p0_2, p0_3]     形状: [1, 4]

这是最终输出Y的第一个累加项。

4.3.2 NPU1上的计算

持有 Z_B（形状[1, 4]）和 w2_B（形状[4, 4]）：

Partial_1 = Z_B @ w2_B
          = [z4, z5, z6, z7] × [4×4矩阵]
          = [p1_0, p1_1, p1_2, p1_3]     形状: [1, 4]

这是最终输出Y的第二个累加项。

4.3.3 结果的最终统筹（All-Reduce Sum）

NPU0和NPU1发起一次 All-Reduce（求和） 通信，该操作由HCCL提供：

Y = Partial_0 + Partial_1
  = [p0_0+p1_0, p0_1+p1_1, p0_2+p1_2, p0_3+p1_3]     形状: [1, 4]

通信完成后：

• NPU0获得：完整的最终输出 Y（形状[1, 4]）
• NPU1获得：完整的最终输出 Y（形状[1, 4]）

通信类型：All-Reduce（求和归约），基于HCCL。
通信量：约 [1, 4] = 4 个浮点数（本例中）。
数学原理：分块矩阵乘法的分配律是TP能够高效运行的理论基石。

4.4 阶段4：数据并行维度的同步执行

在NPU0和NPU1计算X的同时：

• NPU2和NPU3执行完全相同的流程，但输入是独立样本X'（形状[1, 4]）。
• NPU2持有w1_A、w3_A、w2_A，NPU3持有w1_B、w3_B、w2_B。
• 计算得出X'对应的最终输出Y'（形状[1, 4]）。

五、反向传播与梯度更新

5.1 损失回传与梯度计算

假设损失函数L已计算完毕，dY = ∂L/∂Y（形状[1, 4]）开始回传。

5.2 TP组内梯度同步（All-Reduce）

针对w2的梯度：

由于NPU0和NPU1都有完整的Y，各自可以计算出w2_A和w2_B的梯度：

• NPU0计算 d(w2_A) = Z_A^T @ dY（形状[4, 4]）
• NPU1计算 d(w2_B) = Z_B^T @ dY（形状[4, 4]）

为了保持权重一致性，TP组内对梯度进行 All-Reduce，确保两张卡上同一切片的梯度完全一致。

针对X的梯度（传递给上一层）：

根据链式法则，X的梯度需要从w1和w3两条路径回传：

dX = dY @ (w1_A)^T + dY @ (w1_B)^T

• NPU0计算 dX_A = dY @ (w1_A)^T
• NPU1计算 dX_B = dY @ (w1_B)^T
• 发起 All-Reduce Sum：dX = dX_A + dX_B
• 两张卡均获得完整的dX（形状[1, 4]），传给上一层。

注意：反向传播中同样包含一次All-Reduce Sum，与前向形成对称结构。对于一个完整的线性层，TP共需两次All-Reduce Sum（前向一次，反向一次）。 这些操作均由HCCL在NPU间高效完成。

5.3 FSDP梯度分片（Reduce-Scatter）

在DP维度上，FSDP触发post-backward钩子：

• NPU0和NPU2都计算出了针对w1_A的完整梯度（形状[4, 4]）。
• 执行 Reduce-Scatter：将两份梯度（来自不同数据样本）相加并切分。
• 最终：
  • NPU0保存w1_A前半部分的梯度（形状[2, 4]）
  • NPU2保存w1_A后半部分的梯度（形状[2, 4]）

同理，NPU1和NPU3对w1_B的梯度做Reduce-Scatter。

5.4 优化器参数更新

每个NPU只用自己持有的梯度分片更新自己持有的参数分片：

NPU	负责更新的参数	更新依据
NPU0	w1_A前半、w2_A前半	对应的梯度分片
NPU1	w1_B前半、w2_B前半	对应的梯度分片
NPU2	w1_A后半、w2_A后半	对应的梯度分片
NPU3	w1_B后半、w2_B后半	对应的梯度分片

关键优势：优化器状态（如Adam的一阶矩和二阶矩）只需为分片参数保存，相比DDP节省约 3倍 的优化器HBM占用（因为Adam需保存2倍于参数的状态量）。在Atlas 800I A2上，这种节省对于大模型训练至关重要。

六、完整流程一览表

阶段	执行主体	核心计算（形状变化）	通信类型	通信量	结果统筹目的
参数准备	NPU0, NPU1	向DP对端收集权重分片	All-Gather	中等	恢复完整的TP分片
升维（w1/w3）	NPU0	X[1,4] @ w1_A[4,4] = Z_A[1,4]	无	0	自然分割隐藏层输出
升维（w1/w3）	NPU1	X[1,4] @ w1_B[4,4] = Z_B[1,4]	无	0	自然分割隐藏层输出
降维（w2）	NPU0	Partial_0[1,4] = Z_A[1,4] @ w2_A[4,4]	All-Reduce(Sum)	小	分块乘法累加
降维（w2）	NPU1	Partial_1[1,4] = Z_B[1,4] @ w2_B[4,4]	All-Reduce(Sum)	小	两张卡同时获得完整Y[1,4]
反向-w2梯度	NPU0, NPU1	对w2梯度做All-Reduce	All-Reduce	中等	保持TP组内权重同步
反向-dX梯度	NPU0, NPU1	dX_A + dX_B	All-Reduce(Sum)	小	恢复完整dX[1,4]传给上一层
参数更新	NPU0, NPU2	对w1_A梯度做Reduce-Scatter	Reduce-Scatter	中等	分片存储梯度，节省HBM

上表中所有集合通信操作均由昇腾HCCL通信库在NPU间高效实现。

七、补充说明：TP的核心数学原理

7.1 TP能够成立的理论基石

张量并行（TP）能够高效运行，其核心在于分块矩阵乘法的线性分配律：

对于列并行（Column Parallelism，如w1的切分）：

X @ [W1, W2] = [X@W1, X@W2]

输出结果天然按列分割，无需通信即可完成切分。

对于行并行（Row Parallelism，如w2的切分）：

[X1, X2] @ [W1; W2] = X1@W1 + X2@W2

输出结果需要将各NPU的局部累加项相加（All-Reduce Sum），只需一次通信。

7.2 通信次数的精确结论

• 前向传播：行并行层末尾需要 1次 All-Reduce Sum。
• 反向传播：列并行层的梯度计算需要 1次 All-Reduce Sum。
• 总计：每个线性层完整的前向+反向需要 2次 All-Reduce Sum。

这个通信次数不随模型层数增加而增加（每层固定2次），这是TP高效的根源。通信量的大小与隐藏层维度hidden_dim成正比。

7.3 FSDP与TP的本质关系

• FSDP本身不包含TP：两者是正交且独立的并行维度。
• FSDP解决存储问题：将参数、梯度、优化器状态在DP维度分片，节省HBM。
• TP解决计算问题：将单层权重在NPU间切分，加速计算并分摊单层HBM。
• 协同工作：通过All-Gather（计算前恢复TP分片）和Reduce-Scatter（计算后分片存储梯度）完成解耦。

八、特殊情况简述

除了上述标准的1D TP实现外，还存在以下特殊变体，这些变体在Atlas 800I A2上同样可以获得支持：

1. 通信模式变体：在与专家并行（EP）混合使用时，可通过skip_allreduce跳过All-Reduce，推迟合并操作；在2D/3D TP中，通信模式可能扩展为All-Reduce + All-Gather的组合。
2. 切分维度扩展：2D和3D张量并行将权重切分到二维或三维设备网格上，每个NPU只负责更细粒度的小块计算，适用于超大规模模型。
3. 特殊层处理：Embedding层按词表维度切分，使用All-to-All通信；CrossEntropy Loss层按类别维度切分，使用All-Reduce汇总损失。
4. 计算通信融合：在Atlas 800I A2上，可通过算子融合将GEMM计算的尾声（epilogue）与All-Reduce通信融合到单个NPU算子中，将通信开销隐藏在计算过程中，进一步提升性能。

九、Attention计算案例（昇腾NPU环境）

9.1 基础设定（Attention专用）

• 模型维度：d_model = 4，注意力头数 num_heads = 2，每个头的维度 head_dim = 2。
• 输入：X，形状为 [seq_len=2, d_model=4]。包含两个Token：Token_A 和 Token_B。
• TP切分（2路）：按注意力头切分。NPU0负责第0个头，NPU1负责第1个头。
• 权重切分：
  • W_Q, W_K, W_V：原始形状 [4, 4]，列并行切分。每卡分到形状 [4, 2]（负责映射出2维的头向量）。
  • W_O（输出投影）：原始形状 [4, 4]，行并行切分。每卡分到形状 [2, 4]。

9.2 阶段一：Q/K/V 投影（列并行，无通信）

输入 X [2, 4] 同时进入NPU0和NPU1。每张卡拥有完整的序列（2个Token），但只负责计算部分头。

在 NPU0 上（计算 Head 0，输出维度第 0~1 维）：

• 计算 Q：Q_A = X @ W_Q_A → 形状 [2, 4] × [4, 2] = [2, 2]
• 计算 K：K_A = X @ W_K_A → 形状 [2, 2]
• 计算 V：V_A = X @ W_V_A → 形状 [2, 2]

物理意义：NPU0 拥有 Q 矩阵中第 0~1 列（Head 0 的数据）。

在 NPU1 上（计算 Head 1，输出维度第 2~3 维）：

• 计算 Q：Q_B = X @ W_Q_B → 形状 [2, 2]
• 计算 K：K_B = X @ W_K_B → 形状 [2, 2]
• 计算 V：V_B = X @ W_V_B → 形状 [2, 2]

物理意义：NPU1 拥有 Q 矩阵中第 2~3 列（Head 1 的数据）。

关键结论：此阶段 0次通信。因为列并行使得输出张量天然按“头”维度分割，每张卡独立拥有完整的序列长度。

9.3 阶段二：缩放点积注意力（本地计算，无通信）

现在，每张卡在本地独自完成 Attention 的核心公式 Softmax(Q @ K^T / √d) @ V。注意：Softmax 是在 seq_len 维度上对每个 Token 做的，不涉及跨卡交互。

在 NPU0 上（Head 0 的计算）：

1. 计算注意力分数：Scores_A = Q_A [2,2] @ K_A^T [2,2] = [2, 2]。
   • 结果矩阵中，第 i 行第 j 列表示第 i 个 Token 对第 j 个 Token 的注意力分数。
2. 缩放与 Softmax：Attn_A = Softmax(Scores_A / √2) → 形状保持 [2, 2]。
3. 加权求和：Output_Head_A = Attn_A [2,2] @ V_A [2,2] = [2, 2]。
   • 这是 Head 0 的最终输出。

在 NPU1 上（Head 1 的计算）：

• 完全相同的流程：Scores_B = Q_B [2,2] @ K_B^T [2,2] = [2, 2]
• Attn_B = Softmax(Scores_B / √2) → [2, 2]
• Output_Head_B = Attn_B [2,2] @ V_B [2,2] = [2, 2]

物理意义：完整的 Multi-Head 输出（形状应为 [2, 4]）现在被物理切分为：

• NPU0 持有左半部分 [2, 2]（Head 0 的结果）。
• NPU1 持有右半部分 [2, 2]（Head 1 的结果）。

9.4 阶段三：输出投影 W_O（行并行 + All-Reduce Sum）

现在要将两个头的输出拼接起来（[Output_Head_A | Output_Head_B]，形状 [2, 4]）投影回 d_model 维度。W_O 按行切分（W_O_A 为上半行，W_O_B 为下半行）。

根据分块矩阵乘法：
Y = [Head_A | Head_B] @ [W_O_A ; W_O_B] = (Head_A @ W_O_A) + (Head_B @ W_O_B)

在 NPU0 上（计算前半部分贡献）：

• 持有 Output_Head_A [2, 2] 和 W_O_A [2, 4]。
• 计算局部累加项：Partial_0 = Output_Head_A @ W_O_A → 形状 [2, 2] × [2, 4] = [2, 4]。
• 这是最终输出 Y 的第一项。

在 NPU1 上（计算后半部分贡献）：

• 持有 Output_Head_B [2, 2] 和 W_O_B [2, 4]。
• 计算局部累加项：Partial_1 = Output_Head_B @ W_O_B → 形状 [2, 2] × [2, 4] = [2, 4]。
• 这是最终输出 Y 的第二项。

核心通信与统筹（All-Reduce Sum）：

• NPU0 和 NPU1 发起 All-Reduce（求和） 通信，由HCCL完成。
• 通信完成后：
  • NPU0 获得：Partial_0 + Partial_1 = Y（完整的最终输出，形状 [2, 4]）。
  • NPU1 获得：Partial_0 + Partial_1 = Y（同样获得完整的 [2, 4]）。

9.5 阶段四：反向传播（镜像对称）

反向传播是前向的严格逆过程，通信模式对调：

1. W_O 的梯度计算（无需通信）：既然两张卡都有完整的 dY [2,4]，它们各自独立计算 W_O_A 和 W_O_B 的梯度，无需通信。
2. dHead 的梯度回传（需要 All-Reduce）：为了给上一阶段（Q/K/V投影）传梯度，需要计算 dX_proj。根据链式法则，dX = dY @ W_O^T。因为 W_O 是行并行，反向计算拆分：
   • NPU0 算 dX_A = dY @ (W_O_A)^T
   • NPU1 算 dX_B = dY @ (W_O_B)^T
   • 发起 All-Reduce Sum：dX = dX_A + dX_B，两张卡均获得完整梯度，传给前面的层。

9.6 Attention vs. FFN 核心计算对比表

维度	FFN 层 (之前案例)	Attention 层 (本案例)
输入形状	[1, 4] (单Token)	[2, 4] (双Token序列)
第一步切分	w1/w3 列并行（升维）	W_Q/K/V 列并行（按注意力头切分）
中间非线性	SwiGLU (逐元素，无通信)	Softmax (沿序列维度，完全本地，无通信)
第二步切分	w2 行并行（降维）	W_O 行并行（输出投影）
前向通信点	仅 1 次 All-Reduce (在 w2 末尾)	仅 1 次 All-Reduce (在 W_O 末尾)
反向通信点	仅 1 次 All-Reduce (在 w1 梯度回传)	仅 1 次 All-Reduce (在 W_Q/K/V 梯度回传)
核心数学依赖	矩阵乘法分配律	矩阵乘法分配律 + Softmax 对序列维度的独立性

9.7 特殊情况说明

此处对应之前提到的特殊情况，简要说明：

1. Dropout 的位置：在标准 Megatron 实现中，W_O 输出后的 Dropout 层必须在 All-Reduce 之后执行（因为 Dropout 需要完整的全局张量来生成掩码），这会导致两张卡重复计算。高级优化（如序列并行）会将 LayerNorm/Dropout 的输入沿序列维度切分来规避。
2. KV Cache（推理场景）：在自回归推理时，K 和 V 矩阵会不断增长。TP 依然按头切分，每张卡只管理自己负责的头对应的 Cache，互不干扰，无需额外通信。
3. Grouped-Query Attention (GQA)：在 GQA 中，Q 的头数多于 K/V 的头数。TP 切分时通常确保每张卡分到整数个 Q 头，并复制对应的 K/V 头分片，计算逻辑不变，依然遵循“列并行投影 -> 本地算注意力 -> 行并行输出”的范式。

十、总结

通过上述案例可以看到，无论 FFN 还是 Attention，TP 的本质都是“先列切（免通信算中间结果） -> 再行切（All-Reduce 合并最终结果）”。Attention 只是把中间的矩阵乘法换成了 Softmax，而 Softmax 恰好不破坏这种分片独立性，因此整体范式完全通用。这一整套并行策略在Atlas 800I A2 NPU平台上可以无缝运行，HCCL通信库提供了高效的集合通信支持。