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）操作。
参数更新：梯度聚合会让各卡的模型以相同的梯度值同时进入参数更新阶段，然后针对新的数据进行下一轮训练。

¹

模型并行¶

Model Parallelism 思想： (层内切分的并行)将模型的每一层，手动拆分成多分到不同GPU上。例如，将模型张量的参数根据特定的维度进行分割，根据手工设计，由系统分配对应训练数据和特征图（可以是分割或复制的方法）。
难点：需要人为的切分设计。也有必要的数据传输。
优点：解决了数据并行显存利用率低的问题，其通过对模型的切分，每个节点只需要放置一部分的模型参数，从而使得其可以部署更大的模型。

Megatron-LM模型在训练中采用更为手工的方式，将每个transformer块都进行了分割，实现了高性能的计算。

²

张量并行 TP¶

Tensor Parallelism 是模型并行的一种

思想：一个操作中进行并行计算，主要是矩阵-矩阵乘法。张量并行训练是将一个张量沿特定维度分成 N 块，每个设备只持有整个张量的 1/N，同时不影响计算图的正确性。这需要额外的通信来确保结果的正确性。
难点：额外的通信
优点：每个设备只持有整个张量的 1/N

序列并行 SP¶

Sequence Parallelism 是模型并行的另一种

思想：对于LLM的长序列文本，从序列维度进行划分⁷, 有以下几个分类：
- Megatron SP: 将层内切分的并行思想，从TP的对Tensor并行，拓展到对Norm层和Dropout层也并行⁸。
- DeepSpeed-Ulysses: 整体从序列维度进行划分设计¹⁰。
- Ring Attention: 通讯采用了Ring Self-Attention⁵或者Ring Attention⁹.
相关领域的发展可以看⁶.

上下文并行 CP¶

Context Parallelism 是序列并行的一种

流水线并行 PP¶

Pipeline Model Parallelism (层间切分的模型并行)

思想：AI训练是重复的有依赖长过程，可以打散成有依赖的基本单元micro-batch进行流水线调度，提高设备的利用率。
难点：依赖基本单元间的数据传输时间，如何隐藏。流水线并行的方式更复杂，并且micro-batch的方式减少了单节点计算密集度，增加了节点间的信息传递频率，使得取得一个好的加速比成为一个难题。
优点：解决了数据并行显存利用率低的问题，其通过对模型的切分，每个节点只需要放置一部分的模型参数，从而使得其可以部署更大的模型。

GPipe¶

³

1F1B¶

1F1B

PipeDream¶

⁴

VPP¶

Virtual Pipeline Parallelism

混合并行¶

先数据并行 + 后流水线并行
例子：卷积层参数量小，但计算量大（数据并行）。全连接层参数量大，但计算量小（模型并行）。
例子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 内存没有被用于分布式训练。