从一个 Decoder-only Transformer 的输入张量出发:
其中
对应六种主流并行策略
| 并行策略 | 核心切分维度 | 主要解决问题 | 典型通信 |
|---|---|---|---|
| 数据并行 DP | Batch | 提高吞吐、扩展训练规模 | AllReduce / ReduceScatter / AllGather |
| 张量并行 TP | Hidden / Head / FFN | 单层参数或计算过大 | AllReduce / ReduceScatter / AllGather |
| 流水线并行 PP | Layer | 模型层数太多、整模型放不下 | P2P Send / Recv |
| 上下文并行 CP | Sequence / Context | 长上下文 Activation 与 Attention 计算 | Ring P2P / AllGather / AllToAll |
| 专家并行 EP | Expert | MoE Expert 参数规模过大 | AllToAll |
| 序列并行 SP | Sequence | 降低 TP 域中重复 Activation | AllGather / ReduceScatter |
Hugging Face 的 Model Parallelism 把 TP 描述为对 Tensor 的“横向切分”,PP 描述为沿 Layer 的“纵向切分”;进一步组合 DP、TP、PP 后,就得到经典的 3D Parallelism。
并行训练的通信问题
统一的 - 通信模型
分布式通信耗时可以近似写成:
其中:
:每次通信的启动延迟; :消息或 Collective 的次数; :单位字节传输时间,可以近似理解为带宽倒数; :实际传输的数据量。
这意味着通信瓶颈存在两个完全不同的方向。
一类是 Latency Bound:Tensor 不大,但 Collective 调得非常频繁,例如高 TP Degree 下每层频繁执行 Collective;另一类是 Bandwidth Bound:通信次数不多,但一次要搬运大量参数、梯度或 Token,例如 DP Gradient Reduce、FSDP Parameter AllGather、MoE AllToAll。
以 Ring AllReduce 为例,对大小为
原因是 Ring AllReduce 可以拆成:
每一半的数据量约为:
Visible Communication
完全消灭通信通常不现实。更实际的目标是让通信与计算重叠:
如果通信能被 GEMM、Attention 或 Backward 覆盖,那么即使总字节数没有下降,端到端训练时间也可能明显下降。
因此很多工程优化的本质都是:
数据并行 DP
DP 的实现机制
Data Parallelism 的基本思想最简单:每张 GPU 保存相同的完整模型,数据拆成多份不同 GPU 处理不同的数据。因此 DP 主要解决吞吐问题,不能解决单卡模型显存的问题
如果模型参数为
只是每个 GPU 处理的数据不同,DP 切分的是 Batch Dimension 而不是模型本身。
在 Forward 阶段,几乎完全独立,假设一个 Global Batch 为
DP 之后
每张卡独立处理
在 Backward 阶段,每张 GPU 得到的 梯度不同。每张 GPU 独立计算:
为了保证每张 GPU 使用相同的全局梯度,计算
所有卡使用相同的 AllReduce。
下证 DP 在数学上等价于 大 Batch 单卡训练。
损失定义为 Batch 平均
梯度
现在用 D 张 GPU,每张 GPU 拿到
个样本,第 张 GPU 的局部梯度为
然后对所有的 GPU 的梯度求平均
等于单卡直接处理整个 Global Batch 得到的梯度,因此同步数据并行,本质上是在多张 GPU 上分布式计算一个更大的 Batch.
单卡处理 128 个样本与每张卡承担原来
但是实际上还需要梯度通信
因此加速不可能达到
的比例。
每次参数对应只有一个梯度,因此 DP 通信的数据量大致和参数量有关。而计算量和 Local Batch Size 有关,因此
DP 更新欢每张 GPU 都有足够大的 Local Batch,通信更容易被隐藏。如果每张卡计算很少,但是仍然需要同步整个模型的梯度,GPU 数量增加反而没有意义了。
最朴素的计算与通信流程是
graph LR A(Forward) --> B(整个 Backward) --> C(AllReduce 全部 Gradient) --> D(Optimizer)
这样会导致
1 | Layer4 backward |
理想情况下
根据统一的
- 通信模型,不能每算出一个参数梯度就执行一次 AllReduce,DP 会把多个 Gradient 放进一个 Bucket,满了之后执行 AllReduce(bucket).需要平衡的是
Bucket 太小 → 通信启动次数太多 → latency 高
Bucket 太大 → 必须等更多梯度 ready → overlap 机会下降
DP 的通信量
假设模型参数量为
如果直接使用 Ring AllReduce,
ZeRO / FSDP
传统 DDP 最大的问题是每个 Rank 都保存完整 Parameter、Gradient 与 Optimizer State。ZeRO/FSDP 在仍然保持数据并行语义的前提下,把训练状态按 DP Rank 分片。
1 | ZeRO-1: Optimizer State sharded |
ZeRO-3 / FSDP 的 Forward 典型流程是:
1 | Parameter shard |
Backward 则需要再次恢复参数,并用 ReduceScatter 聚合并重新分片梯度。因此 ZeRO/FSDP 的本质是:
Hugging Face 的 Model Parallelism 用了一个非常形象的比喻:每个参与者只背一部分装备,需要使用某层时临时共享并重建完整参数,用完后再释放。
DP 的优缺点
优点:模型计算图基本不需要修改;GPU 之间在 Forward/Backward 主体中相互独立;扩展到更多节点相对容易;通信可以通过 Gradient Bucket 与 Backward 高度重叠。
缺点:普通 DDP 不降低模型状态显存;DP Degree 过大时 Global Batch 容易变得过大;使用 FSDP/ZeRO-3 后虽然显存下降,但 Parameter AllGather 与 Gradient ReduceScatter 增加了通信压力。
张量并行 TP
数据并行 DP 是 模型复制、数据切分;张量并行 TP 是 数据共享、模型切分,也就是说把一个 Transformer Layer 内部的大矩阵拆到多张 GPU 上,让多张 GPU 共同完成同一个样本的计算。
考虑线性层:
其中:
Column Parallel Linear
按输出维切
每块:
于是每个 TP Rank 独立计算:
得到:
因为输出维本来就被切开,所以这一阶段不需要立刻做求和通信。
Row Parallel Linear
第二个 Linear 可以按输入维切:
其中:
如果输入本来已经被 TP 切成
完整结果是:
因此需要 AllReduce;启用 Sequence Parallel 后,这里通常进一步改写成 ReduceScatter,把“求和”和“Sequence 分片”合并到同一个 Collective 中。
以两层 MLP 为例:
让
1 | X |
中间的 GELU/SwiGLU 可以直接对各自 Shard 独立执行,因此不用在两个 Linear 中间重建完整 Tensor,Column Parallel 后不 Gather,而是直接接 Row Parallel,直到最后才 Reduce。
Hugging Face 的 Model Parallelism 也用相同的矩阵视角解释 TP:先按列切第一个权重矩阵,使不同 GPU 的输出可以独立经过 GeLU,再在需要恢复完整语义的位置进行同步。
Multi-Head Attention:
多个 Attention Head 本身相互独立。若总 Head 数为
个 Head。
QKV Projection 使用 Column Parallel,Output Projection 使用 Row Parallel:
1 | X |
因此 Transformer Block 中 TP 的 Collective 通常集中在 Attention Output 与 MLP Output 附近,而不是每一个 Linear 都 Gather 完整 Tensor。
TP 依赖高速互联
TP 的最大问题不是总通信量,而是通信发生在 每一层的 Critical Path 上:
1 | GEMM |
几十甚至上百层会不断重复这一模式。因此TP 需要非常快的网络,实践中通常尽量不要把一个 TP Group 跨越慢速节点互联。TP 应优先放进 NVLink/NVSwitch 域,并只扩到足够解决单层显存与计算问题的程度。
TP Degree 继续增大时,每个 Rank 上 GEMM 变小,Tensor Core 利用率下降,而 Collective Latency 并不会同比下降,最终会出现 Scaling Saturation。
讨论
评论