大模型训练并行策略

date

Mar 29, 2025

slug

parallel_in_transformer_training

TODO

补充pytorch的FSDP

补充CP的内容: ‣

🥰总结

⁉️问题

🧐内容

背景

近几年模型的参数数量不断攀升，从最初的数千扩展到如今数千亿乃至数万亿级别。大模型虽然带来了前所未有的应用效果，但与此同时，也引发了计算资源、内存管理和训练稳定性等一系列严峻挑战。因此本博客总结了一些常用分布式并行训练和内存管理技术，希望能够帮助大家更好地训练和优化大模型。

大模型的训练挑战

参数规模爆炸式增长

随着对模型容量和性能的不断追求，神经网络的参数数量呈现出指数级增长。现今从百万级到数十亿、数千亿甚至数万亿参数的模型层出不穷。例如，Llama 3.1 405B 拥有约 4,050 亿参数，而据传 GPT-4 的参数量可能达到 1.7 万亿级别。这种庞大的参数规模使得计算和内存需求急剧上升，给训练过程带来了前所未有的压力。

计算复杂度剧增

参数数量的急速增加直接导致整体计算复杂度大幅上升。训练一次大型模型可能需要耗费数周甚至数月的时间，即便采用大规模高性能 GPU 集群，训练周期仍难以令人满意，从而严重制约了模型迭代速度和研究效率。

内存瓶颈日益凸显

除了需要存储庞大的模型参数之外，大模型在训练过程中还必须保存中间激活值、梯度以及优化器状态，这些数据对 GPU 显存构成了巨大挑战。即使是配备 A100、H100(80GB 显存)、H200(141GB 显存)或 GB200(384GB 显存)等高端 GPU，单卡内存往往也难以满足数千亿甚至数万亿级模型的需求，“Out of Memory(OOM)”错误频发。

通信开销成为瓶颈

在多 GPU 分布式训练环境中，节点间需要频繁进行数据同步(如梯度汇总)。随着模型规模和 GPU 数量的增加，这种通信量急剧上升，即使在高带宽网络中，All-Reduce 操作传输海量数据也会消耗大量时间，成为整体并行效率的主要瓶颈之一。

训练稳定性挑战

超大规模模型在训练过程中更容易遭遇梯度消失或梯度爆炸问题，导致训练过程不稳定、难以收敛。虽然混合精度训练可以在一定程度上加速训练并降低显存占用，但同时也可能引入新的数值稳定性问题，要求研究人员投入更多精力进行细致调优。

分布式训练的必要性

面对上述挑战，分布式训练技术成为支撑大模型训练的关键方案。通过将训练任务拆分并分配到多台 GPU 或计算节点上，分布式训练能够充分利用并行计算和集群内存资源，从而突破单 GPU 的局限，主要优势体现在以下几个方面：

突破单 GPU 算力限制

单个 GPU 的计算能力终究有限，无法应对万亿级别参数模型的庞大计算需求。借助数据并行与模型并行技术，训练任务可以均匀分布至多个 GPU，从而大幅缩短整体训练时间。

克服单 GPU 内存瓶颈

通过将模型参数、中间激活值和优化器状态分散存储在多个 GPU 显存中，分布式训练有效扩展了可用内存容量。典型技术如 ZeRO，通过对模型参数、梯度以及优化器状态进行分片，大幅降低单卡的显存负担，使得超大规模模型的训练成为可能。

加速模型迭代与研发周期

分布式训练的高并行度使得原本需要数周甚至数月完成的训练任务有望在数天内完成，从而大幅提升模型迭代速度，使得新架构和新策略能够更快得到验证与应用。

支持更大规模模型探索

分布式训练为探索更大规模、更复杂的神经网络架构提供了坚实基础。正是在这种技术支持下，万亿参数级别的模型(如 Switch Transformer)才得以成功训练并投入实际应用。

提高训练系统的鲁棒性与可扩展性

分布式系统具备出色的容错能力，当某个 GPU 节点出现故障时，其他节点可迅速接管任务，确保训练过程不被中断。同时，集群规模可以根据具体需求灵活扩展或缩减，满足不同规模模型的训练要求。

并行训练策略

下图直观展示了多种并行训练策略的不同之处。不同颜色代表不同的模型层(例如三层)，虚线将不同的 GPU 区分开。从左到右分别是数据并行、模型并行(含流水线并行和张量并行)以及专家并行(MoE)。

Fig. 1. An illustration of various parallelism strategies on a three-layer model. Each color refers to one layer and dashed lines separate different GPUs. (Image source: OpenAI Blog, 2022) — **Fig. 1. An illustration of various parallelism strategies on a three-layer model. Each color refers to one layer and dashed lines separate different GPUs. (Image source:** **OpenAI Blog, 2022**)

数据并行

完整模型会被拷贝到每个 GPU 上，数据集则被切分为不同批次分配给各个 GPU 并行计算，最终在参数更新时聚合所有 GPU 的梯度。

模型并行

将模型划分到不同的 GPU 上，每个 GPU 只负责模型的一部分计算；可进一步分为以下两类：

流水线并行：按层(垂直方向)拆分模型，不同 GPU 负责不同的层，通过微批次(micro-batch)在流水线中传递来并行执行前向和反向计算。
张量并行：在层内(水平方向)对大规模张量操作(如大矩阵乘法)进行切分，各 GPU 并行完成这部分运算并在必要时进行聚合。

专家并行

通过门控策略，让每个输入样本只经过部分专家(子网络)，从而将整个模型按“专家模块”分布到不同 GPU。常见于 Mixture-of-Experts(MOE) 结构，可实现超大参数规模但推理/训练时仅激活部分专家。

下面我将对多种并行方式进行详细说明。

数据并行

Fig. 2. Data Parallelism. (Image source: Clolossal-AI Documentation) — **Fig. 2. Data Parallelism. (Image source:** **Clolossal-AI Documentation**)

在深度学习训练中，数据并行(Data Parallelism, DP) 是最常用的并行策略，其核心思路是：

复制模型参数：在每个计算设备(通常是 GPU)上都放置一份完整的模型参数。

划分训练数据：将大规模的数据集按样本维度拆分为多个子集，不同子集分配给不同的 GPU 进行处理。

局部前向与反向传播：每个 GPU 独立计算损失及对应的局部梯度。

梯度/参数同步：将各 GPU 的梯度聚合后更新模型参数，保证在每一次迭代后所有 GPU 的模型副本保持一致。

下面展示了 数据并行 工作流程：

数据集划分

将训练数据集 D 划分为 N 个互不重叠的子集，其中 N 是 GPU 数量。通常会确保各子集大小相近，以实现负载均衡。

模型复制

在每个 GPU 上复制一份完整的模型参数 θ。在训练开始时，这些参数在各 GPU 上都是相同的。

数据分发

将子集分发给第 i 张 GPU，让其在本地存储并供后续计算使用。

局部前向传播

每个 GPU 基于其本地数据子集做前向传播，得到局部损失。

局部反向传播

每个 GPU 基于局部损失进行反向传播，计算局部梯度

梯度同步

各 GPU 之间执行梯度同步(常用 All-Reduce)，将所有局部梯度 g1,g2,…,gN 汇总得到全局平均梯度

参数更新

每个 GPU 使用全局平均梯度, 更新本地模型参数：

其中为学习率(learning rate)。

迭代循环

重复步骤 4 - 7，直至模型达到收敛或达到预设的训练轮数(epochs)。

批量同步并行与异步并行

在上面的第 6 步“梯度同步”中，如何以及何时进行“同步”是影响数据并行性能和收敛行为的重要因素之一。一般分为以下两大类：

批量同步并行(Bulk Synchronous Parallel, BSP) 是数据并行中最常见、也是最易理解的同步模式。其特点可概括为「在每一次小批量(mini-batch)迭代结束后，全局同步一次梯度并更新参数」。具体流程：

局部计算：各 GPU 基于其数据子集分别做前向与反向传播，得到局部梯度。

全局通信：所有 GPU 同步(如通过 All-Reduce)计算。

参数更新：每个节点均使用更新本地参数副本。

迭代：所有节点完成上述操作后，再进入下一个迭代。

异步并行(Asynchronous Parallel, ASP) 旨在摆脱 BSP 的全局同步点，让各节点独立进行计算和参数更新。其典型实现是「参数服务器」(Parameter Server, PS) 架构下的 异步 push-pull 过程：

各节点在本地计算得到梯度 , 然后 push 到参数服务器；

参数服务器一旦收到梯度，立即更新全局模型参数；

其他节点在需要最新参数时，会 pull 下来继续下一步计算。

BSP vs. ASP

下表总结了在数据并行环境下，同步并行与异步并行的主要差异：

对比维度	同步并行(BSP)	异步并行(ASP)
参数更新时机	每个小批量或一定迭代后，全局同步一次	各节点独立更新参数，无需与他人保持同一时间步
收敛稳定性	高。使用的梯度均为最新，收敛路径可控，易于分析	较低。存在过时梯度，收敛速率与稳定性可能受影响
通信需求	高度依赖 All-Reduce，同步时所有节点都需要等待和交换数据	每个节点向参数服务器异步推送/拉取，通信更为灵活，但参数服务器可能成为瓶颈
硬件资源利用	若有慢节点或网络延迟，则其他节点需等待，资源利用率可能降低	无需等待慢节点，可高效使用计算资源
实现复杂度	相对较低，主流框架(PyTorch DDP、Horovod 等)有内置支持	相对更高，需要参数服务器等组件，需处理更多的同步逻辑与数据一致性
适用场景	同构硬件、网络带宽良好、追求较高收敛质量	异构硬件、网络不稳定或带宽较低、需要极高的吞吐量且能容忍一定收敛风险
典型实现	PyTorch DDP、TensorFlow MirroredStrategy	Parameter Server 架构(MXNet、TensorFlow ParameterServer 模式等)

梯度累积

当批量大小较大或通信成为主要瓶颈时，可以采用 梯度累积(Gradient Accumulation) 来减少同步频率。其核心思路是：

连续计算多个小批量(mini-batch)的局部梯度，并将它们累加到本地的累积缓冲区中；

当累积的 mini-batch 数量达到 K 时，再触发一次全局梯度同步与参数更新。

设第 j 个 mini-batch 的梯度为，则在一个「累积周期」内得到

再用学习率 η 更新：

由于梯度同步不再是每个 mini-batch 都进行，而是每累计 K 个 mini-batch 执行一次，通信开销可显著降低。但参数更新频率降低也可能导致训练收敛速度放缓，需在吞吐量与收敛性能之间做权衡。

分布式数据并行

分布式数据并行(Distributed Data Parallel, DDP) 是 PyTorch v1.5(Li et al. 2020)在 BSP 思想下的高度优化实现，为单机多 GPU 乃至多机多 GPU 的数据并行提供便利。其主要优化包括：

梯度 Bucketing(梯度桶化)：将模型参数分为多个「桶」(bucket)；反向传播时一旦某个桶内所有梯度都已计算完，就立即启动一次针对该桶的 All-Reduce，而不是等到所有梯度都算完后再一次性同步。

通信与计算重叠：DDP 通过异步通信和非阻塞操作，尽可能地将梯度同步(通信)与前向传播、反向传播(计算)重叠，从而减少了通信开销。这种重叠策略提升了整体的并行效率。

梯度累积：DDP 也能方便地与梯度累积相结合，结合使用，通过增加每次同步的梯度更新间隔，从而减少同步频率。这在大规模分布式训练中有助于进一步降低通信开销，提高训练效率。

Fig. 3. Pseudo code for Pytorch DDP. (Image source: Li et al. 2020) — **Fig. 3. Pseudo code for Pytorch DDP. (Image source:** **Li et al. 2020**)

Ring All-Reduce

在多 GPU(尤其是单机多 GPU)环境下，若有高速互联(如 NVLink、PCIe 交换机等)，可使用 Ring All-Reduce 来显著降低通信开销。其思路是：

将 k 个节点组织成一个环，并把梯度向量等分成 k 份。

在reduce-scatter，每个节点分别向下一个节点发送其本地的一部分梯度，并与收到的梯度相加；该过程循环若干次后，每个节点会持有完整的「聚合后」梯度。

在allgather，再将最终结果沿环路分发给所有节点。

理想情况下，Ring All-Reduce 的通信代价与节点数量近似无关(可以视为 O(1))，非常适合多 GPU 环境下的梯度同步，是 Horovod、NCCL 等库中广泛使用的核心通信模式。

参数服务器

当集群规模扩展至多机多 GPU 时，若简单地采用单点聚合(例如一台中心服务器)往往难以支撑海量数据的并行训练。参数服务器(Parameter Server, PS)(Li, et al. 2014)是为可扩展分布式训练而设计的一种典型架构：

参数分片：将模型参数按键值对(key-value) 的形式进行拆分，不同 PS 节点只管理特定分片的参数；

push-pull 语义：计算节点在本地得到梯度后，push 到相应的 PS；PS 更新完该分片参数后，计算节点可在需要时 pull 下最新版本进行下一步计算。

灵活容错与扩展：通过增加或移除 PS 节点，可在带宽或计算需求上灵活扩容；在 PS 上也能实现备份与容错策略。

这种 PS + Worker 模式可以同时结合数据并行和模型并行，将超大模型拆分到多个 PS 上存储，并对超大数据进行分布式训练。PS 本身也能根据负载情况做拆分与合并，形成更加复杂的层次化拓扑结构。

模型并行

模型并行(Model Parallelism, MP) 是一种将模型本身分割到多个计算设备(GPU) 上进行训练的并行方式。当模型参数规模超过单个 GPU 的内存容量时，模型并行成为必要的选择。模型并行主要分为两种类型：流水线并行(Pipeline Parallelism) 和张量并行(Tensor Parallelism)。

朴素模型并行与气泡问题

Fig. 4. A naive model parallelism setup where the model is vertically split into 4 partitions. Data is processed by one worker at a time due to sequential dependency, leading to large “bubbles” of idle time. (Image source: Huang et al. 2018) — Fig. 4. A naive model parallelism setup where the model is vertically split into 4 partitions. Data is processed by one worker at a time due to sequential dependency, leading to large “bubbles” of idle time. (Image source: **Huang et al. 2018**)

朴素的模型并行实现，即将模型简单地按层划分，并顺序地在不同 GPU 上执行，会遇到严重的 “气泡”(bubble) 问题。由于层之间的依赖关系，当一个 GPU 在处理某个数据样本的某个阶段时，其他 GPU 可能处于空闲状态，等待前一个 GPU 的输出或者后一个 GPU 的输入。这种 GPU 空闲时间被称为 “气泡”，严重降低了流水线并行的效率。

其中，表示的前向传播，表示的反向传播。可以看到，在朴素流水线并行中，大部分时间只有一个 GPU 在工作，其他 GPU 处于空闲状态，效率低下。

气泡问题产生的原因:

层间依赖: 神经网络的层之间存在顺序依赖关系，后一层的计算必须依赖于前一层的输出。

顺序执行: 朴素模型并行按照层顺序依次执行，导致 GPU 之间无法充分并行工作。

流水线并行

Fig. 5. Pipeline Parallelism. (Image source: Clolossal-AI Documentation) — **Fig. 5. Pipeline Parallelism. (Image source:** **Clolossal-AI Documentation**)

流水线并行(Pipeline Parallelism, PP) 将模型按层划分为多个阶段(stage)，每个阶段分配到一个 GPU 上。数据像流水线一样在不同 GPU 之间传递，前一个 GPU 的输出作为后一个 GPU 的输入。流水线并行旨在提高模型并行训练的效率，减少 GPU 空闲时间。

GPipe

GPipe(Huang et al. 2018) 是 Google 提出的一个高效的流水线并行训练系统，旨在解决朴素流水线并行的气泡问题。GPipe 的核心思想是将 mini-batch 划分为多个 micro-batch，并采用同步梯度聚合的方式来缓解气泡问题，提高流水线效率。

以下是 GPipe 调度策略：

Micro-batch 划分: 将一个 mini-batch 划分为 m 个 micro-batch。划分后的每个 micro-batch 的大小为原 mini-batch 的。

流水线阶段划分: 将模型按层划分为 d 个阶段，每个阶段分配到一个 GPU 上。

流水线执行: 依次处理每个 micro-batch，在流水线中进行前向和反向传播。具体流程如下：

前向传播(Forward Propagation): 对于每个 micro-batch，依次在 Stage 1, Stage 2 , …, Stage 上进行前向传播。Stage 的输出作为 Stage的输入。

反向传播(Backward Propagation): 当所有 micro-batch 的前向传播都完成后，开始反向传播。对于每个 micro-batch，依次在 Stage , Stage , …, Stage 1 上进行反向传播。Stage 的梯度作为 Stage的输入。

同步梯度聚合(Synchronous Gradient Aggregation): 在所有 micro-batch 的反向传播都完成后，将所有 micro-batch 的梯度进行聚合(例如求平均)，得到全局平均梯度。

参数更新(Parameter Update): 每个 GPU 使用全局平均梯度更新本地模型参数。

GPipe 气泡比例公式

假设每个 micro-batch 的前向和反向传播时间均为 1 单位，流水线深度为 d，micro-batch 数量为 m，则 GPipe 的气泡比例为：

当 micro-batch 数量 m 远大于流水线深度 d 时(m≫d)，气泡比例趋近于 0，流水线效率接近线性加速。GPipe 论文中指出，当 m>4d 时，气泡开销几乎可以忽略不计(在激活重计算的情况下)。因此有以下好处：

减少气泡: GPipe 通过 micro-batch 划分和流水线调度，显著减少了朴素流水线并行的气泡问题，提高了 GPU 利用率和训练效率。

同步梯度聚合: GPipe 采用同步梯度聚合，保证了训练过程的同步性，模型收敛性较好。

线性加速潜力: 在 micro-batch 数量足够大的情况下，GPipe 可以实现接近线性的加速效果。

PipeDream(1F1B)

Fig. 7. Illustration of 1F1B microbatch scheduling in PipeDream. (Image source: Harlap et al. 2018) — **Fig. 7. Illustration of 1F1B microbatch scheduling in PipeDream. (Image source:** **Harlap et al. 2018**)

PipeDream(Harlap et al. 2018)是另一种高效的流水线并行训练系统，它采用了 1F1B(1-Forward-1-Backward) 调度策略，并引入了权重暂存(Weight Stashing) 技术，进一步减少气泡，提高流水线效率，并解决 1F1B 调度可能导致的权重版本不一致问题。

PipeDream 的 1F1B 调度策略的核心思想是，每个 GPU(Stage) 交替执行前向传播和反向传播，尽可能地并行工作，减少 GPU 空闲时间。具体流程如下：

Micro-batch 划分: 将一个 mini-batch 划分为 m 个 micro-batch。

流水线阶段划分: 将模型按层划分为 d 个阶段，每个阶段分配到一个 GPU 上。

1F1B 调度执行: 每个 GPU 轮流执行前向传播和反向传播。

权重暂存

由于 1F1B 调度中，前向传播和反向传播可能使用不同版本的模型权重，会导致权重版本不一致问题，影响训练的正确性和收敛性。PipeDream 引入了 权重暂存(Weight Stashing) 技术来解决这个问题。权重暂存的核心思想是，每个 GPU 维护多个版本的模型权重，并确保前向传播和反向传播使用同一版本的权重。

权重暂存实现方式:

版本管理: 每个 GPU 维护一个权重版本队列，存储多个版本的模型权重。

版本选择: 在进行前向传播时，选择当前最新的权重版本。在进行反向传播时，选择与对应前向传播相同的权重版本。

版本更新: 在完成一个 mini-batch 的所有 micro-batch 的反向传播后，更新模型权重，并生成新的权重版本。

为了进一步优化 PipeDream 的内存使用，尤其是在权重暂存方面，PipeDream 衍生出了 PipeDream-flush 和 PipeDream-2BW 两种内存优化变体。

PipeDream-flush

PipeDream-flush 在 PipeDream 的基础上，周期性地进行全局同步的流水线刷新(flush)，类似于 GPipe 的同步梯度聚合。通过定期刷新，PipeDream-flush 可以大幅减少权重暂存所需的内存空间，只需维护单个版本的模型权重，但会牺牲少量吞吐量。

PipeDream-2BW

PipeDream-2BW(Double-Buffered Weights) 维护两个版本的模型权重，即 “双缓冲权重”。它每 k 个 micro-batch 更新一次模型版本，其中 k 大于流水线深度 . 新更新的模型版本不会立即完全替换旧版本，因为可能还有一些剩余的反向传播操作仍然依赖于旧版本。通过双缓冲权重，PipeDream-2BW 可以将权重暂存的内存开销降低到只维护两个版本的模型权重，显著减少内存占用。

PipeDream-2BW 策略有以下优点：

更低的气泡开销: 1F1B 调度策略相比 GPipe 可以进一步减少气泡，提高 GPU 利用率和训练效率。

权重暂存解决版本一致性: 权重暂存技术保证了前向传播和反向传播使用同一版本的权重，解决了 1F1B 调度可能导致的权重版本不一致问题。

内存优化变体: PipeDream-flush 和 PipeDream-2BW 进一步优化了内存使用，降低了权重暂存的内存开销，使得流水线并行更适用于内存受限的场景。

张量并行

张量并行(Tensor Parallelism, TP) 是一种将模型中的张量(通常是权重矩阵) 沿着特定维度切分，并将切分后的分片分配到不同的 GPU 上进行计算的并行方式。张量并行有以下几点优势：

突破单 GPU 显存限制: 张量并行可以将模型参数分散存储在多个 GPU 上，突破单 GPU 显存容量限制，支持训练更大规模的模型。

层内并行: 张量并行可以实现模型层内部的并行化，例如矩阵乘法操作的并行计算，提高计算效率。

与数据并行和流水线并行结合: 张量并行可以与数据并行和流水线并行等其他并行技术结合使用，形成多维并行策略，进一步提高训练效率和可扩展性。

Megatron-LM

Megatron-LM(Shoeybi et al. 2019) 是 NVIDIA 提出的一个用于训练超大型语言模型的系统，它采用了张量并行技术，将 Transformer 模型层内部的矩阵乘法操作并行化，包括 self-attention 和 MLP 中的矩阵乘法。详细信息见

Transformer 的 MLP 层通常包含两个线性层，第一个线性层的计算可表示为 Y=GeLU(XA)，其中 X 是输入矩阵，A 是权重矩阵，GeLU 是激活函数。Megatron-LM 将权重矩阵 A 沿着列维度切分为 P 个分片，其中 P 是 GPU 的数量。每个 GPU 负责存储和计算权重分片。

MLP 层张量并行计算流程:

权重分片: 将权重矩阵沿着列维度切分为个分片，并将分片分配到 GPU 上。

局部矩阵乘法: 每个 GPU 使用输入矩阵 X 和权重分片进行矩阵乘法计算，得到局部输出。

全局拼接(All-Gather): 所有 GPU 通过 All-Gather 操作，将局部输出拼接成完整的输出矩阵。

PTD-P

PTD-P(Pipeline, Tensor, and Data Parallelism)(Narayanan et al. 2021)是一个结合了流水线并行、张量并行和数据并行的多维并行策略。PTD-P 旨在充分利用各种并行技术的优势，提高超大型模型训练的效率和可扩展性。

PTD-P 的特点:

多维并行结合: PTD-P 同时使用了流水线并行、张量并行和数据并行三种并行技术，可以从多个维度对训练过程进行并行化。

Interleaved 1F1B 调度: PTD-P 采用了 interleaved 1F1B 调度策略，与传统的流水线并行不同，它将模型划分为多个不连续的层块(model chunk)，并将多个层块分配给每个 GPU。这种调度策略可以进一步减少气泡，提高流水线效率。

灵活的并行配置: PTD-P 允许用户根据模型结构和硬件资源灵活配置各种并行技术的组合方式，例如可以只使用张量并行和数据并行，也可以同时使用流水线并行、张量并行和数据并行。

传统的流水线并行通常将模型划分为连续的层块，每个 GPU 负责一个连续的层块。PTD-P 的 interleaved 1F1B 调度则将模型划分为多个不连续的层块，例如，GPU 1 负责层 1, 2, 9, 10，GPU 2 负责层 3, 4, 11, 12，依此类推。每个 GPU 负责多个不连续的层块，可以更有效地利用 GPU 资源，减少气泡开销。

Fig. 11.(Top) Default 1F1B pipeline schedule as in PipeDream-flush.(Bottom) Interleaved 1F1B pipeline schedule. First model chunks are in dark colors and second chunks are in light colors. (Image source: Narayanan et al. 2021) — Fig. 11.(Top) Default 1F1B pipeline schedule as in PipeDream-flush.(Bottom) Interleaved 1F1B pipeline schedule. First model chunks are in dark colors and second chunks are in light colors. (Image source: **Narayanan et al. 2021**)

专家并行

详情见

序列并行

序列并行(Sequence Parallelism, SP) 是针对长序列模型(如 Transformer)提出的一种并行化策略，通过在序列维度上对输入进行划分，大幅降低激活内存占用并提高训练效率。它常与数据并行、张量并行或流水线并行结合使用，尤其适合处理超长文本或其他序列数据。

Colossal-AI 序列并行

Fig. 17. The overall architecture of the proposed sequence parallelism and existing parallel approaches. For sequence parallelism, Device 1 and Device 2 share the same trainable parameters. (Image source: Li, et al. 2021) — Fig. 17. The overall architecture of the proposed sequence parallelism and existing parallel approaches. For sequence parallelism, Device 1 and Device 2 share the same trainable parameters. (Image source: **Li, et al. 2021**)

自注意力(self-attention) 的计算复杂度和内存开销与序列长度 s 的平方成正比，长序列数据将增加中间 activation 内存使用量，从而限制设备的训练能力。Colossal-AI 序列并行(Li, et al. 2021)从系统角度提出拆分超长序列到多卡，具体的解决步骤如下。

序列分块

将输入序列划分为若干块，每个块由不同 GPU 保存和计算；因此每张卡只需存储自己对应的序列块激活，避免单卡内存爆炸。

环状通信 + 自注意力

提出环自注意力(Ring Self-Attention, RSA) 机制：各 GPU 先本地计算局部注意力，然后依次向相邻 GPU 传递(环状结构)Key/Value 分块，多轮迭代后保证每个 GPU 能获取全局序列信息。

与其他并行方式结合

不受注意力头数、层数等超参数限制，可配合数据并行、张量并行、流水线并行等技术，共同突破大规模模型的序列长度限制。

Fig. 18. Ring Self-Attention. (Image source: Li, et al. 2021) — **Fig. 18. Ring Self-Attention. (Image source:** **Li, et al. 2021**)

Megatron-LM 序列并行

Megatron-LM(Shoeybi et al. 2019) 原本使用张量并行分担部分激活值，但 Transformer 中的 LayerNorm、Dropout 等操作的激活值仍需完整保存在单卡，显存消耗依旧庞大。因此 NVIDIA 提出 Megatron-LM 序列并行(Korthikanti, et al. 2022)在序列维度对这些激活值进行切分，大幅降低占用。

Fig. 19. Transformer layer with tensor and sequence parallelism. (Image source: Korthikanti, et al. 2022) — **Fig. 19. Transformer layer with tensor and sequence parallelism. (Image source:** **Korthikanti, et al. 2022**)

Fig. 20. MLP layer with tensor and sequence parallelism. (Image source: Korthikanti, et al. 2022) — **Fig. 20. MLP layer with tensor and sequence parallelism. (Image source:** **Korthikanti, et al. 2022**)

序列维度切分 针对 LayerNorm、Dropout 等难以在张量维度切分的激活，将其沿序列维度划分，使每个 GPU 只处理一部分序列的非线性操作。

张量并行仍保留 注意力(Attention)、MLP 等线性操作继续使用张量并行；序列并行的激活需要在前后进行对应的 All-Gather 或 Reduce-Scatter 以交换数据。

选择性激活重计算(Selective Activation Recomputation) 针对部分计算量小但激活量大的操作，选择在反向传播时临时重算，以进一步节省显存。

DeepSpeed-Ulysses 序列并行

DeepSpeed-Ulysses(Jacobs et al. 2023) 针对超长序列训练提出了一种高效的序列并行方案，通过在序列维度对输入进行划分，并结合两阶段的全对全通信，有效降低通信量和激活内存，从而支持训练百万 token 级别的长序列 Transformer 模型。

Fig. 21. DeepSpeed sequence parallelism(DeepSpeed-Ulysses) design. (Image source: Jacobs et al. 2023) — **Fig. 21. DeepSpeed sequence parallelism(DeepSpeed-Ulysses) design. (Image source:** **Jacobs et al. 2023**)

序列划分 + 全对全通信

将输入序列沿序列维度划分到 P 张 GPU 上，每个 GPU 只处理局部的序列；在注意力计算前，通过 All-to-All 操作交换查询()、键()和值()，使得每个 GPU 获得完整序列信息，但仅计算分配到的注意力头。

双阶段通信优化

第一次 All-to-All： 在注意力计算前对进行全对全交换，分散激活计算并降低每卡内存压力；

第二次 All-to-All： 在注意力计算后收集输出上下文，将其重新映射为局部序列分区，既恢复原始序列结构，又显著减少了通信数据量。

高效通信与通用性

利用全对全通信，使得通信量降为，相比传统的 All-Gather 方法(通信量 )节省了近 P 倍的带宽；同时，该方案适用于密集和稀疏注意力，并可与 ZeRO-3 内存优化无缝集成，从而支持更大模型和更长序列的高效训练。

Fig. 22. DeepSpeed-Ulysses vs Megatron LM. (Image source: DeepSpeed Blogs) — **Fig. 22. DeepSpeed-Ulysses vs Megatron LM. (Image source:** **DeepSpeed Blogs**)

在 64 卡 A100 环境下，吞吐量较 Megatron-LM 序列并行最高提升 2.5 倍，并可处理更长序列(百万级 token)；

收敛性能与原模型无差别，可轻松集成到 Megatron-DeepSpeed 框架。

多维度并行

多维度并行(Multi-dimensional Parallelism) 是指在分布式训练中将数据并行、模型并行和流水线并行等多种并行技术有机结合，以充分利用现代 GPU 集群的计算资源。通过这种“3D 并行”或“4D 并行”策略，不仅能提高内存效率，还能提升计算效率，从而实现超大规模(甚至万亿参数级别)模型的高效训练。

3D 并行

随着 GPU 集群计算能力的迅速提升，训练万亿参数级别的模型不再遥不可及。DeepSpeed 将数据并行、模型并行与流水线并行三种技术融合，构建了一种“3D 并行”策略。该策略主要解决训练超大模型所面临的两大挑战：

内存效率：

模型层被划分到不同的流水线阶段，每个阶段内部又通过模型并行进一步分割，减少了模型、优化器和激活值占用的内存量。但需要注意，模型分割不能无限制进行，否则通信开销会显著增加，进而影响计算效率。

计算效率：

为了让计算工作者数量超越单纯模型和流水线并行的限制，同时保证计算效率，DeepSpeed 借助 ZeRO-DP(基于优化器状态分片的数据并行)进行扩展。ZeRO-DP 不仅进一步优化内存使用，还通过拓扑感知映射将数据并行组分配到局部高带宽通信的设备上，极大降低了通信开销。

下面的图示展示了 3D 并行的整体策略：

Fig. 24. Example 3D parallelism with 32 workers. Layers of the neural network are divided among four pipeline stages. Layers within each pipeline stage are further partitioned among four model parallel workers. Lastly, each pipeline is replicated across two data parallel instances, and ZeRO partitions the optimizer states across the data parallel replicas. (Image source: Majumder et al. 2020) — Fig. 24. Example 3D parallelism with 32 workers. Layers of the neural network are divided among four pipeline stages. Layers within each pipeline stage are further partitioned among four model parallel workers. Lastly, each pipeline is replicated across two data parallel instances, and ZeRO partitions the optimizer states across the data parallel replicas. (Image source: **Majumder et al. 2020**)

每个并行维度(数据、模型、流水线)均经过精心映射，以充分利用节点内和节点间的通信带宽。具体策略包括：

优化节点内通信： 由于模型并行的通信开销最大，优先将模型并行组安排在同一节点内，以利用较高的节点内带宽(例如采用 NVIDIA Megatron-LM 的张量切分方式)；

数据并行与流水线并行： 当模型并行不覆盖整个节点时，数据并行组尽可能安排在同一节点内；而流水线并行由于通信量较小，可灵活安排跨节点调度。

通过减少每个数据并行组中通信数据量以及提高局部并行通信的并行度，整体通信带宽得到有效放大。

Fig. 25. Mapping of workers in Figure 24 to GPUs on a system with eight nodes, each with four GPUs. Coloring denotes GPUs on the same node. (Image source: Majumder et al. 2020) — **Fig. 25. Mapping of workers in Figure 24 to GPUs on a system with eight nodes, each with four GPUs. Coloring denotes GPUs on the same node. (Image source:** **Majumder et al. 2020**)

4D 并行

为了进一步扩展模型规模，Llama3(Grattafiori et al. 2024) 训练的时候采用了 4D 并行，它结合了四种并行方法，将模型进行更细粒度的分片，使每个 GPU 上的模型参数、优化器状态、梯度和激活值均能适配高带宽内存(HBM)的容量限制。这四种并行方法分别是：

张量并行(Tensor Parallelism, TP)： 将单个权重张量划分为多个块，分布在不同设备上；

流水线并行(Pipeline Parallelism, PP)： 将模型垂直划分为多个阶段，各阶段在不同设备上并行处理不同微批次；

上下文并行(Context Parallelism, CP)： 将输入上下文划分成多个段，从而缓解长序列输入时的内存瓶颈；

数据并行(Data Parallelism, DP)，通常采用完全分片的数据并行(FSDP)： 对模型、优化器状态和梯度进行分片，并在每个训练步后同步。

下面的图示展示了 4D 并行在 16 个 GPU 上的实现示例，各 GPU 的位置用向量 [D1, D2, D3, D4] 表示，其中每个维度对应一种并行策略。GPU 按照 [TP, CP, PP, DP] 四个维度分组，每个维度的组内大小均为 2。例如，GPU0 和 GPU1 属于同一张量并行组；GPU0 和 GPU2 属于同一上下文并行组；GPU0 和 GPU4 属于同一流水线并行组；GPU0 和 GPU8 属于同一数据并行组：

Fig. 26. Illustration of 4D parallelism. (Image source: Grattafiori et al. 2024) — **Fig. 26. Illustration of 4D parallelism. (Image source:** **Grattafiori et al. 2024**)

通过 4D 并行策略，Llama3 在训练时能够充分利用多个 GPU 的计算资源，同时有效减少内存占用，支持训练超大规模的模型。

业界大模型混合并行策略

模型	DP	TP	PP	ZeRO Stage	GPU类型	GPU数量	计算精度
Bloom-176B	8	4	12	ZeRO-1	A100 80GB	384	BF16
CodeGeeX-13B	192	8	-	ZeRO-2	Ascend 910 32GB	1,536	FP16
GLM-130B	24	4	8	ZeRO-1	A100 40G	768	FP16
OPT-175B	124	8	-	-	80GB A100	992	FP16
Megatron-Turing NLG-530B	16	8	35	N/A	A100 80G	4480	BF16
GPT-NeoX-20B	12	2	4	ZeRO-1	A100 40G	96	FP16

CodeGeeX(13B)

CodeGeeX 采用华为 MindSpore 框架实现，在鹏城实验室"鹏城云脑II"中的192个节点（共1536个国产昇腾910 AI处理器）上训练而成。CodeGeeX 历时两个月在20多种编程语言的代码语料库（> 8500 亿 Token）上预训练得到。

CodeGeeX采用8路模型并行组和192路数据并行组进行混合并行训练；同时，启用 ZeRO-2 来进一步减少优化器状态的内存消耗。

GPT-NeoX（20B）

在12 台 Supermicro AS-4124GO-NART 服务器上进行训练；其中，每台服务器配备 8 个 NVIDIA A100-SXM4-40GB GPU，并配置了两个 AMD EPYC 7532 CPU。所有 GPU 都可以通过用于 GPUDirect RDMA 的四个 ConnectX-6 HCA 之一直接访问 InfiniBand 交换结构（switched fabric）。两台 NVIDIA MQM8700-HS2R 交换机（通过 16 个链路连接）构成了该 InfiniBand 网络的主干，每个节点的 CPU 插槽有一个链路连接到每个交换机。每个训练节点的架构图如下所示：

GPT-NeoX-20B 采用了数据并行、流水线并行和张量并行相结合的方式进行训练。

GLM(130B)

GLM-130B 的预训练持续了 60 天，使用 96 个 DGX-A100（40G）节点，共 768 张 GPU 卡。采用了流水线模型并行与张量并行、数据并行策略相结合的方式，形成 3D并行策略。

为了进一步减少流水线引入的气泡，利用 DeepSpeed 的 PipeDream-Flush 实现来训练具有相对较大的全局批量大小 (4,224) 的 GLM-130B，以减少时间和 GPU 内存浪费。通过数值和实证检验，采用4路张量并行组和8路流水线并行组，达到每张 GPU（40G）135 TFLOP/s。

OPT(175B)

OPT-175B 是 Meta AI 开源的一个拥有 1750 亿参数的语言模型，利用完全分片数据并行（FSDP）与 Megatron-LM 张量并行（8路组）在 992 个 80GB A100 GPU 上训练了 OPT-175B。训练数据包含180B个token，对应800GB的数据，持续训练了约33天。

Bloom（176B）

在 384 张 NVIDIA A100 80GB GPU (48 个节点) 上使用了 3D 并行（数据并行、流水线并行、张量并行）策略，针对 350B 个Token 训练了大约 3.5 个月。

Megatron-Turing NLG（530B）

在训练时，每个模型副本跨越 280 个 NVIDIA A100 GPU，节点内采用Megatron-LM 的 8 路张量并行组，节点间采用 35 路流水线并行组。整个训练过程一共使用了 4480 块英伟达 A100 GPU，在 2700 亿个 Token 上面训练。

参考链接

训练大模型并行和内存优化技术 https://syhya.github.io/zh/posts/2025-03-01-train-llm/