1. 分布式训练概述

1.概述.md

1.1. 数据并行

在数据并行训练过程中，整个数据集被划分为多个小块，每个小块被独立分配给不同的计算设备。这种并行化策略主要在批次（Batch）维度上展开，意味着训练过程在多个设备上同时进行。每个设备上都运行着模型的一个完整副本，并且仅针对其分配的数据块进行训练。训练完成后，通过反向传播得到的梯度会在所有设备间汇总，以确保各设备上的模型参数保持一致性。 PyTorch的分布式数据并行（DDP）是实现数据并行的一个典型例子。

1.2. 模型并行

在数据并行训练中，一个显著特征是每个GPU都拥有模型权重的完整副本，这不可避免地导致了资源冗余。与之相对的是模型并行策略，该策略将模型分割并跨多个设备分布。

模型并行主要分为两大类：张量并行和流水线并行。

张量并行：专注于在单个操作中实现并行化，例如在矩阵-矩阵乘法中同时处理多个数据点（可视为层内并行，因为它在单个层内分配计算任务）。
流水线并行：在模型的不同层之间进行并行处理，允许各层同时进行计算（相当于层间并行，因为它在模型的不同层之间分配计算任务）。

1.2.1. 张量并行

张量并行训练技术通过沿特定维度将一个张量划分为N个部分，使得每个设备仅存储整个张量的1/N，而不破坏计算图的完整性。为了确保计算结果的准确性，需要进行额外的通信步骤。

以矩阵乘法为例，考虑计算 $C = AB$ 。我们可以将矩阵 $B$ 按列分割成多个子矩阵 $[B0, B1, B2, ..., Bn]$ ，每个设备负责一个子矩阵。接着，每个设备将本地的子矩阵 $Bi$ 与矩阵 $A$ 相乘，得到一系列部分乘积 $[AB0, AB1, AB2, ..., ABn]$ 。此时，每个设备仅持有部分结果，例如，设备0（rank=0）持有 $AB0$ 。为了获得完整的结果 $C$ ，我们需要将所有设备上的部分乘积收集起来，并沿着列维度进行拼接。这样，我们就能在多个设备间分配张量计算任务，同时确保整个计算过程的正确性。

典型的张量并行实现：Megatron-LM（1D）、Colossal-AI（2D、2.5D、3D）。

1.2.2. 流水线并行

流水线并行的核心理念在于将模型分解成多个层级，并分别分配给不同的计算设备。具体来说：

模型分割：模型被划分为若干独立的层块，每个层块负责处理模型的一部分计算任务。
前向传播优化：在前向传播阶段，每个设备负责计算其分配的层块，并将其产生的中间激活值传递给流水线中的下一个设备，以此类推。
后向传播优化：在后向传播阶段，每个设备接收来自下一个阶段的输入张量梯度，并将其回传至前一个阶段，以此类推。
并行计算优势：通过这种方式，各个设备可以并行执行各自的计算任务，显著提高了模型训练的并行度和吞吐量，从而加快训练速度。

流水线并行训练的一个明显缺点是训练设备容易出现空闲状态（因为后一个阶段需要等待前一个阶段执行完毕），导致计算资源的浪费，加速效率没有数据并行高。

典型的流水线并行实现：GPipe、PipeDream、PipeDream-2BW、PipeDream Flush（1F1B）。

1.3. 优化器相关并行

目前随着模型越来越大，单个GPU的显存目前通常无法装下那么大的模型了。那么就要想办法对占显存的地方进行优化。

通常来说，模型训练的过程中，GPU上需要进行存储的参数包括了模型本身的参数、优化器状态、激活函数的输出值、梯度以及一些零时的Buffer。

模型参数仅占模型训练过程中所有数据的一部分，当进行混合精度运算时，模型状态参数（优化器状态+梯度+模型参数）占到了一大半以上。因此，我们需要想办法去除模型训练过程中的冗余数据。

优化器相关的并行就是一种去除冗余数据的并行方案，目前这种并行最流行的方法是 ZeRO（即零冗余优化器）。 ZeRO 通过分片技术优化模型状态的存储，确保每张 GPU 仅存储模型状态的 1/N 部分，从而在整个系统中只维护一份模型状态。 ZeRO 分为三个级别，它们分别对应不同程度的模型状态分片：

ZeRO-1 : 对优化器状态分片（Optimizer States Sharding）
ZeRO-2 : 对优化器状态和梯度分片（Optimizer States & Gradients Sharding）
ZeRO-3 : 对优化器状态、梯度分片以及模型权重参数分片（Optimizer States & Gradients & Parameters Sharding）

1.4. 多维混合并行

多维混合并行指将数据并行、模型并行和流水线并行等多种并行技术结合起来进行分布式训练。通常，在进行超大规模模型的预训练和全参数微调时，都需要用到多维混合并行。

为了充分利用带宽，通常情况下，张量并行所需的通信量最大，而数据并行与流水线并行所需的通信量相对来说较小。因此，同一个服务器内使用张量并行，而服务器之间使用数据并行与流水线并行。

1.5. MoE 并行/专家并行

通常来讲，模型规模的扩展会导致训练成本显著增加，计算资源的限制成为了大规模密集模型训练的瓶颈。为了解决这个问题，一种基于稀疏 MoE 层的深度学习模型架构被提出，即将大模型拆分成多个小模型(专家，expert)，每轮迭代根据样本决定激活一部分专家用于计算，达到了节省计算资源的效果；并引入可训练并确保稀疏性的门(gate)机制，以保证计算能力的优化。

使用 MoE 结构，可以在计算成本次线性增加的同时实现超大规模模型训练，为恒定的计算资源预算带来巨大增益。而 MOE 并行，本质上也是一种模型并行方法。下图展示了一个有六个专家网络的模型被两路专家并行地训练。其中，专家 1-3 被放置在第一个计算单元上，而专家 4-6 被放置在第二个计算单元上。

分布式训练