单卡放不下？数据并行太慢？模型并行利用率低？3D并行一次性解决所有问题

前言

前两节我们分别学习了分布式训练的三种基本策略：数据并行、模型并行、流水线并行。但现实中的大模型训练（如GPT-3、LLaMA、Mixtral）用的不是单一策略，而是3D并行——三种策略的组合。

为什么需要组合？因为：

• 仅用数据并行：模型太大放不进单卡（70B模型需要140GB+）
• 仅用模型并行：GPU利用率低（某卡计算时其他卡空闲）
• 仅用流水线并行：仍有流水线气泡，且单卡仍存整层参数

今天，我们将深入两个工业级框架的实现：

• Megatron-LM（NVIDIA）：张量并行 + 流水线并行 + 数据并行
• DeepSpeed（微软）：ZeRO + 流水线并行 + 3D并行

通过剖析它们的实现，你将掌握3D并行的核心技术，能够为大模型设计合理的并行策略。

一、Megatron-LM：张量并行 + 流水线并行 + 数据并行

1.1 Megatron-LM的核心理念

Transformer模型具有天然的可并行结构：

• 多头注意力：各头可以并行