分布式训练编排实战 - 第4节
4. 混合并行（Hybrid Parallelism）：数据并行+模型并行的组合策略
聊到混合并行，我得先说说自己的经历。几年前我接手一个百亿参数模型训练任务，单卡显存根本放不下，数据并行又因为通信开销太大导致效率极低。那时候我就在想，能不能把几种并行方式组合起来用？说白了，这就是混合并行的核心思路——不把鸡蛋放在一个篮子里。
4.1 为什么需要混合并行？
你想想看，数据并行虽然简单，但每个GPU都要存一份完整的模型。模型大到一定程度，比如GPT-3这种1750亿参数的怪兽，单卡显存根本装不下。模型并行呢？虽然能拆分模型，但计算效率往往不如数据并行高。
所以实际工程中，我们通常把两者结合起来。我习惯的做法是：先做模型并行把模型拆开，再做数据并行把数据分片。这样既解决了显存瓶颈，又能充分利用多卡的计算能力。
核心原则：混合并行的目标是在显存约束和计算效率之间找到最佳平衡点。不是简单叠加，而是根据模型结构和硬件拓扑做精细化设计。
4.2 混合并行的典型架构
我在项目中常用的混合并行架构是这样的：
- 层内并行（张量并行）：把单个Transformer层的参数切分到多卡上
- 层间并行（流水线并行）：把不同层分配到不同设备上
- 数据并行：在以上基础上，复制多份模型副本处理不同数据
举个例子，假设我们有64张GPU，训练一个12层的Transformer模型。我可能会这样分配：
| 并行策略 | 切分方式 | GPU数量 |
|---|---|---|
| 张量并行 | 每层切4份 | 4 |
| 流水线并行 | 12层分3组 | 3 |
| 数据并行 | 复制多份 | 64/(4*3)≈5 |
嗯，这里要注意：张量并行的通信量最大，所以通常只在单机内使用（NVLink连接）。流水线并行次之，数据并行通信量最小。
3D并行：数据并行+张量并行+流水线并行
3D并行，说白了就是把三种并行方式组合成一个三维网格。我最早接触这个概念是在Megatron-LM项目中，当时觉得这设计真巧妙。
4.3 3D并行的三维视角
你可以把3D并行想象成一个三维坐标系：
- X轴：数据并行 — 复制模型副本，处理不同数据批次
- Y轴：张量并行 — 切分单个层的参数
- Z轴：流水线并行 — 切分模型的层
每个GPU在这个三维网格中都有一个坐标(x, y, z)，负责处理特定的数据分片、特定的张量分片和特定的层。
我的经验：3D并行最头疼的是通信拓扑设计。我建议把通信密集的张量并行放在节点内部，流水线并行放在节点之间，数据并行则全局使用。这样能最大化利用带宽。
4.4 3D并行的通信模式
3D并行涉及三种通信模式，每种的开销都不一样：
| 并行类型 | 通信模式 | 通信量 | 频率 |
|---|---|---|---|
| 张量并行 | All-Reduce | 高（模型参数大小） | 每个计算步骤 |
| 流水线并行 | P2P Send/Recv | 中（激活值大小） | 每个微批次 |
| 数据并行 | All-Reduce | 低（梯度大小） | 每个训练步骤 |
我曾经踩过一个坑：把张量并行的通信域设得太大，导致跨节点通信成了瓶颈。后来我改成只在单机8卡内做张量并行，性能直接翻倍。
ZeRO优化器：内存节省原理
ZeRO（Zero Redundancy Optimizer）是DeepSpeed推出的内存优化技术。我刚开始用的时候，觉得这玩意儿简直就是显存救星。
4.5 ZeRO-1：优化器状态分区
ZeRO-1只分区存储优化器状态（比如Adam的动量和方差）。每个GPU只保存自己那部分优化器状态，而不是完整复制。
内存节省效果：
- 原始：每个GPU存完整优化器状态（4倍模型参数大小）
- ZeRO-1：每个GPU只存1/N（N为GPU数量）
- 节省比例：约75%（假设N=4）
关键点：ZeRO-1不减少模型参数和梯度的内存占用，只优化优化器状态。如果你的模型已经能放进显存，但优化器状态放不下，ZeRO-1是首选。
4.6 ZeRO-2：梯度分区
ZeRO-2在ZeRO-1基础上，进一步分区存储梯度。每个GPU只保存自己那部分梯度，计算完就丢弃其他部分。
内存节省效果：
- 原始：每个GPU存完整梯度（1倍模型参数大小）
- ZeRO-2：每个GPU只存1/N
- 加上ZeRO-1：总共节省约87.5%的优化器相关内存
我记得有次训练一个30亿参数的模型，用ZeRO-2后，原本需要8张A100才能跑的任务，4张就搞定了。嗯，真香。
4.7 ZeRO-3：参数分区
ZeRO-3是终极方案，连模型参数都分区存储。每个GPU只保存自己那部分参数，计算时通过All-Gather通信获取完整参数。
内存节省效果：
- 原始：每个GPU存完整参数（1倍模型参数大小）
- ZeRO-3：每个GPU只存1/N
- 加上ZeRO-1/2：总共节省约95%的内存
注意：ZeRO-3虽然省内存，但通信开销最大。每个前向/反向传播都需要All-Gather参数。我建议只在模型实在放不下时才用ZeRO-3，否则优先用ZeRO-2。
4.8 ZeRO各版本对比
| 版本 | 分区内容 | 内存节省 | 通信开销 | 适用场景 |
|---|---|---|---|---|
| ZeRO-1 | 优化器状态 | ~75% | 低 | 优化器状态占大头 |
| ZeRO-2 | 优化器状态+梯度 | ~87.5% | 中 | 梯度也放不下 |
| ZeRO-3 | 优化器状态+梯度+参数 | ~95% | 高 | 模型完全放不下 |
4.9 ZeRO与混合并行的结合
实际工程中，ZeRO经常和混合并行一起用。我常用的组合是：
- 小模型（<10B）：数据并行 + ZeRO-2
- 中模型（10B-100B）：3D并行 + ZeRO-1
- 大模型（>100B）：3D并行 + ZeRO-3
为什么会这样？因为模型越大，参数本身占的内存就越多，ZeRO-3的优势就越明显。但也要权衡通信开销，毕竟All-Gather参数可不是免费的。
避坑指南：我曾经在ZeRO-3下使用流水线并行，结果发现通信重叠做得不好，导致GPU利用率只有40%。后来我调整了微批次大小和流水线深度，才把利用率提到80%以上。记住：ZeRO-3的通信和计算要尽量重叠。
4.10 知识体系图
下面这张图展示了混合并行、3D并行和ZeRO优化器的关系：
这张图清晰地展示了混合并行、3D并行和ZeRO优化器的层级关系。从底层的数据并行、张量并行、流水线并行，到上层的3D并行组合，再到ZeRO优化器的内存优化，每一层都有其独特的作用。
我个人建议，在实际项目中先评估模型大小和硬件资源，再选择合适的并行策略组合。不要一上来就上3D并行+ZeRO-3，有时候简单的数据并行+ZeRO-2反而更高效。
想要更多资料，联系可可