AI Traning Parallism
导言
- AI 训练时，有些分布式训练的常见并行概念需要了解。
- 例如，TP, VP, SP, VPP
数据并行 DP
- (Data Parallelism)思想：把同一个模型放在多个GPU上，batch数据平均分布到各个GPU上，并行计算。
- 难点：注意参数同步和信息过期问题。
- 优点：加速比线性。部署简单工作量小，每个节点内的计算效率高。适合规模小，计算密集度高的模型。由于部署简单，是最先采用的并行方式。
- 缺点：需要在每个节点复制所有模型参数，显存重复度高，利用率低，并不适合大模型的部署。数据分布在不同机器，但是需要allreduce同步权重 FS（fully shared）DP
特点+具体操作：
- 同构模型，不同数据：每个节点都包含完整的模型，以及模型的参数（weight，parameter），输入数据则根据模型的并行度进行拆分，分别被每个节点读取；
- 假设我们有8张GPU卡或者昇腾的NPU卡来训练图片分类的模型，训练的批量为160，那么每张卡上面分到的批量数据（min-batch）为20，每张卡基于样本数据完成训练。
- 独立运行：每个节点读取相应的输入数据后，分别独自处理模型的前向和反向传播，并得到Gradients，归并所有的梯度并更新梯度；
- 梯度聚合：因为各张卡上处理的数据样本不同，所以获得的梯度会有些差别。因此，需要对梯度进行聚合（求和、均值）等计算来保持和单卡训练相同的结果，最后再更新参数。
- 统一通讯：所有节点之间的通信，主要包括前向传播的Loss归并以及反向传播的gradient归并以及更新，这些通信则是通过相应的通信原语（gather/reduce/broadcast）操作。
- 参数更新：梯度聚合会让各卡的模型以相同的梯度值同时进入参数更新阶段，然后针对新的数据进行下一轮训练。
[^1]
模型并行
- Model Parallelism 思想： (层内切分的并行)将模型的每一层，手动拆分成多分到不同GPU上。例如，将模型张量的参数根据特定的维度进行分割，根据手工设计，由系统分配对应训练数据和特征图（可以是分割或复制的方法）。
- 难点：需要人为的切分设计。也有必要的数据传输。
- 优点：解决了数据并行显存利用率低的问题，其通过对模型的切分，每个节点只需要放置一部分的模型参数，从而使得其可以部署更大的模型。
张量并行 TP
Tensor Parallelism 是模型并行的一种
- 思想：一个操作中进行并行计算，主要是矩阵-矩阵乘法。张量并行训练是将一个张量沿特定维度分成 N 块，每个设备只持有整个张量的 1/N，同时不影响计算图的正确性。这需要额外的通信来确保结果的正确性。
- 难点：额外的通信
- 优点：每个设备只持有整个张量的 1/N
序列并行 SP
Sequence Parallelism 是模型并行的另一种
- 思想：对于LLM的长序列文本，从序列维度进行划分[^7], 有以下几个分类：
- Megatron SP: 将层内切分的并行思想，从TP的对Tensor并行，拓展到对Norm层和Dropout层也并行[^8]。
- DeepSpeed-Ulysses: 整体从序列维度进行划分设计[^10]。
- Ring Attention: 通讯采用了Ring Self-Attention[^5]或者Ring Attention[^9].
- 相关领域的发展可以看[^6].
上下文并行 CP
Context Parallelism 是序列并行的一种
流水线并行 PP
Pipeline Model Parallelism (层间切分的模型并行)
- 思想：AI训练是重复的有依赖长过程，可以打散成有依赖的基本单元micro-batch进行流水线调度， 提高设备的利用率。
- 难点：依赖基本单元间的数据传输时间，如何隐藏。流水线并行的方式更复杂，并且micro-batch的方式减少了单节点计算密集度，增加了节点间的信息传递频率，使得取得一个好的加速比成为一个难题。
- 优点：解决了数据并行显存利用率低的问题，其通过对模型的切分，每个节点只需要放置一部分的模型参数，从而使得其可以部署更大的模型。
GPipe
1F1B
PipeDream
VPP
混合并行
- 先数据并行 + 后流水线并行
- 例子：卷积层参数量小，但计算量大（数据并行）。全连接层参数量大，但计算量小（模型并行）。
- 例子2：2021年10月，微软和英伟达联合提出了 PTD-P(Inter-node Pipeline Parallelism, Intra-node Tensor Parallelism, and Data Parallelism)训练加速方法，通过数据并行、张量并行和 Pipeline 并行“三管齐下”的方式，将模型的吞吐量提高 10%以上。该并行方法可以在3072个GPU 上，以502P的算力对一万亿参数的GPT 架构模型进行训练，实现单GPU吞吐量52%的性能提升。
专家并行 MoE
- 思想：一种基于稀疏 MoE(Mixture-of-Experts) 层的深度学习模型架构被提出，即将大模型拆分成多个小模型(专家，expert)， 每轮迭代根据样本决定激活一部分专家用于计算，
- 优点：只计算一部分，达到了节省计算资源的效果；
- 实现：MoE 将模型的某一层扩展为多个具有相同结构的专家网络(expert)，并由门(gate)网络决定激活哪些 expert 用于计算，从而实现超大规模稀疏模型的训练。
ZeRO
- Zero 优化方法有三个层次，分别是 ZeRO-1、ZeRO-2 和 ZeRO-3。
- 它们是由微软提出的 ZeRO（Zero Redundancy Optimizer） 优化技术的不同阶段，旨在减少大规模分布式训练中的内存占用。
1. ZeRO-1：优化器状态分区
- 目标：减少优化器状态的内存占用。
- 实现方式：将优化器状态（如动量、梯度方差等）分布在不同的 GPU 上，而不是在每个 GPU 上保存完整的副本。
- 优点：显著减少内存占用，同时通信开销较小。
- 适用场景：适合中等规模的模型训练。
2. ZeRO-2：梯度分区
- 目标：进一步减少梯度存储的内存占用。
- 实现方式：将梯度分区存储在不同的 GPU 上，每个 GPU 只保存一部分梯度。
- 优点：内存占用进一步降低，但通信开销有所增加，因为需要在反向传播后聚合梯度。
- 适用场景：适合大规模模型训练。
3. ZeRO-3：参数分区
- 目标：最大化内存节省，支持超大规模模型训练。
- 实现方式：将模型参数分区存储在不同的 GPU 上，每个 GPU 只保存一部分参数。
- 优点：内存占用大幅降低，可以训练非常大的模型，但通信开销最大，因为需要在每次前向和反向传播时聚合参数。
- 适用场景：适合超大规模模型训练（如 GPT、BERT 等）。
对比总结
| 特性 | ZeRO-1 | ZeRO-2 | ZeRO-3 |
|---|---|---|---|
| 分区对象 | 优化器状态 | 梯度 | 模型参数 |
| 内存节省 | 中等 | 较大 | 最大 |
| 通信开销 | 最小 | 中等 | 最大 |
| 适用场景 | 中等规模模型 | 大规模模型 | 超大规模模型 |
分层ZeRo
- 分层Zero（Hierarchical Zero）是一种用于大规模AI模型训练的优化方法，旨在解决传统Zero（Zero Redundancy Optimizer）在大规模分布式训练中的局限性。
- 它通过分层通信和计算优化，提升训练效率和扩展性。
核心思想是将计算和通信任务分层处理，减少通信开销，提高资源利用率。具体包括：
- 分层通信：将通信任务分为多个层次，优先在低层次（如节点内）完成，减少高层次（如跨节点）的通信。
- 分层计算：将计算任务分层处理，优先在低层次完成，减少高层次的计算负担。
主要特点
- 减少通信开销：通过分层通信，降低跨节点通信频率，提升效率。
- 提高资源利用率：分层计算使资源分配更合理，减少闲置。
- 增强扩展性：优化后的方法更适合大规模分布式训练，支持更多计算节点。
异构系统的并行
人们思考为什么 CPU 内存没有被用于分布式训练。
参考文献
[^1]: MindSpore (master) 分布式并行原生
[^2]: Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
[^3]: GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism
[^4]: PipeDream: Generalized Pipeline Parallelism for DNN Training
[^5]: Sequence Parallelism: Long Sequence Training from System Perspective
[^6]: DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training
[^7]: Paradigms of Parallelism
[^8]: Reducing Activation Recomputation in Large Transformer Models
[^9]: Ring Attention with Blockwise Transformers for Near-Infinite Context