跳转至

AI Training Parallelism

导言

  • 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

  • 思想： 通过将模型切分成不同的部分分别在多个设备上进行计算，从而使得其可以部署更大的模型。
  • 难点：
    • 需要人为的切分设计。
    • 也有必要的数据传输。
    • 切分的部分能有线性加速比，
  • 优点：
    • 适用于模型较大无法单独放入一个设备内存的情况，通过将模型切分到多个设备上进行计算，从而使得其可以部署更大的模型。

张量并行 TP

Tensor Parallelism 是模型并行的一种
  • 思想：通过在多个计算设备上分片模型参数来实现并行。
  • 举例：一个操作中进行并行计算，主要是矩阵-矩阵乘法。张量并行训练是将一个张量沿特定维度分成 N 块，每个设备只持有整个张量的 1/N，同时不影响计算图的正确性。这需要额外的通信来确保结果的正确性。
  • 底层逻辑是矩阵乘法的拆分计算, 矩阵乘法中列并行与行并行这两种张量并行的方式，以及它们在前向传播和反向传播中的区别。14
  • 难点：额外的通信
  • 优点：每个设备只持有整个张量的 1/N

张量并行

张量并行要具体场景具体设计

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

LLM11

序列并行 SP

Sequence Parallelism 不是一种单独算法

SP 只说明输入或激活沿序列维切分，没有说明 Attention 如何取得其他分片的上下文。真正决定显存、通信、头数限制和迁移成本的是 Attention 内部选择 AllGather、All-to-All、P2P Ring，还是它们的二维组合。

两条技术路线

同一个名称下实际有两条不同路线：

  1. TP 配套的 Megatron SP：只把 LayerNorm、Dropout 等逐 token 激活沿序列维切开，Attention 和 MLP 仍由 TP 负责。它的目标是消除 TP 区域之间的冗余激活，不能单独解决长上下文 Attention。8
  2. 长上下文 SP/CP：网络输入和各层激活始终保持 N/P 的序列分片；Attention 再通过 AllGather、All-to-All 或 Ring 取得全局上下文。Ulysses、Ring Attention、Megatron Context Parallelism（CP）和 USP 都属于这条路线。

三个容易混淆的名称

  • Megatron SP 是 TP 配套的激活优化。
  • Megatron CP 是贯穿所有层的长上下文切分，Attention 默认以 Ring 交换 KV。
  • USP 在原论文中指 Unified Sequence Parallelism，即 Ulysses × Ring 的二维组合；部分框架也把 Ulysses Sequence Parallel 简写成 USP，阅读配置时需要回到具体通信原语。

长上下文序列并行方案选择示意图

自绘示意图：AllGather、Ulysses、Ring 与 USP 的核心区别分别是整段复制、按头换轴、KV 分块环传和二维组合。

演进脉络

方案 出现时间 当时的设计初衷 核心数据移动 论文或官方效果
Ring Self-Attention（早期 SP） 2021-05 单卡不再持有完整序列，突破自注意力的长序列显存上限 固定本地 Q，K/V 分块沿 Ring 传递 论文在 64 张 P100 上相对 TP 支持 13.7× 最大 batch、3.0× 最大序列；稀疏 Attention 超过 114K token。5
Megatron SP 2022-05 减少 TP 区域之间 LayerNorm、Dropout 激活的重复存储，少做全量重计算 TP 的 AllReduce 拆成 ReduceScatter + AllGather 与选择性重计算合用时，论文报告激活显存减少 5×；530B GPT 训练达到 54.2% MFU，比全量重计算的 42.1% 快 29%。这些数字不是 SP 单项消融。8
DeepSpeed-Ulysses 2023-09 用高带宽集合通信替代低效 P2P，把长序列 Attention 包装成可复用模块 Attention 前后各一次 All-to-All，在序列维与头维之间换轴 官方实验报告比既有基线支持 4× 更长序列、通信量降低超过 10×、吞吐最高 2.5×，并训练百万 token 序列。1023
Ring Attention 2023-10 用 blockwise Attention 把 KV 通信隐藏在计算后面，并让可处理长度随设备数扩展 每轮只收发一个 KV block，在线合并 softmax 统计量 论文支持百万 token，并给出最多随设备数线性扩展的上下文长度；“无额外暴露通信”依赖 block 计算足以覆盖 P2P。9
Megatron CP 2023 Q3；MCore 0.5.0+ 把 Ring Attention 工程化到 Megatron/Transformer Engine，兼顾因果负载均衡与现有 Flash/Fused Attention 默认 P2P Ring；当前也支持 AllGather、All-to-All 和分层 A2A+P2P 官方文档给出每卡激活显存约按 CP 度缩小，并在 175B GPT 示例中优于全量重计算；页面未给出可脱离图表复述的统一加速数字。2024
USP（Unified SP） 2024-05 不在 Ulysses 与 Ring 间二选一，而是适配头数和节点内/节点间拓扑 二维进程网格：一维 All-to-All，另一维 P2P Ring 论文在 2×8 A800 上训练 LLaMA3-8B：208K 序列达到 147.26 TFLOPS/GPU、47% MFU；64K/80K 时最佳二维配置比纯 Ring 高 13%/12%。18
AllGather CP 2024-10 进入 MCore cp_comm_type 为 CP 提供最直接、最容易复用标准 Attention kernel 的通信后端 前向 AllGather 完整 K/V；反向 ReduceScatter K/V 梯度 它不是独立论文方案，MCore 明确标注当前 AllGather 不能异步、不能与 Attention 计算重叠；没有独立的统一性能数字。2122

效果数字的边界

上表来自不同年份、模型、GPU、网络、精度和 baseline，只能证明各方案在原实验中的可行性，不能横向排名。同一集群上的最终选择应固定模型、序列长度、CP 度和重计算策略，比较峰值显存、step time、MFU 与通信暴露时间。

通信机制

1. Ulysses：用头维换回完整序列。 每张卡先持有 [N/P, H] 的 Q/K/V；All-to-All 后改为 [N, H/P]，因此本地仍可直接调用 FlashAttention，输出再通过一次 All-to-All 回到序列分片。其优势是集合通信带宽利用率高、Attention kernel 保持大块计算、代码可包装成 DistributedAttention。其硬约束是 Ulysses 度不能超过可切分头数；GQA 通常受 KV head 数限制，MQA 的单 KV head 尤其不适合纯 Ulysses。19

2. Ring Attention：固定本地 Q，流动 K/V。 每张卡在 P-1 轮中接收一个远端 KV block，计算局部 QKᵀ，并用 online softmax 的最大值与归一化因子合并局部结果。它不按头切分，因此没有 Ulysses 的头数上限，也不需要在 Attention 内物化完整 KV；代价是更小的矩阵块、P2P 带宽利用率和因果负载均衡会决定实际效率。

3. Megatron CP：Ring 的框架化实现。 CP 让 Q、MLP、Norm 等都维持 N/P，只有 Attention 需要跨 rank 的 KV。Megatron 默认的 P2P CP 采用 Ring，并通过负载均衡避免因果下三角导致部分 rank 空转；MQA/GQA 的 KV head 更少，反而可以降低通信量。它能与 TP、PP、DP 组成独立维度。20

4. USP：把 Ulysses 和 Ring 放到二维网格。 设总 SP 度为 P = P_u × P_r：节点内或高带宽域使用 P_u 路 Ulysses，跨节点使用 P_r 路 Ring。于是只有 P_u 受头数约束，总 SP 度仍可通过增大 P_r 扩展。这个设计比纯 Ring 更容易保留大块 Attention 的效率，也比纯 Ulysses 更适合 GQA/MQA 和分层网络。1819

USP 论文中的 Ulysses 与 Ring Attention 机制

来自 USP 论文 Figure 2：左侧 All-to-All 把序列分片转换为头分片；右侧 Ring 让 KV block 在设备间逐轮流动。

5. AllGather CP：先复制完整 KV，再调用标准 Attention。 它保留本地 Q_local，一次性聚合 K_full/V_full：

# 前向机制示意；具体 API 取决于框架
k_full = all_gather(k_local, dim="sequence")
v_full = all_gather(v_local, dim="sequence")
o_local = attention(q_local, k_full, v_full, global_mask_for_local_q)
# 反向通过 ReduceScatter 把 K/V 梯度归并回所属序列分片

这种写法最容易复用现有 MHA、GQA、MQA 或自定义 Attention kernel，且不受头数限制；但每张卡在 Attention 时都临时持有完整 K/V，AllGather 还不能与计算重叠。因此它更适合 CP 度较小、先打通正确性、Ring kernel 尚未适配，或本地 Attention kernel 明显更高效 的场景，不适合作为超长序列的大规模默认方案。MCore 的 Python 配置名是 cp_comm_type="all_gather"，命令行枚举写作 allgather。21

USP 在 208K 序列长度下的性能上界实验

来自 USP 论文 Table 7：2×8 A800、LLaMA3-8B 的 TP × Ulysses × Ring 组合；208K 序列达到 147.26 TFLOPS/GPU、47% MFU。

工程联评

下表中的迁移成本与可读性是基于数据布局、通信组、mask/position 处理和 kernel 依赖做出的工程判断，不是论文测量值。

方案 头数与拓扑限制 代码复杂度与可读性 新框架迁移成本 新模型是否重新适配 主要适用边界
Megatron SP 必须依附 TP；不负责 Attention 全局上下文 框架内高可读；通信藏在 TP Linear 边界 MCore 内低，跨框架中 通常不改 Attention；需保证 Norm、Dropout、残差和 RNG 的序列分片语义 用于 TP 激活优化，可与 CP 叠加；不能替代长上下文 CP
Ulysses P_u 受可切分头数，GQA 常受 KV heads 限制；偏好高带宽 All-to-All 中等；Attention 前后各一个换轴 中等 标准 MHA 较容易；GQA/MQA、TP 共存、packed sequence、cross-attention 需重新核对布局 头数充足、节点内高带宽、希望复用现有 FlashAttention
Ring Attention 无头数上限；依赖 P2P 与足够大的计算块覆盖通信 较高；分块调度、online softmax 与反向较难读 高 新 mask、RoPE、变长序列、cross-attention 或特殊 Attention kernel 通常需适配 GQA/MQA、SP 度大于头数、显存优先的超长序列
Megatron CP（P2P） 无头数上限；当前高性能路径依赖 Transformer Engine/支持的 Attention 应用层低，后端高 MCore GPT/LLaMA 路径低；移植到新框架高 共享 GPT 路径的模型多为配置接入；新 Attention 类型仍需实现 CP 语义 NVIDIA/Megatron 训练栈的默认长上下文方案
USP 仅 P_u 受头数限制；可按 NVLink/RDMA 拆二维组 较高；两套通信组和布局转换降低直观性 中高 标准 Attention 可集中在 wrapper；特殊 mask、GQA 映射和训练框架 group 管理仍需适配 多节点分层网络、GQA/MQA、既要集合通信效率又要扩展 SP 度
AllGather CP 无头数限制；要求每卡能临时容纳完整 K/V 最高可读性，最接近单卡 Attention 低到中 通常只需全局 KV、局部 Q 的 mask/position 与反向归并；新 kernel 适配最少 正确性基线、小 CP 度、特殊 Attention 快速接入；超长序列受 KV 峰值显存限制

选型建议

  1. 先区分目标：只想减少 TP 的冗余激活，使用 Megatron SP；真正被上下文长度卡住，才增加 CP/SP Attention。
  2. Megatron + Transformer Engine 的常规模型：优先从 P2P CP 开始；AllGather CP 用作正确性对照、较小 CP 度或尚未支持 Ring 的 Attention 类型。
  3. 头数充足且通信域带宽高：纯 Ulysses 通常代码更薄、kernel 粒度更大；但需要验证 P_u 能整除 Q/KV heads，并检查 TP 是否也在切头维。
  4. GQA/MQA 或总 SP 度超过 KV head 数：优先 Ring 或 USP；多节点时可把 Ulysses 放在节点内、Ring 放在节点间，再用 profile 选择二维比例。
  5. 不要只看最长序列：固定相同模型与 batch 语义，同时记录峰值显存、吞吐/MFU、通信暴露、因果负载均衡、数值一致性，以及变长/packed sequence 的有效 token 比例。

流水线并行 PP

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

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

层间拆分来减小参数压力的思想

GPipe

简单流水线GPipe

3

1F1B

1F1B

PipeDream

Generalized流水线PipeDream

4

VPP

  • 假定当前模型网络共16层（编号 0-15），4个Device，
  • 前述GPipe模式和PipeDream是分成4个stage， 按编号0-3层放Device1，4-7层放Device2，并以此类推。
  • virtual pipeline则是按照文中提出virtual_pipeline_stage概念减小切分粒度，
  • 以virtaul_pipeline_stage=2为例，将0-1层放Device1，2-3层放在Device2，...，6-7层放到Device4，8-9层继续放在Device1，10-11层放在Device2，...，14-15层放在Device4。

VPP

DualPipe

idea

传统流水线并行卡之间串行，效率低。可以0号和7号卡同时都是layer0和layer7,这样做的目的：真个流水线能同时运行两批次的数据！比如第1批从gpu0卡的layer0开始forward，第10批数据从gpu7卡的layer0开始forward，提升效率！

DualPipe 是一种创新的双向管道并行算法，在 DeepSeek-V3 技术报告中提出。实现了正向和反向计算-通信阶段的完全重叠，同时也减少了管道气泡时间。

混合并行

  • 2021年10月，微软和英伟达联合提出了 PTD-P(Inter-node Pipeline Parallelism, Intra-node Tensor Parallelism, and Data Parallelism)训练加速方法，
  • 通过数据并行、张量并行和 Pipeline 并行“三管齐下”的方式，将模型的吞吐量提高 10%以上。
  • 该并行方法可以在3072个GPU 上，以502P的算力对一万亿参数的GPT 架构模型进行训练，实现单GPU吞吐量52%的性能提升。

Picture curtesy of Hugging Face

Bloom-176B

DP8，TP4, PP12：

megatron 并行默认维度

采用的顺序是tp-cp-ep-dp-pp，我们认为越靠前的并行组，通讯量越大，所以尽量安排在一台机器内的更近的维度。12

专家并行 MoE(EP)

MOE的概念

  • 思想：一种基于稀疏 MoE(Mixture-of-Experts) 层的深度学习模型架构被提出，即将大模型拆分成多个小模型(专家，expert)， 每轮迭代根据样本决定激活一部分专家用于计算，
  • 优点：只计算一部分，达到了节省计算资源的效果；
  • 实现：MoE 将模型的某一层扩展为多个具有相同结构的专家网络(expert)，并由门(gate)网络决定激活哪些 expert 用于计算，从而实现超大规模稀疏模型的训练。

将中间层扩展为具有n个expert的MoE结构

MOE 与 EP的关联

17

  • MOE层如果不开EP， MoE的结构由多个 expert 构成，每次只选择一个专家执行；
  • MOE层开启EP，MOE层的多个专家就能同时并行，处理不同的tokens，从而达到并行的效果。为了实现这点需要对topk的token来dispatch和combine

EP 与 TP的区别

  • TP开启时：每个 EP rank 上只包含一部分 expert，而每个 EP rank 上的 token（即 token 对应的 hidden state） 会根据 gating 结果分发到其他 EP rank 上的 expert。这个过程通过 all-to-all 通信完成。
  • 只开启EP时，GPU会拥有多个完整的专家，每个矩阵计算都是完整的，是大矩阵计算。
  • 不同情况各有优劣。

EP

红色和绿色方块表示非-MOE层，e0、e1表示MOE层（总共4n个专家）

一共16块GPU：

  • ep_world_size = 4：表示我们希望用4块GPU装下一套完整的专家。确定这个数值后，我们就能确认ep_groups
  • local_expert_num：expert_num / ep_world_size，其中expert_num表示每层专家的总数。
    • 假设每层专家数量是4，那么1块gpu上就放一个专家；
    • 假设每层专家数量是8，那么1块gpu上就放2个专家。
    • 所以图中的e0等符号并不绝对表示这里只有1个专家，只是对local_expert的统称。
  • ep_dp_world_size：类比于non-MoE层，MoE层同样也有数据并行的概念。例如图中[g0, g4, g8, g12]上都维护着e0，所以它们构成一个ep_dp_group。这个group的作用是当我们在计算bwd时，它们之间是需要做梯度的allreduce通讯的，我们会在下文详细图解这一点。另外需要注意的是，构成ep_dp_group的条件不仅是e相同，还需要每个e吃的batch的数据不同（类比于一个普通的dp_group，组内的每张卡吃的是不同的小batch）。现在你可能无法具象化感受这点，我们在后文将ep+tp+dp并行的时候再细说。
  • ep_tp_world_size：类比于non-MoE层，MoE层同样也有张量并行的概念，即一个专家可以纵向切割成若干份.

代码实现

参考DeepSeek-VL2实现， 当前代码实现有几个特点：

  1. EP=x，就将world_size拆分成几份，每份里的机器上拥有 experts/EP 的专家数的副本。
  2. self.shared_experts = MLP, forward里和路由专家结果相加。

ZeRO

  • ZeRO通过在多个设备上分片优化器状态、梯度和参数来减少每个设备的存储需求。
  • Zero 优化方法有三个层次，分别是 ZeRO-1、ZeRO-2 和 ZeRO-3。13
  • 它们是由微软提出的 ZeRO（Zero Redundancy Optimizer） 优化技术的不同阶段，旨在减少大规模分布式训练中的内存占用。

https://blog.csdn.net/qq_38563206/article/details/133792668

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）在大规模分布式训练中的局限性。
  • 它通过分层通信和计算优化，提升训练效率和扩展性。

核心思想是将计算和通信任务分层处理，减少通信开销，提高资源利用率。具体包括：

  1. 分层通信：将通信任务分为多个层次，优先在低层次（如节点内）完成，减少高层次（如跨节点）的通信。
  2. 分层计算：将计算任务分层处理，优先在低层次完成，减少高层次的计算负担。

主要特点

  1. 减少通信开销：通过分层通信，降低跨节点通信频率，提升效率。
  2. 提高资源利用率：分层计算使资源分配更合理，减少闲置。
  3. 增强扩展性：优化后的方法更适合大规模分布式训练，支持更多计算节点。

异构系统的并行

人们思考为什么 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 ↩

  10. DEEPSPEED ULYSSES: SYSTEM OPTIMIZATIONS FOR ENABLING TRAINING OF EXTREME LONG SEQUENCE TRANSFORMER MODELS ↩

  11. Efficient large-scale language model training on gpu clusters using megatron-lm ↩

  12. https://zhuanlan.zhihu.com/p/5502876106 ↩

  13. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models ↩

  14. 深度学习的分布式训练与集合通信（二） ↩

  15. 深度学习的分布式训练与集合通信（三） ↩

  16. NV ↩

  17. https://developer.nvidia.com/zh-cn/blog/demystifying-ai-inference-deployments-for-trillion-parameter-large-language-models/ ↩

  18. USP: A Unified Sequence Parallelism Approach for Long Context Generative AI ↩↩

  19. YunChang: Unified Sequence Parallel Attention ↩↩

  20. Megatron Core Context Parallel Package ↩↩

  21. Megatron Core TransformerConfig.cp_comm_type ↩↩

  22. Megatron-LM commit: configure per-layer CP communication type ↩

  23. DeepSpeed-Ulysses official release and evaluation ↩

  24. Megatron-LM initial context-parallel process group ↩

评论