前言
曾在游戏世界挥洒创意,也曾在前端和后端的浪潮间穿梭,如今,而立的我仰望AI的璀璨星空,心潮澎湃,步履不停!愿你我皆乘风破浪,逐梦星辰!
好问题!你想深入理解 NVIDIA 提出的三种并行策略,以及在工程实践中怎么选择——这就是从「理解理论」到「工程落地」最关键的一步了。
✅ 一、NVIDIA 提出的三种并行策略
这三种并行策略主要用于训练超大规模模型,尤其在 Megatron-LM 和后续如 GPT-3 的训练中:
1️⃣ 数据并行(Data Parallelism)
👉 原理:
每张卡都拥有完整的模型副本,但处理不同的数据子集。
每轮训练后,所有卡计算自己的梯度,然后用 All-Reduce 操作平均梯度,同步更新参数。
✅ 优点:
-
实现简单,PyTorch DDP/TF MirroredStrategy 都支持;
-
适用于模型可以放进单卡显存的场景;
-
扩展性强,适合水平扩展训练。
❌ 缺点:
-
受限于单卡显存,模型参数不能太大;
-
每轮都要 All-Reduce,通信量跟模型大小成正比。
2️⃣ 张量并行(Tensor Parallelism / 纵向切分)
👉 原理:
将模型中的某些大张量(如线性层的权重矩阵)按列或行方向切分,每张卡只负责计算一部分权重对应的输出。
Megatron-LM 中举的例子:
-
把 transformer 中的 QKV weight
[hidden, hidden]
拆成[hidden, hidden / n]
分给 n 张卡; -
每张卡输出一部分结果,在最后一个维度拼接起来。
✅ 优点:
-
能训练远超单卡显存的模型,比如百亿、千亿参数;
-
层内部并行,能细粒度拆分大参数层。
❌ 缺点:
-
通信频繁(forward 和 backward 时都要通信),性能依赖高带宽网络;
-
编程和调试比数据并行复杂不少;
-
只适用于结构规则的网络(如 Transformer),不适用于 CNN 等。
3️⃣ 流水线并行(Pipeline Parallelism / 横向切分)
👉 原理:
将模型按「层」划分成若干段,分别部署在不同设备上,像流水线一样传递前向 & 反向的中间输出。
比如:
-
GPU0 负责 Layer1~4
-
GPU1 负责 Layer5~8
输入 batch 被切分成 micro-batches,在 pipeline 中流动,形成并行。
✅ 优点:
-
可以训练比单卡大但没有超大参数的模型;
-
通信量少(只在层与层之间传中间激活值);
-
适合和张量并行组合使用(如 DeepSpeed/GPipe/Megatron-LM)。
❌ 缺点:
-
存在 pipeline bubble(流水线空闲时间);
-
不支持跨层依赖的模型结构(如残差连接跨层);
-
通常 batch 越大、pipeline 越长,效率越高,但越复杂。
🧠 二、那工程落地时怎么选?
✅ 我会基于以下几个关键指标来做决策:
决策因素 | 说明 | 推荐策略 |
---|---|---|
模型参数量大小 | 能否放进单卡显存? | 小模型 → 数据并行;超大模型 → 张量/流水线 |
网络结构复杂度 | 是否容易张量拆分?跨层依赖多吗? | 简单如 Transformer → 可用张量并行 |
GPU 数量与带宽 | GPU 多?NVLink 好? | GPU 多 + NVLink → 适合张量/流水线并行 |
批次大小 | 是否能提供足够多 micro-batches? | Batch 小 → 数据并行更合适 |
开发复杂度 / 时间压力 | 团队是否熟悉复杂并行策略? | 快速迭代 → 数据并行优先 |
通信开销 vs 计算开销 | 网络带宽够用吗? | 带宽低 → 优先考虑流水线并行 |
框架支持 | 用的是哪种框架?是否有 Megatron/DeepSpeed 等工具? | Megatron → 用张量并行;DeepSpeed → 灵活选择 |
🧪 实际落地场景示例:
场景 | 选择策略 |
---|---|
训练 GPT-2 / BERT,小模型(<500M) | 数据并行就够了 |
GPT-3 级别大模型(百亿参数) | Tensor Parallel + Pipeline Parallel + Data Parallel(Megatron-LM 经典组合) |
微调小模型用于下游任务 | 纯数据并行即可,部署方便 |
显存限制严重,且通信带宽高 | Tensor Parallel 合理拆层内张量 |
🧾 最后总结:
NVIDIA 提出的三种并行策略:
数据并行处理不同数据,适合小模型;
张量并行拆分层内参数,适合超大模型;
流水线并行分层部署模型,适合深网络。
实际工程中我会根据模型规模、显存限制、网络带宽、框架支持、团队经验来灵活组合这三种策略。