📚分布式训练系列文章
随着模型参数量快速增长，超过单张 GPU 显存承载能力的情况越来越常见。对于超大规模模型（如百亿参数以上），仅靠数据并行已无法训练。这时，模型并行（Model Parallelism） 应运而生，它通过将模型本身拆分到多张 GPU 上进行计算，实现显存扩展和训练加速。本文将详细解析张量并行、流水线并行和混合并行的原理、实践方法以及优化策略。
所有相关源码示例、流程图、模型配置与知识库构建技巧，我也将持续更新在Github：AIHub，欢迎关注收藏！
希望大家带着下面的问题来学习，我会在文末给出答案。
- 模型并行和数据并行有什么区别？
- 张量并行、流水线并行和混合并行的原理是什么？
- 如何选择合适的模型并行策略来训练超大模型？
1. 模型并行概述
模型并行的核心思想是将模型本身拆分到多张 GPU 上，每张 GPU 只计算模型的一部分，模型并行能很好的避免单卡显存溢出，使训练超大模型成为可能，并且模型并行可以与数据并行结合，形成混合并行策略，提高训练效率。
区别于数据并行：
- 数据并行是 复制模型到每张 GPU 并分批处理数据；
- 模型并行是 拆分模型到不同 GPU，每张 GPU 只计算部分模型；
- 模型并行通信压力主要来自 模型层之间的张量传输，而数据并行通信主要是 梯度同步。
2. 张量并行（Tensor Parallel）
张量并行将模型层中的矩阵或张量切分到多张 GPU 上，每张 GPU 只计算张量的一部分，然后通过 AllReduce 或通信操作汇总结果，张量并行常用于全连接层、注意力机制中权重矩阵很大的情况。
张量并行进行精细化切分计算，可减少单卡显存压力，但是对通信效率要求较高，需要高带宽网络。
以全连接层为例，下面是一个张量并行的示例：
# 假设有两张 GPU，要对权重矩阵 W 进行切分
# GPU 0: W[:, :hidden//2]
# GPU 1: W[:, hidden//2:]
# 前向计算分别计算部分输出，最后合并结果
output = torch.cat([output_gpu0, output_gpu1], dim=-1)
3. 流水线并行（Pipeline Parallel）
流水线并行将模型按层或阶段划分到不同 GPU 上，形成流水线，每张 GPU 处理输入的不同阶段，并在计算完成后传递给下一阶段，可以同时处理多个 mini-batch，通过流水线提高 GPU 利用率。
流水线并行可以训练比单卡显存大得多的模型，但是可能产生 pipeline bubble（GPU 空闲等待），需要通过 micro-batch 技术优化。
下面是流水线并行的实践示例：
# 假设4层模型分到2张GPU
# GPU0: layer1 + layer2
# GPU1: layer3 + layer4
for micro_batch in micro_batches:
out0 = gpu0_forward(micro_batch)
out1 = gpu1_forward(out0)
loss = compute_loss(out1)
backward_pass(loss)
4. 混合并行（Hybrid Parallel）
混合并行顾名思义，将数据并行和模型并行结合，在每个节点内部使用模型并行（张量并行/流水线并行），跨节点使用数据并行，常用于训练百亿甚至千亿参数模型。
混合并行适合超大规模模型训练，通信和调度复杂，需要精心设计 GPU 分配和批次调度策略。
最后，我们回答文章开头提出的问题。
- 模型并行和数据并行有什么区别？
数据并行复制模型到每张 GPU 并分批处理数据；模型并行拆分模型到不同 GPU，每张 GPU 只计算部分模型。
- 张量并行、流水线并行和混合并行的原理是什么？
- 张量并行：拆分层内权重矩阵或张量；
- 流水线并行：按层划分模型形成流水线；
- 混合并行：结合模型并行和数据并行，提高大模型训练效率。
- 如何选择合适的模型并行策略？
- 模型稍大但单卡不足：张量并行；
- 模型超大层数多：流水线并行；
- 模型极大（百亿参数以上）且有多节点：混合并行。
关于深度学习和AI大模型相关的知识和前沿技术更新，请关注公众号coting！
以上内容部分参考了开源框架文档和论文，非常感谢，如有侵权请联系删除！
参考链接
