模型训练的并行化策略详解

大语言模型(LLM)的训练本质上是一个不断试错并自我修正的循环过程. 无论是几十亿参数的 Qwen 还是千亿参数的超大模型, 底层的基本动作都逃不出这四个步骤的循环.

以下是这四个阶段的核心归纳:

模型训练的核心阶段回忆

1. 样本划分 (Data Partitioning)

在这个阶段, 海量的文本数据被准备好并切割成模型可以消化的"小块(Chunk)".

• Token 化 (Tokenization): 文本首先被切分成模型能看懂的词元(Token). 比如"我爱机器学习"可能被切分成 [我, 爱, 机器, 学习] 并转化为对应的数字 ID.
• 构建输入与标签 (Inputs & Labels): 大模型通常采用"下一个词预测(Next-Token Prediction)"任务. 如果输入序列是 [我, 爱, 机器], 那么它的真实标签(Ground Truth)就是 [学习].
• 批处理 (Batching): 模型不是看一句话训练一次, 而是同时看很多句话.
  • Global Batch Size: 全局每一轮看多少数据. 比如设为 1024, 意味着训练一次要看 1024 条长句子.
  • Micro-Batch Size: 在分布式训练中, 这 1024 条数据会被分发到各个 GPU/NPU 卡上, 单张卡一次处理的数据量就叫微批次(比如每张卡分到 4 条).

2. 损失计算 (Loss Calculation / 前向传播)

这是模型做"预测"并评估"自己有多笨"的阶段.

• 前向传播 (Forward Pass): 数据进入模型. 输入的 Token 经过多层 Transformer(包括自注意力层和全连接层)的复杂矩阵乘法, 最终在最后一层输出一个词表大小的概率分布.
• 计算概率: 模型会给词表里的每一个词打分. 比如预测下一个词是"学习"的概率为 20%, 是"苹果"的概率为 5%.
• 计算损失 (Loss): 使用交叉熵损失函数(Cross-Entropy Loss). 系统会将模型的预测概率与真实标签(100% 应该是"学习")进行对比. 预测越准, Loss 越小; 预测越离谱, Loss 越大. 最终得出一个标量数值, 代表模型当前表现的误差.

3. 梯度计算 (Gradient Calculation / 反向传播)

算出误差后, 模型需要知道"究竟是哪个参数出了问题, 该往哪个方向改". 这个过程依靠微积分中的链式法则(Chain Rule), 这里通常需要使用矩阵微积分.

• 反向传播 (Backward Pass): 计算过程从最后一层往回倒推, 回到第一层.
• 求导数 (Gradients): 系统会对模型中的每一个参数(权重 \(W\)) 计算它对最终 Loss 的偏导数. 这个偏导数就是梯度\((\nabla W)\).
• 梯度的物理意义: 梯度就像是指南针. 它告诉我们, 如果想让当前的误差(Loss)变小, 这个特定的参数 \(W\) 应该增大一点点还是减小一点点.

4. 参数更新 (Parameter Update)

这是模型真正"变聪明"的时刻, 由优化器(Optimizer, 如 AdamW)主导.

• 应用规则: 优化器拿着刚刚算出的所有梯度, 决定最终要怎么修改参数.
  • 如果你用的是最基础的 SGD, 公式大概是: 新参数 = 老参数 - (学习率 × 梯度).
  • 在大模型中通常使用 AdamW 优化器(Adam with Decoupled Weight Decay), 它还会参考历史梯度(一阶动量和二阶方差), 防止参数更新过猛或陷入死胡同.
• 更新权重: 优化器在内存中把新的数值写入模型参数矩阵中. 至此, 模型完成了一次"学习".
• 清空梯度 (Zero Grad): 更新完后, 必须把所有的梯度清零, 干干净净地迎接下一个批次的数据.

模型训练中GPU显存具体存储的内容

在基础训练模式下, 显存并非存储所有数据的"一部分", 除了训练数据外, 它必须完整吞下所有模型训练过程中的状态:

• 100% 完整模型参数: 整个神经网络的所有权重矩阵(FP16/BF16 格式)必须全部加载在显存中.
• 100% 完整梯度: 模型有多少个参数, 显存里就必须有对应数量的梯度(FP16/BF16 格式).
• 100% 完整优化器状态: 如果使用 AdamW 优化器, 还需要包含模型参数的 FP32 备份(Master Weights)、一阶动量(FP32)和二阶方差(FP32).
• 当前的微批次(Micro-Batch): 海量的完整数据集存放在 CPU 内存或硬盘(SSD)中. 只有当前这一步需要计算的那一小批数据(例如 4 条句子)会被拷贝进显存.

显存的分布与占比(混合精度训练)

基于上述内容, 我们可以算出一笔清晰的显存账. 假设我们使用主流的 FP16/BF16 混合精度, 并配合 AdamW 优化器训练一个参数量为 \(\Phi\) 的模型.

显存内容	占用大小	典型占比(大致)	说明
优化器状态	12\(\Phi\) 字节	~60%	显存最大的消耗者, 用于保存 FP32 备份及动量参数.
激活值 (Activations)	动态变化	~20% - 30%	前向传播的中间结果(如 Attention 的 Q/K/V 矩阵).
模型参数	2\(\Phi\) 字节	~5% - 10%	网络本体.
模型梯度	2\(\Phi\) 字节	~5% - 10%	反向传播算出的误差方向.
当前 Batch 数据 / 显存碎片	极小	< 5%	当前处理的文本输入以及框架内部缓存.

数据并行(Data Parallelism)

数据并行的核心思想多个batch的样本同时训练, "切分数据, 不切分模型". 每个GPU上都复制一份完整的模型, 但各自处理不同的数据批次(Batch). 计算出梯度后, 所有GPU会进行通信, 同步梯度并更新参数.

数据并行是如何计算损失的？

数据并行的损失是各自独立计算的, 并没有一台专门的中心机器来统揽全局. 假设全局批次大小(Global Batch Size)是 1024, 有 4 台机器(GPU):

1. 数据加载器会将数据等分为 4 份, 每台机器分到 256 条数据(称为 Micro-batch).
2. 每台机器使用自己内存中的完整模型, 对自己分到的这 256 条数据进行前向传播(Forward Pass).
3. 每台机器基于这 256 条数据的预测结果和真实标签, 独立计算出属于自己的局部损失(Local Loss).

如何保证四台机器上的模型完全一致

这是数据并行的核心机制, 依赖于反向传播后的梯度同步(Gradient Synchronization). 具体步骤如下:

1. 局部反向传播: 每台机器根据刚才算出的局部 Loss, 独立进行反向传播, 计算出自己这份数据的梯度(局部梯度).
2. All-Reduce 聚合(关键): 在更新模型权重之前, 4 台机器会通过网络执行一个叫做 All-Reduce 的通信操作. 这个操作会将 4 台机器上的局部梯度加总并求平均, 然后把这个全局平均梯度广播回每一台机器.
3. 同步更新: 此时, 4 台机器上虽然最初处理的数据不同, 但它们最终拿到的梯度是完全一致的. 因为它们在训练之初的参数(\(W\))是相同的, 现在用来更新的梯度\(\nabla W\)也是相同的, 经过优化器(如 Adam)计算后, 更新后的模型权重必然保持绝对一致.

数据并行带来的显存增长

在传统的数据并行中, 如果你有 64 张卡, 模型的状态(包含模型参数、梯度、Adam 优化器状态)会被原封不动地复制 64 份. 对于大模型, 单卡根本装不下一份完整的状态. 尤其是 Adam 优化器需要保存一阶动量和二阶方差, 其占用的显存是模型参数本身的数倍. 传统 DP 极度浪费显存, 遇到大模型会立刻 OOM(内存溢出).

.svg)

ZeRO 根据"切碎"的激进程度, 分为三个阶段(Stage):

• ZeRO-1(切分优化器状态): 模型参数和梯度每张卡都有全量, 但最庞大的优化器状态被切分成 \(N\) 块, 每张卡只维护 \(\frac{1}{N}\). 需要更新权重时, 每张卡只负责更新属于自己的那一小块参数, 更新完后再互相广播同步(All-Gather).
• ZeRO-2(切分优化器状态 + 梯度): 在 Stage 1 基础上, 梯度也被切分. 算完反向传播并执行 Reduce-Scatter 后, 每张卡只保留自己负责那一小块的梯度, 其余部分的梯度直接从显存中丢弃.
• ZeRO-3(切分优化器状态 + 梯度 + 参数): 终极形态. 每张卡平时只保留 \(\frac{1}{N}\) 的参数. 实现方式是动态借用: 当卡 A 计算前向或反向传播到某一层网络, 发现自己缺少这层的参数时, 它会通过网络向其他卡临时索要(All-Gather). 算完这一层后, 立刻把借来的参数丢弃(释放显存), 再去索要下一层的参数.

张量并行(Tensor Parallelism)

张量并行的本质是"把一个巨大的矩阵乘法, 强行掰成几块, 让多张卡一起算, 算完再拼起来". 在大模型训练中, 就是每张卡处理相同的输入 \(X\), 不同的卡计算不同的矩阵部分, 计算完之后再将模型拼起来. 而它切分模型的维度为矩阵维度. 它由 NVIDIA 的 Megatron-LM 框架发扬光大.

简单的场景下, 我们假设对于输入 \(X\), 经过参数矩阵 \(A\), 得到输出为 \(Y\), 使用表达式就是 \(Y = X \times A\). 当 \(A\) 矩阵极其庞大时, 我们可以利用线性代数的分块矩阵乘法法则将其切开.

对于输入的矩阵乘法, 假设有两张卡(GPU 0 和 GPU 1), 我们有两种切法:

按列切分 (Column Parallelism)

把参数矩阵 A 垂直切成两半: \(A=\begin{bmatrix}A_1,A_2\end{bmatrix}\). \(Y = X \times A =X \times \begin{bmatrix}A_1,A_2\end{bmatrix}\)

• GPU 0 计算: \(Y_1=X×A_1\)
• GPU 1 计算: \(Y_2=X×A_2\)
• TP并行策略: 两张卡分别算出了输出矩阵的左半边和右半边. 最后只要把这两个结果水平拼起来(All-Gather), 就得到了完整的 \(Y=[Y_1,Y_2]\).

按行切分 (Row Parallelism)

把参数矩阵 A 水平切成两半: \(A=\begin{bmatrix}A_1 \\ A_2\end{bmatrix}\). 为了能乘起来, 输入 \(X\) 也必须在特征维度上切成两半: \(X=[X_1,X_2]\).

• GPU 0 计算: \(Y_1=X_1×A_1\)
• GPU 1 计算: \(Y_2=X_2×A_2\)
• TP并行策略: 注意这里的数学关系, 根据矩阵乘法规则, \(Y=Y_1+Y_2\). 所以两张卡算完后, 必须把各自的结果加起来(All-Reduce), 才能得到最终的 \(Y\).

\[Y = X \times A = \begin{bmatrix}X_1, X_2\end{bmatrix} \times \begin{bmatrix} A_1 \\ A_2 \end{bmatrix} = (X_1 \times A_1) + (X_2 \times A_2) = Y_1 + Y_2 \]

Megatron 绝妙的设计

上述的两种切法与神经网络中的全连接层(\(MLP\)) 有着十分精妙的配合, 因为 \(MLP\) 通常是两个线性层的串联. Megatron-LM 天才般地将这两种切法组合了起来, 把通信开销降到了最低, 我们以基础的两层全连接神经网络为例, 对于输入 \(X\), 输出 \(Z = X \times MLP\):

1. 第一层(按列切): \(Y_1,Y_2=X\times [A_1,A_2]\). 算完之后, 坚决不通信拼接. 让 \(GPU_0\) 继续拿着 \(Y_1\), \(GPU_1\) 拿着 \(Y_2\).
2. 第二层(按行切): 刚好, 第二层需要输入被切分, 而我们手里刚好拿着切分好的 \(Y_1\)和 \(Y_2\). 于是直接计算 \(Z_1=Y_1×B_1\) 和 \(Z_2=Y_2×B_2\).
3. 最后汇总: 只有在两层全都算完之后, 才执行一次 All-Reduce(求和 \(Z=Z_1+Z_2\)).

\[Z = X \times MLP \\ MLP = \begin{bmatrix}A_1, A_2\end{bmatrix} \times \begin{bmatrix}B_1 \\ B_2\end{bmatrix} \\ Z = X \times MLP = X \times \begin{bmatrix}A_1, A_2\end{bmatrix} \times \begin{bmatrix}B_1 \\ B_2\end{bmatrix} = \begin{bmatrix}Y_1, Y_2\end{bmatrix} \times \begin{bmatrix}B_1 \\ B_2\end{bmatrix} = \begin{bmatrix} Y_1 \times B_1 + Y_2 \times B_2 \end{bmatrix} = \begin{bmatrix} Z_1 + Z_2 \end{bmatrix} \]

结论: 在一个由两个线性层组成的庞大 MLP 模块里, TP 只需要在最后做一次通信 这极大降低了网络压力.

Tranformer 中 Docoder 层的 TP 实例

在大模型训练与推理中, 几乎都会使用到 TP并行, 因为大部分大模型都是基于 Transformer 架构的, 我们使用下面的实例, 假设输入 \(X\in\mathbb{R}^{B\times S\times D}\) 已在4卡上同步.

多头 Attention 层

QKV线性(列切分)

每卡持有 \(W_Q\), \(W_K\), \(W_V\) 的列片段

\[Q_i = X \times W_{Q_i} \\ K_i = X \times W_{K_i} \\ V_i = X \times W_{V_i} \]

局部Attention计算, 每卡独立完成自己的头, 这里不涉及行列切分, 也不涉及对输入 \(X\) 的直接计算, 因此 \(Z_i\) 依然可以不进行 All-Reduce 的操作.

\[A_i = softmax(\frac {Q_i K_i^T} {sqrt(d)}) \\ Z_i = A_i \times V_i \]

输出投影(行切分)

每卡持有 \(W_O\) 的行片段

\[O_{partial} = Z_i \times W_{O_i} \]

所有卡求和, 得到完整(B,S,D), 然后计算残差与正则化

\[O_{full} = AllReduce\_sum(O_{partial}) \\ H = X + O_{full} \\ X_2 = LayerNorm(H) \]

前馈网络子层

MLP第一层(列切分)

\[U_{partial_i} = X_2 \times W_{1_i} \\ U = GELU(U_{partial_i}) \]

#MLP第二层(行切分)

\[V_{partial_i} = U \times W_{2_i} \]

所有卡求和, 得到完整(B,S,D), 然后计算残差与正则化, 至此, 得到这一层 Decoder 一层的输出

\[V_{full} = AllReduce\_sum(V_{partial}) \\ H_2 = X_2 + V_{full} \\ X = LayerNorm(H_2) \\ \]

反向传播时的绝妙巧合

反向传播是从后往前的. 假设我们已经拿到了来自更后方网络的梯度 \(\frac{\partial L}{\partial Z}\)(由于前向做过了 AllReduce, 此时每张卡上的 \(\frac{\partial L}{\partial Z}\) 是一模一样的).

第一步: 穿过第二层(原行切分层)

根据我们在上一轮推导的偏导公式, 对输入的梯度为:

\[\frac{\partial L}{\partial Y} = \frac{\partial L}{\partial Z}W_2^T \]

• 由于 \(W_2\) 前向是行切分的, 它的转置矩阵 \(W_2^T\) 自然就变成了列切分.
• \(GPU_0\) 拥有 \(W_{20}\), 反向传播时计算: \(\frac{\partial L}{\partial{Y_0}}=\frac{\partial L}{\partial Z}W_{20}^T\)
• \(GPU_2\) 拥有 \(W_{21}\), 反向传播时计算: \(\frac{\partial L}{\partial{Y_1}}=\frac{\partial L}{\partial Z}W_{21}^T\)
• \(GPU_0\) 拥有 \(W_{22}\), 反向传播时计算: \(\frac{\partial L}{\partial{Y_2}}=\frac{\partial L}{\partial Z}W_{22}^T\)
• \(GPU_2\) 拥有 \(W_{23}\), 反向传播时计算: \(\frac{\partial L}{\partial{Y_3}}=\frac{\partial L}{\partial Z}W_{23}^T\)
• 奇妙的事情发生了: 这里的计算过程完全独立, 梯度输出天然就是列切分的格式 \(\begin{bmatrix}\frac{\partial L}{\partial{Y_0}},\frac{\partial L}{\partial{Y_1}},\frac{\partial L}{\partial{Y_2}},\frac{\partial L}{\partial{Y_3}} \end{bmatrix}\). 这跟前向的"列切分"逻辑一模一样, 没有任何通信开销.

第二步: 穿过第一层(原列切分层)

继续往前传, 我们需要求对最开始输入 X 的梯度:

\[\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y}W_1^T \]

• 由于 \(W_1\) 前向是列切分的, 它的转置矩阵 \(W_1^T\) 自然就变成了行切分.

• 同时, 上一步传下来的梯度 \(\frac{\partial L}{\partial Y}\) 刚好是列切分的(各自独立持有 \(\begin{bmatrix} \frac{\partial L}{\partial Y_0}, \frac{\partial L}{\partial Y_1},\frac{\partial L}{\partial Y_2},\frac{\partial L}{\partial Y_3} \end{bmatrix}\)).

• 要想求出完整的 \(\frac{\partial L}{\partial X}\) , 每个GPU计算出自己的那部分, 然后两者必须加起来:

  \[\frac{\partial L}{\partial X} =\frac{\partial L}{\partial Y}W_1^T= \begin{bmatrix} \frac{\partial L}{\partial Y_0}, \frac{\partial L}{\partial Y_1},\frac{\partial L}{\partial Y_2},\frac{\partial L}{\partial Y_3} \end{bmatrix} \times \begin{bmatrix} W_{10}^T \\ W_{11}^T \\ W_{12}^T \\ W_{13}^T \end{bmatrix} = \frac{\partial L}{\partial Y_0}W_{10}^T+ \frac{\partial L}{\partial Y_1}W_{11}^T+\frac{\partial L}{\partial Y_2}W_{12}^T+\frac{\partial L}{\partial Y_3}W_{13}^T \]

• 结论: 这里必须执行一次 AllReduce 通信. 这跟前向的"行切分"逻辑一模一样.

这就是为什么在阅读大模型分布式框架源码(如 Megatron-LM)时, 我们会看到极其对称的代码结构:

网络阶段	网络层设计	切分方向 (原矩阵)	计算过程 (转置后)	通信操作
前向传播	第一层 (Col Parallel)	按列切分 \(W_1\)	矩阵乘法	无需通信
	第二层 (Row Parallel)	按行切分 \(W_2\)	向量内积形式	AllReduce (求和输出)
反向传播	第二层 (原 Row Parallel)	按列切分 \(W_2^T\)	矩阵乘法	无需通信
	第一层 (原 Col Parallel)	按行切分 \(W_1^T\)	向量内积形式	AllReduce (求和梯度)

流水线并行策略(Pipeline Parallelism)

TP 并行是比 DP 并行更加细粒度的并行策略, TP 并行是在对大模型每一层 Transformer 的内部计算进行并行计算, 而 PP 并行则是在不同的 transformer 层之间进行并行计算.

PP 并行基础原理: 宏观切分模型

**TP并行是在层内切分, 而流水线并行(PP)的核心思想是"按层切分, 纵向接力". **

• 物理切分: 将一个极深的模型(如 64 层的 Qwen3-32B)沿着网络层横向切断. 如果 PP=4, 模型就被均分为 4 块, 分别交给 4 台机器(或 4 个计算节点/Stage)独立负责.
• 显存独占: 每个 Stage 的 GPU 只加载并保存属于自己的那部分模型参数. 这彻底打破了单卡显存无法容纳大模型的瓶颈.
• 接力通信: 前向传播时, Stage 0 算完前 16 层, 将"边界激活值(中间输出)"通过网络传给 Stage 1; 反向传播时, Stage 1 算完, 将"边界梯度"传回给 Stage 0. 这种通信数据量极小, 非常适合跨服务器(Node)互联.

PP 的调度魔法: 解决"发呆挂起"问题

如果只切分模型, 一次输入一整个大批次(Batch), 会导致严重的"流水线空泡(Bubble)"——一台机器满载时, 其余机器都在发呆.

为了提高利用率, 必须引入 微批次(Micro-batching): 将一个大 Batch 切碎, 源源不断地送入流水线, 让多台机器的计算重叠起来. 围绕微批次, 业界演进出了两种经典的调度算法:

1. 标准 1F1B (One-Forward-One-Backward)

这是最经典、也是目前最常用的非交错调度策略(Megatron-LM 默认方式).

• 运行机制: 在流水线"火力全开"的稳态下, 每个 Stage 的 GPU 严格遵守"算一个微批次的前向 (1F), 接着立刻算这个微批次的反向 (1B)"的节奏交替工作.
• 显存优势: 它的核心目的是尽早释放显存. 只要一个微批次算完了反向, 它在前向传播时积累的庞大"激活值(草稿)"就可以立刻丢弃.
• 缺陷: 依然存在不可忽视的 Pipeline Bubble. 在流水线刚启动(预热)和快结束(排空)时, 大量 GPU 处于等待状态, 造成算力浪费.

2. 交错式 1F1B (Interleaved 1F1B)

为了进一步压缩气泡, Megatron 后期提出了更加极限的交错调度策略.

• 切分更碎: 在标准 PP 中, 一个 GPU 独占连续的 16 层. 而在 Interleaved 1F1B 中, 框架将每张卡负责的层"打散". 比如让 GPU 0 负责第 1-4 层, 同时还要负责第 33-36 层(形成虚拟阶段 Virtual Stages).
• 运行机制: 微批次在流水线里变成了"折返跑". 数据在 GPU 0, 1, 2, 3 之间传一圈后, 还会回到 GPU 0 算后半段.
• 优势: 通过这种多段交错, 能大幅缩短流水线预热和排空的时间, 显著减小 Bubble 的比例, 提升算力利用率(MFU).
• 代价: 通信次数直接翻倍(数据要在节点间多传好几轮), 并且代码实现逻辑极其复杂, 对网络通信带宽要求更高.

PP 并行如何计算损失

流水线并行的运行看似是各个微批次在"各自为战", 但它们在数学上必须坚守 梯度累加(Gradient Accumulation) 原则, 以保证模型收敛.

损失(Loss)与局部梯度的计算

• 计算单位: 以 Micro-batch 为基本计算单位.
• 过程: 当最后一个 Stage(如 GPU 3)跑完一个微批次的前向传播后, 它会利用这部分数据的输出计算局部 Loss, 并立刻开启反向传播, 将误差一层层向前推(1B).

优化模型(参数更新)与全局同步

• 只存不更: 当任意一个 Stage 的 GPU 算出某个微批次的梯度时, 绝对不能立刻更新模型参数(否则会导致同一批次的前后数据遭遇不同版本的权重, 即权重撕裂). 它只能将这个梯度暂存在显存里.
• 累积叠加: 随着后续微批次的反向传播不断完成, GPU 会将它们产生的梯度在显存中不断相加求和(\(\nabla W = \nabla W_1 + \nabla W_2 + \dots\)).
• 最终裁决 (Optimizer Step): 只有当这个大 Batch 下的所有微批次都在整个流水线中跑完了反向传播(流水线完全排空), 每个 GPU 手里才凑齐了真正的"全局总梯度".
• 跨组通信与更新: 此时, 如果结合了数据并行(DP/ZeRO), 框架会在跨组的相同 Stage 之间执行通信(All-Reduce), 平均各自的总梯度. 最后, 所有的 GPU 统一在这一瞬间调用优化器更新参数, 干净利落地进入下一轮训练循环.

3D并行: DP + TP + PP 的巅峰协同

在工业界, 训练千亿甚至万亿参数的超大模型(如 GPT-3, 千亿 Qwen 等)往往无法仅靠单一的并行策略. 单卡装不下(需要切分参数), 通信开销大(需要优化拓扑), 这时就需要将数据并行(DP)、张量并行(TP)和流水线并行(PP)结合起来, 形成三维并行(3D Parallelism).

3D 并行的场景与拓扑设计

在 3D 并行中, 集群的 GPU 会被组织成一个三维网格. 假设总 GPU 数量为 \(N\), 那么 \(N = DP \times PP \times TP\). 不同的并行策略对网络通信的带宽要求大不相同, 因此其在物理硬件上的映射极度讲究:

1. TP (张量并行) 放节点内: TP 的前向和反向每算一层都需要 All-Reduce 或者 All-Gather 通信, 数据量大且频繁. 因此, TP 通常被限制在单个物理服务器(Node)内部(例如一台 8 卡 A100/H800 服务器), 利用超高速的 NVLink / NVSwitch(或华为的 HCCS)进行通信. 如果跨节点做 TP, 网络延迟会瞬间拖垮整个训练.
2. PP (流水线并行) 放跨节点: PP 只需要在 Stage 的边界传递激活值(Activations)或者梯度, 通信量相对较小(点对点 Send/Recv). 因此, PP 最适合放在节点之间, 通过 InfiniBand / 高速以太网进行互联.
3. DP (数据并行) 放跨节点组合: 剩下的维度留给 DP. 不管怎么切分, 模型的每个 "切片组合" 都会在全局中有多份副本. DP 组的通信仅在反向传播结束后的 Optimizer Step 时发生(一次全局 All-Reduce), 相比 TP 频率极低, 也可以从容应对跨节点的带宽. (工业界通常会在 DP 维度上开启 ZeRO-1 甚至 ZeRO-2 进一步节省显存).

3D 并行协同运行的数值案例

场景假设: 我们有 64 张 GPU(分布在 8 台八卡服务器上), 训练一个 1000 亿参数的大模型.
策略配置: 设定 \(TP = 8\), \(PP = 4\), \(DP = 2\). (\(8 \times 4 \times 2 = 64\) 卡). 全局 Batch Size 为 1024, Micro-Batch Size 为 4.

它的运行过程会像精密的齿轮一样咬合:

1. 数据切分 (DP维度):
   因为 \(DP=2\), 整个 64 卡的集群相当于分为两套"完全独立的完整模型流水线"(每套 32 卡). 全局的 Batch (1024条数据) 会被均分为两份, 每一条流水线负责处理其中的 512 条数据. 这 512 条数据会被切分成一个个大小为 4 的 Micro-batch, 源源不断地送入各自的流水线.
2. 层间接力 (PP维度):
   看其中一条流水线(32 张卡). 按照 \(PP=4\), 整个大模型(假设 64 层)被切分成 4 个大块(Stage 0 到 Stage 3), 每个 Stage 分到 16 层大模型, 并占用 8 张卡. 当数据进入 Stage 0 时, 它算完前 16 层的计算, 会通过跨节点网络(如 InfiniBand)把边界处的激活值打包 Send 给远端的 Stage 1 节点.
3. 层内爆算 (TP维度):
   现在让我们用放大镜看向某个具体的 Stage (比如 Stage 0 中的那 8 张卡). 这 8 张卡正好处于同一台服务器内. 它们在计算第一层(比如 Attention 计算的矩阵乘法)时, 参数 \(W_Q, W_K, W_V\) 已经被按列/按行切成了 8 份.
   • 8 张卡同时执行本地的算子计算(如 cublasLt 的分片矩阵乘法).
   • 算完后, 瞬间通过本地巨大的 NVLink 带宽执行 All-Reduce 或 All-Gather, 组合拼凑结果.
   • 继续往下进入当前 Stage 的下一层直至完成16层.
4. 全局参数更新 (统一时刻):
   当一套流水线把那 512 条数据的所有 Micro-batch 都跑完了反向传播, 并在各自显存里暂存并累加好了梯度. 此时, PP 和 TP 的接力任务暂时结束, 开始进行 DP 维度的最后收尾.
   • 跨越两套流水线的"相同位置的卡"(比如流水线 1 的 \(GPU_0\) 和流水线 2 的 \(GPU_0\))会执行一次全局维度的 All-Reduce.
   • 此时这两套流水线求出了当前 step 的 1024 条数据的平均总梯度.
   • 最终, 所有卡同时调用 AdamW 刷新本地保存的那一小部分参数切片. 内存清理, 干净利落地进入下一轮循环.

业界落地实现 (NVIDIA vs 华为 Ascend)

• NVIDIA 生态 (Megatron-LM + DeepSpeed):
  这是当前训练超大模型最主流的开源打法. Megatron-LM 贡献了极致的 TP 和 PP 切分实现, 以及高度融合的底层 CUDA Kernel(如 Fused LayerNorm, Custom Softmax 等); 而 DeepSpeed 则补充了强大的 ZeRO 系列优化机制和高效的流水线引擎(Pipeline Engine). 两方结合即是我们常说的 "Megatron-DeepSpeed" 方案.
• 华为昇腾 (MindSpore + Ascend):
  依托于 Ascend 芯片和 HCCL(华为集合通信库), MindSpore 框架原生提供了强大的自动并行(Auto-Parallel)能力. 开发者往往无需手动硬编码通信逻辑, 框架编译器会自动进行算子级的分布式切分图生成. 在 NPU 执行时, 会自动做大量算子层面的融合(例如将 Transformer 中的大量碎片化运算组成单一大算子)去提升计算密度, 并掩盖 HCCL 的通信延迟.