浅谈AI工程落地 - 分布式训练 - Nvidia Megatron LM - 三种训练策略

前言

曾在游戏世界挥洒创意，也曾在前端和后端的浪潮间穿梭，如今，而立的我仰望AI的璀璨星空，心潮澎湃，步履不停！愿你我皆乘风破浪，逐梦星辰！

好问题！你想深入理解 NVIDIA 提出的三种并行策略，以及在工程实践中怎么选择——这就是从「理解理论」到「工程落地」最关键的一步了。

---

✅ 一、NVIDIA 提出的三种并行策略

这三种并行策略主要用于训练超大规模模型，尤其在 Megatron-LM 和后续如 GPT-3 的训练中：

---

1️⃣ 数据并行（Data Parallelism）

👉 原理：

每张卡都拥有完整的模型副本，但处理不同的数据子集。
每轮训练后，所有卡计算自己的梯度，然后用 All-Reduce 操作平均梯度，同步更新参数。

✅ 优点：

• 实现简单，PyTorch DDP/TF MirroredStrategy 都支持；

• 适用于模型可以放进单卡显存的场景；

• 扩展性强，适合水平扩展训练。

❌ 缺点：

• 受限于单卡显存，模型参数不能太大；

• 每轮都要 All-Reduce，通信量跟模型大小成正比。

---

2️⃣ 张量并行（Tensor Parallelism / 纵向切分）

👉 原理：

将模型中的某些大张量（如线性层的权重矩阵）按列或行方向切分，每张卡只负责计算一部分权重对应的输出。

Megatron-LM 中举的例子：

• 把 transformer