单卡放不下?数据并行太慢?模型并行利用率低?3D并行一次性解决所有问题
前言
前两节我们分别学习了分布式训练的三种基本策略:数据并行、模型并行、流水线并行。但现实中的大模型训练(如GPT-3、LLaMA、Mixtral)用的不是单一策略,而是3D并行——三种策略的组合。
为什么需要组合?因为:
- 仅用数据并行:模型太大放不进单卡(70B模型需要140GB+)
- 仅用模型并行:GPU利用率低(某卡计算时其他卡空闲)
- 仅用流水线并行:仍有流水线气泡,且单卡仍存整层参数
今天,我们将深入两个工业级框架的实现:
- Megatron-LM(NVIDIA):张量并行 + 流水线并行 + 数据并行
- DeepSpeed(微软):ZeRO + 流水线并行 + 3D并行
通过剖析它们的实现,你将掌握3D并行的核心技术,能够为大模型设计合理的并行策略。
一、Megatron-LM:张量并行 + 流水线并行 + 数据并行
1.1 Megatron-LM的核心理念
Transformer模型具有天然的可并行结构:
- 多头注意力:各头可以并行