模型训练并行方案讲解：从基础到进阶

单张GPU显存和算力往往无法满足大模型需求。并行方案的核心是把”计算任务”或”模型/数据”拆分到多个设备，突破硬件瓶颈。

纯数据并行（DP）

核心定义

所有设备保存完整模型参数、梯度和优化器状态，每个设备处理不同数据分片；前向计算后汇总梯度，反向更新时保证参数同步。

技术原理

假设总批次$N$，有$k$个GPU，每个GPU数据量$n=N/k$：

• 前向：第$i$个GPU处理$X_i$，计算损失$L_i$
• 梯度计算：$g_i={\nabla }_{\theta }{L}_i$
• 梯度汇总：$g=\frac{1}{k}{\sum }_{i=1}^k{g}_i$
• 参数更新：$\theta =\theta -\eta \cdot g$

优缺点：

• 优点：实现简单、通用性强、通信开销小
• 缺点：显存瓶颈——模型参数量超单卡显存时无法使用
• 适用：中小模型（≤单卡显存）

完全分片数据并行（FSDP / ZeRO）

核心定义

将模型参数、梯度、优化器状态分片存储到不同GPU，每个GPU只存一部分；计算时按需聚合分片。

ZeRO分片粒度：

• ZeRO-1：仅分片优化器状态
• ZeRO-2：分片优化器状态+梯度
• ZeRO-3：分片优化器状态+梯度+模型参数

FSDP/ZeRO-3单步训练时间线

阶段	第$i$个GPU行为	所有GPU协同
前向	聚合所有${\theta }_j$得到完整$\theta$，计算$L_i$	同步发送各自参数分片
梯度计算	计算完整梯度$g$，拆分为$g_1$到$g_k$，仅保留$g_i$	各GPU独立完成
梯度汇总	交换$g_i$，计算全局梯度分片$G_i$	全量交换梯度分片
参数更新	仅用$G_i$更新${\theta }_i$	各GPU独立更新

优缺点：

• 优点：单卡显存仅为DP的$1/k$，能训练超大模型
• 缺点：通信开销略高
• 适用：大模型（参数量超单卡显存）

张量并行（TP）

核心定义

按模型的”张量维度”拆分单个层的权重——把线性层、注意力层的权重张量拆分成子张量，不同GPU负责计算该层的一部分。

比喻：做一道菜，其中”炒肉丝”步骤需要切、炒、调味三个子步骤，3个厨师各负责一个子步骤，做完后拼接结果。

技术原理

线性层$Y=X\cdot W+b$，将$W$按输出维度$V$拆分为$k$个子矩阵：

1. 第$i$个GPU持有$W_i$，接收完整输入$X$
2. 计算局部输出：$Y_i=X\cdot W_i+b_i$
3. 通信拼接：$Y=[Y_1,Y_2,...,Y_k]$

优缺点：

• 优点：解决单层层权重超单卡显存问题
• 缺点：实现复杂、通信开销高
• 适用：超大模型的核心层（如LLaMA-70B）

流水线并行（PP）

核心定义

按模型的”层维度”拆分——把连续层分配到不同GPU，数据按”流水线”方式依次经过各GPU。

比喻：工厂流水线，模型的每一层是一个工位，数据是产品；产品依次经过各工位处理。

技术原理

模型有$L$层，拆分为$k$个GPU，第$i$个GPU负责层$L_{(i-1)\cdot L/k+1}$到$L_{i\cdot L/k}$：

• 前向：数据从GPU 1流向GPU k
• 反向：梯度从GPU k流回GPU 1

优缺点：

• 优点：拆分逻辑直观
• 缺点：存在”流水线气泡”，部分GPU空闲等待
• 适用：超深模型，常与TP/FSDP组合（3D并行）

总结

方案	核心思想	显存占用	适用场景
DP	数据分片，模型完整	高	中小模型
FSDP/ZeRO	模型状态分片	低（$1/k$）	大模型
TP	张量维度拆分	中	超大模型核心层
PP	层维度拆分	中	超深模型

TP和PP有什么具体不同