跳到正文
← 返回文章目录
技术与趋势

分布式训练并行策略详解:DP/TP/PP/SP 如何选择

全面解析四种并行策略的适用场景与组合方案

分布式训练并行策略详解:DP/TP/PP/SP 如何选择

训练大模型,单机远远不够。本文全面解析四种并行策略,帮你选择最优方案。


一、为什么需要并行?

单机限制

显存限制:

  • 单卡显存:40-80GB(主流)
  • 70B 模型 FP16:140GB+
  • 结论:单卡放不下

计算限制:

  • 训练时间过长
  • 无法在合理时间完成

并行目标

  1. 突破显存限制 - 训练更大模型
  2. 加速训练 - 缩短训练时间
  3. 提高资源利用率 - 多卡协同

二、四种并行策略

1. 数据并行(DP)

原理:

  • 模型完整复制多份
  • 数据分片到各卡
  • 梯度同步更新

优点:

  • 实现简单
  • 通信量小(只传梯度)
  • 适合中小模型

缺点:

  • 每卡需放下完整模型
  • 显存效率低

适用:

  • 模型 < 单卡显存
  • 卡数较少(< 8)
# PyTorch DDP 示例
model = torch.nn.parallel.DistributedDataParallel(model)

2. 张量并行(TP)

原理:

  • 单层内切分矩阵
  • 每卡计算部分
  • 结果合并

优点:

  • 突破单卡显存限制
  • 适合超大层

缺点:

  • 通信量大(每层都通信)
  • 实现复杂
  • 需要高速互联

适用:

  • 超大模型(70B+)
  • 卡间带宽高(NVLink)

切分方式:

  • 行切分(Row Parallel)
  • 列切分(Column Parallel)

3. 流水线并行(PP)

原理:

  • 按层切分模型
  • 每卡负责部分层
  • 数据像流水线传递

优点:

  • 显存效率高
  • 适合超深网络
  • 可跨机器

缺点:

  • Pipeline Bubble(气泡)
  • 通信频繁
  • 实现复杂

适用:

  • 超深网络(100+ 层)
  • 多机训练

优化:

  • 1F1B 调度(减少 bubble)
  • 微批次(Micro-batch)

4. 序列并行(SP)

原理:

  • 沿序列维度切分
  • 每卡处理部分序列
  • 针对 Attention 优化

优点:

  • 降低显存(Attention 是 O(n²))
  • 适合长序列

缺点:

  • 只适用于特定层
  • 实现复杂

适用:

  • 长序列场景
  • 配合 TP 使用

三、混合并行策略

为什么需要混合?

单一并行的局限:

  • DP:受限于单卡显存
  • TP:通信开销大
  • PP:Bubble 问题
  • SP:适用范围有限

混合方案:

  • 组合优势
  • 灵活适配
  • 最优效率

常见组合

2D 并行(DP+TP):

  • 适用:8-64 卡
  • 配置:DP=4, TP=2
  • 场景:中等模型

3D 并行(DP+TP+PP):

  • 适用:64-512 卡
  • 配置:DP=8, TP=4, PP=2
  • 场景:大模型

4D 并行(DP+TP+PP+SP):

  • 适用:512+ 卡
  • 场景:超大模型 + 长序列

配置示例

训练 70B 模型(64 卡):

总卡数:64
DP=8, TP=4, PP=2

每卡显存需求:
- 模型参数:70B / (8×4×2) = 1.1B
- 显存占用:约 10GB
- 加上激活、梯度:约 30GB
- A100 80GB:✅ 足够

四、框架实现

DeepSpeed

支持: ZeRO + 3D 并行

配置:

{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 4,
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu"
    }
  },
  "tensor_parallel": {
    "tp_size": 4
  },
  "pipeline": {
    "enabled": true,
    "num_stages": 2
  }
}

Megatron-LM

支持: TP+PP+DP

启动:

python train.py \
  --tensor-model-parallel-size 4 \
  --pipeline-model-parallel-size 2 \
  --num-layers 24 \
  --hidden-size 1024

PyTorch FSDP

支持: 类似 ZeRO-3

代码:

from torch.distributed.fsdp import FullyShardedDataParallel

model = FSDP(model, sharding_strategy=FULL_SHARD)

五、性能优化

1. 通信优化

梯度压缩:

  • 减少通信量
  • 精度损失可控

通信重叠:

  • 计算与通信并行
  • 隐藏通信延迟

拓扑感知:

  • 同机优先 TP
  • 跨机优先 DP/PP

2. 显存优化

激活重计算:

  • 牺牲计算换显存
  • 显存减少 60-70%

ZeRO 优化:

  • ZeRO-1:优化器分片
  • ZeRO-2:梯度分片
  • ZeRO-3:参数分片

3. 负载均衡

问题: 层大小不均导致负载不均衡

解决:

  • 均衡切分
  • 动态调度
  • 虚拟流水线

六、选型指南

决策树

模型能放下单卡?
├─ 是 → DP 即可
└─ 否 → 
    需要跨机器?
    ├─ 是 → DP+PP(主)+TP(辅)
    └─ 否 → 
        卡间带宽高?
        ├─ 是 → DP+TP
        └─ 否 → DP+PP

推荐配置

模型规模 推荐方案 卡数
7B 以下 DP 1-8
7B-70B DP+TP 8-64
70B+ DP+TP+PP 64-512
100B+ + 长序列 DP+TP+PP+SP 512+

避坑指南

❌ 盲目增加 TP - 通信开销可能超过收益 ✅ 先测试再扩展 - 从小规模验证

❌ 忽视网络拓扑 - 跨机 TP 效率低 ✅ 同机 TP,跨机 DP/PP

❌ 过度并行 - 管理复杂度高 ✅ 够用就好,逐步优化


总结

并行策略选择是系统工程:

  1. 理解原理 - 四种并行的特点
  2. 评估需求 - 模型规模、硬件条件
  3. 选择方案 - 匹配最优组合
  4. 持续优化 - 根据实际效果调整

核心原则: 简单优先,逐步复杂化。


有问题欢迎交流讨论!

继续阅读

公众号 · 技术与趋势

智能简史:AI也许在倒着进化

孟岩的播客推荐了这本书,他评价:启发很大但受关注度很小。是的,豆瓣9.0分,但打分人数只有800多人。 这本书从神经演化的角度谈智能进化。首先作者有意思,班内特不是学院派,而是一线AI创业者,他比神经科学家更理解智能演化与AI的关系。其次视