为什么要搞分布式训练？一切还得从“太大放不下”说起

本文链接：https://blog.csdn.net/qq_54445177/article/details/147225853

Okay，各位同学，大家好！我是你们的老朋友小雲。最近总有同学问我，现在大模型动不动就几千亿、上万亿参数，我们实验室/公司的卡根本顶不住啊，这到底是怎么训出来的？是不是有什么魔法？

魔法谈不上，但确实有“黑科技”——那就是分布式训练。今天，我用大白话给大家捋一捋大模型分布式训练中常见的几种并行技术，争取让大家不仅听得懂，面试的时候也能跟面试官掰扯几句，听起来有理有据，不露怯！

这篇文章算是咱们分布式训练系列的第一篇，先给大家打个基础，介绍几种主流的并行“武功秘籍”。

为什么要搞分布式训练？一切还得从“太大放不下”说起

咱们先聊聊为啥非得用分布式训练。几年前，大家可能还在玩几亿、几十亿参数的模型，那时候单机单卡或者单机八卡努努力，可能还能应付。但自从Transformer架构（尤其是Attention机制）火了之后，模型参数就像坐上了火箭，噌噌往上涨。再加上后来MoE（Mixture of Experts，混合专家）这种架构的出现，更是把参数量推向了万亿级别。

你想想，一个GPT-3 175B模型，光参数就得几百GB（用float16也得350GB），这还没算训练过程中产生的梯度、优化器状态（比如Adam里的动量和方差）等等。现在主流的GPU显存也就几十GB（比如A100/H100的40GB/80GB），单卡肯定装不下。就算侥幸装下了，训练速度也慢得让人想哭。

这时候，AI集群就成了救世主。简单说，就是把一堆机器（每台机器上可能还有多张GPU卡）连起来，大家一起使劲儿，共同完成一个超大模型的训练任务。怎么让这些机器和GPU高效地协同工作呢？这就需要对计算任务、训练数据、甚至是模型本身进行“分而治之”，也就是我们今天要聊的分布式并行训练技术。

一、数据并行（Data Parallelism, DP）：最简单也最常用的“人海战术”

核心思想： 模型我复制，数据你来分。

数据并行（DP）是最直观、最容易理解的一种并行方式。想象一下，你要训练一个模型，但是数据量太大了，一个人处理不过来。怎么办？找一群人（GPU），每个人都拿到一份完整的操作手册（模型），然后把任务（数据集）分成几份，每人领一份去干活（训练）。

图1：数据并行示意图

具体流程：

分发： 把一个大的Batch数据切分成N个小的mini-batch，每个GPU分到一个mini-batch。
复制： 每个GPU上都保留一份完整的模型副本。
计算： 每个GPU独立地根据自己的mini-batch数据进行前向传播计算损失，然后进行反向传播计算梯度。
同步： 这是关键！在反向传播之后，所有GPU计算出的梯度需要进行一次聚合（Aggregation），通常是做一次All-Reduce操作（所有GPU把自己的梯度发给其他所有GPU，并计算平均值），保证所有GPU上的模型副本在参数更新后仍然保持一致。
更新： 每个GPU使用聚合后的梯度来更新自己的模型参数。

优点：

实现简单，容易理解和上手。PyTorch的DistributedDataParallel (DDP) 就是典型的实现，几行代码就能搞定。
通常能获得不错的加速比，尤其是在网络带宽足够的情况下。
适用于数据量大，但模型可以放入单张 GPU 的情况。

缺点：

内存冗余： 每个GPU都要存一份完整的模型参数、梯度和优化器状态。当模型非常大时，即使一个副本也可能塞不进单张卡的显存，这时候DP就无能为力了。

面试小贴士： DDP是面试常考点。可以提一下它相比于更早的DataParallel(DP) 的优势，比如DDP使用多进程，避免了Python GIL（全局解释器锁）的瓶颈，并且通信效率更高（直接进行All-Reduce，而不是先把梯度汇总到主GPU再分发）。

二、模型并行（Model Parallelism, MP）：模型太大？拆！

当模型大到单卡显存装不下时，数据并行就没辙了。这时候就得请出模型并行（MP）。

核心思想： 模型你来拆，数据我来传。

模型并行，顾名思义，就是把一个大模型拆分成多个部分，分别放到不同的GPU上。这样每个GPU只需要承担模型的一部分，显存压力就小多了。

模型并行主要有两种玩法：张量并行（Tensor Parallelism, TP） 和 流水线并行（Pipeline Parallelism, PP）。你可以粗略地理解为：张量并行是在层内拆，流水线并行是在层间拆。

2.1 张量并行（Tensor Parallelism, TP）：把“大算子”切开算

核心思想： 一个矩阵乘法太大？拆成小块并行算。

张量并行主要针对模型中的单个大运算（比如巨大的nn.Linear层或Attention层里的矩阵乘法）进行拆分。它把一个大的张量（Tensor）沿着某个维度切成N块，每个GPU只持有和计算其中的1/N块。

举个栗子（矩阵乘法）： 假设我们要计算 C = A * B。

按列切分 B： 我们可以把矩阵B沿着列（Column）切成 [B0, B1, ..., Bn-1]。然后把A广播给所有N个GPU，第 i 个GPU计算 Yi = A * Bi。
按行切分 A： 或者，我们可以把矩阵A沿着行（Row）切成 [A0^T, A1^T, ..., An-1^T]^T。然后把B广播给所有N个GPU，第 i 个GPU计算 Ci = Ai * B。最后需要通过All-Gather操作把各个GPU上的 Ci 收集起来，拼接成完整的C。

图2：张量并行（按列切分）示意图

Transformer中的应用： 在Transformer的MLP（多层感知机）部分，通常是先一个nn.Linear（记为A），然后接一个激活函数，再接一个nn.Linear（记为B）。Megatron-LM提出了一种巧妙的1D张量并行方法：

第一个Linear层A按列切分。计算时，输入X需要通过All-Gather在TP组内同步，然后每个GPU计算 X * Ai，得到 Yi。
第二个Linear层B按行切分。输入是Y = [Y0, Y1, ..., Yn-1]，每个GPU持有Yi和Bi。计算 Zi = Yi * Bi。最后，需要对所有GPU的输出 Zi 进行一次All-Reduce，得到最终结果 Z = sum(Zi)。
通过这种方式，两次大的矩阵乘法都被有效地并行化了，并且巧妙地安排了通信操作（一次All-Gather，一次All-Reduce）。

优点：

可以有效降低单张卡上的峰值显存（尤其是激活值的显存）。
允许训练单个层就非常巨大的模型。

缺点：

通信开销大： 需要在每次前向和反向传播过程中进行额外的通信（如All-Gather, All-Reduce）。
实现复杂： 需要深入模型内部，修改算子的实现。
通常局限于节点内： 因为通信量大，对网络带宽要求极高，一般只在节点内部（比如通过NVLink高速互联的GPU）使用。

实现： Megatron-LM (1D TP), Colossal-AI (支持1D, 2D, 2.5D, 3D TP)。

面试小贴士： 面试官可能会问张量并行和数据并行的区别。核心在于切分的对象不同（模型 vs 数据）以及通信模式不同（TP通信更频繁，通常在算子内部）。可以结合Transformer的MLP或Attention解释TP的具体切分和通信过程。

2.2 流水线并行（Pipeline Parallelism, PP）：像工厂流水线一样处理

核心思想： 把模型按层分组，不同组交给不同GPU，串行处理。

流水线并行是把模型的不同**层（Layers）**分配到不同的GPU上。比如一个有40层的模型，用4张卡做流水线并行，可以把0-9层放GPU 0，10-19层放GPU 1，以此类推。

图3：流水线并行示意图

朴素流水线的问题（Naive Pipeline Parallelism）：

最简单的想法是：一个batch的数据先在GPU 0上完成前10层的前向计算，然后把输出（激活值）传给GPU 1；GPU 1计算10-19层，再传给GPU 2…直到最后一个GPU完成计算。反向传播类似地倒着来一遍。

这种方式的问题在于**“流水线气泡（Bubble）”**：在任何一个时间点，只有一个GPU在干活，其他GPU都在等待。这就像一个效率低下的工厂流水线，工人大部分时间在发呆。

图4：朴素流水线并行的“气泡”问题

改进：微批次（Micro-batching）

为了减少气泡，提高设备利用率，现代流水线并行（如GPipe, PipeDream）引入了**微批次（Micro-batch）**的概念。将一个大的训练批次（mini-batch）进一步切分成更小的微批次。然后，让这些微批次像流水一样流过各个GPU阶段（Stage）。

当前一个GPU处理完一个微批次后，立刻将其输出传给下一个GPU，然后自己开始处理下一个微批次。这样，在理想情况下，多个GPU可以同时处理不同的微批次，大大减少了空闲时间。

GPipe vs PipeDream (1F1B):

GPipe: 先把所有微批次的前向传播都做完，然后再做所有微批次的后向传播。这种方式需要缓存每个微批次在前向传播时的激活值，内存开销较大。
PipeDream (及其变种，如1F1B - One Forward, One Backward): 采用更优化的调度策略，例如“1F1B”。一个GPU完成一个微批次的前向计算后，一旦后续阶段也完成了该微批次的前向，并且计算出了梯度，这个梯度就会被立刻传回给当前GPU，使其可以尽快开始计算该微批次的后向传播。这样可以显著减少需要缓存的激活值数量，降低内存压力。

优点：

显著降低单卡内存占用（只存部分层和少量激活值）。
可以训练非常深的模型。
相比TP，对节点间带宽要求稍低。

缺点：

存在流水线气泡，理论加速比不如数据并行。气泡大小和微批次数有关，需要调优。
负载均衡问题：需要尽量让每个阶段的计算量和参数量均衡，否则慢的阶段会成为瓶颈。
实现相对复杂，需要对模型进行切割，处理跨设备的数据传输。

实现： GPipe, PipeDream (及其变种), Megatron-LM, DeepSpeed Pipeline Parallelism。

面试小贴士： 理解“气泡”的成因以及微批次如何缓解气泡是关键。可以对比GPipe和PipeDream（或1F1B）在调度和内存管理上的差异。思考流水线并行如何与数据并行结合（后面会讲）。

三、优化器状态并行 / ZeRO：专治内存“大头”

我们前面提到，训练时的显存占用不仅有模型参数，还有梯度和优化器状态（尤其是Adam这种带一阶、二阶动量的优化器，状态量可能是模型参数的两倍！）。数据并行虽然加速了计算，但这些状态在每个GPU上都存了一份，造成了巨大的内存冗余。

ZeRO（Zero Redundancy Optimizer） 就是来解决这个问题的。

核心思想： 既然冗余，那就分片（Shard），每人只存一部分。

ZeRO本质上是一种增强的数据并行，它的目标是消除数据并行中的内存冗余。它不是复制完整的模型状态（参数、梯度、优化器状态），而是将这些状态分割成N份（N是参与ZeRO的GPU数量），每个GPU只负责存储和更新其中的1/N份。

图5：ZeRO 不同阶段的分片策略

ZeRO的三个阶段：

ZeRO-Stage 1 (ZeRO-1)：优化器状态分片 (Optimizer States Sharding)
- 每个GPU只保存和更新它负责的那部分参数对应的优化器状态。
- 模型参数和梯度在每个GPU上仍然是完整的副本。
- 在优化器更新步骤（optimizer.step()）时，需要一次通信（Reduce-Scatter）将梯度聚合到对应的GPU上，然后每个GPU更新自己负责的那部分参数对应的优化器状态和参数。之后再通过一次通信（All-Gather）将更新后的完整参数同步给所有GPU。
- 效果： 显著减少优化器状态的内存占用（约减少 (N-1)/N）。
ZeRO-Stage 2 (ZeRO-2)：优化器状态 + 梯度分片 (Optimizer States & Gradients Sharding)
- 在ZeRO-1的基础上，梯度也进行分片。
- 在反向传播计算梯度时，每个GPU只保留自己负责的那部分参数对应的梯度，其他参数的梯度计算完后就丢弃。通过一次Reduce-Scatter操作，将梯度直接聚合到目标GPU上。
- 模型参数在每个GPU上仍然是完整的副本（在前向/反向计算时需要）。
- 效果： 进一步减少梯度占用的内存（约减少 (N-1)/N）。优化器更新时的通信模式与ZeRO-1类似。
ZeRO-Stage 3 (ZeRO-3)：优化器状态 + 梯度 + 模型参数分片 (Optimizer States & Gradients & Parameters Sharding)
- 终极形态！把模型参数也分片了。每个GPU平时只持有自己负责的那部分模型参数。
- 关键操作： 在执行前向或反向计算需要用到某个Layer的参数时，才通过All-Gather操作从其他GPU那里把完整的参数临时收集过来。计算完成后，非自己负责的参数就可以丢弃，释放显存。
- 效果： 最大程度地消除了内存冗余，使得N张卡理论上可以训练N倍大的模型（相比单卡）。
- 代价： 通信量大大增加，因为每次需要用到参数时都要进行All-Gather。对网络带宽要求很高。

图6：训练时显存占用构成（混合精度）

ZeRO-Offload： ZeRO还可以结合CPU和NVMe（高速SSD）内存。可以将分片的优化器状态、梯度甚至参数临时“卸载（Offload）”到CPU内存或NVMe硬盘上，进一步释放宝贵的GPU显存，使得在有限的GPU资源下也能训练超大模型。当然，这会带来额外的拷贝开销，牺牲训练速度。

优点：

极大降低了数据并行下的内存冗余，能训练更大的模型或使用更大的Batch Size。
与数据并行的编程模型类似，对用户相对友好（尤其是DeepSpeed等框架封装后）。
ZeRO-3可以看作是一种动态的模型参数重组，非常灵活。

缺点：

增加了通信开销，尤其是ZeRO-3。
需要高效的通信库支持。

实现： DeepSpeed ZeRO (Stages 1, 2, 3, Offload), PyTorch FSDP (Fully Sharded Data Parallel, 类似ZeRO-3)。

面试小贴士： 理解ZeRO三个阶段分别解决了哪部分内存冗余，以及引入了什么样的通信开销是重点。能解释清楚ZeRO-3中参数的动态收集（All-Gather）和释放机制会很加分。知道ZeRO-Offload是利用异构内存进一步扩大模型规模的手段。

四、异构系统并行：榨干CPU和硬盘的价值

前面提到的ZeRO-Offload其实就是异构系统并行的一种体现。

核心思想： GPU显存不够用？把暂时不用的东西（参数、优化器状态、激活值）先挪到内存大得多的CPU或速度还行的NVMe硬盘上。

图7：利用CPU/NVMe进行Offload

这种方式主要是为了突破硬件显存的物理限制，让你有可能在有限的GPU资源（比如单机几张卡）上，也能把一个巨大无比的模型给跑起来（哪怕慢一点）。

优点：

极大扩展了可训练模型的规模上限。

缺点：

速度慢： GPU与CPU之间（PCIe总线）、CPU与NVMe之间的数据传输速度远低于GPU之间（NVLink）或GPU内部的内存访问速度。频繁的Offload/Reload会严重拖慢训练速度。
实现复杂，需要精细管理数据的换入换出。

应用场景： 主要用于显存极度受限，但又必须训练超大模型的场景，或者是一些对训练时间不那么敏感的研究探索。

五、多维混合并行：取长补短，强强联合

单一的并行策略往往有其局限性。比如DP不能解决模型过大的问题，TP通信开销大且受限于节点内，PP有气泡问题。为了训练动辄万亿参数的巨无霸模型，通常需要将多种并行策略组合起来使用，形成多维混合并行。

图8：混合使用数据、流水线、张量并行

常见组合策略：

DP + PP：
- 将模型按层切分到不同的流水线阶段（PP）。
- 每个流水线阶段内部，再使用数据并行（DP）来处理更多数据，加速训练。
- 例如，你有16张卡，可以做4个流水线阶段，每个阶段内用4张卡做数据并行（4x4=16）。
DP + TP：
- 对于模型中特别大的层（比如Attention或MLP），使用张量并行（TP）将其切分到几张卡上。
- 在TP组之外，再使用数据并行（DP）来扩大训练规模。
- 例如，16张卡，可以做4路TP（处理大层），同时做4路DP（4x4=16）。
PP + TP：
- 先将模型按层切分成流水线阶段（PP）。
- 对于每个阶段内部比较大的层，再使用张量并行（TP）进行切分。
- 例如，16张卡，可以做4个流水线阶段，每个阶段内部用4张卡做TP（4x4=16）。这种组合可以处理又深又宽的模型。
DP + PP + TP (3D并行)：
- 终极组合！同时使用数据并行、流水线并行和张量并行。
- 典型部署：
  - 节点内（Intra-node）： 使用TP，因为节点内GPU通过NVLink高速连接，能承受TP的大通信量。
  - 节点间（Inter-node）： 使用PP和DP，因为它们对网络带宽的要求相对较低。
- 例如，你有64张卡，分布在8台机器上（每台8卡）。可以：
  - 设置TP size = 8 (节点内做张量并行)
  - 设置PP size = 4 (跨4台机器做流水线并行)
  - 设置DP size = 2 (剩下2组机器做数据并行)
  - 总卡数 = 8 * 4 * 2 = 64

图9：3D并行与硬件拓扑的匹配

再结合ZeRO： 上述组合还可以和ZeRO（尤其是ZeRO-1，有时也用ZeRO-2）结合使用，进一步优化显存占用。比如在DP组内使用ZeRO。

优点：

集各家之长，可以最大化利用集群资源，训练更大规模的模型，并获得更好的加速效果。

缺点：

极度复杂： 配置和调试非常困难，需要对各种并行策略和硬件拓扑有深入理解。

实现： DeepSpeed (支持灵活的3D并行配置), Megatron-LM (也有自己的实现), Colossal-AI, FairScale等。

面试小贴士： 理解为什么要混合并行（单一策略的局限性）。能解释清楚一种典型的3D并行配置方式（比如节点内TP，节点间PP+DP）及其原因（匹配通信需求和硬件带宽）。

六、自动并行：把复杂性交给框架

手动配置和优化混合并行策略实在是太难了！能不能让框架自动帮我们搞定？自动并行（Automatic Parallelism） 应运而生。

核心思想： 用户只管定义模型和指定硬件资源，框架自动分析计算图、通信开销和内存占用，然后生成一个（近似）最优的并行执行计划。

图10：自动并行的目标

实现方式： 通常涉及复杂的技术，如：

计算图分析： 理解模型的结构和依赖关系。
设备拓扑感知： 了解集群中节点、GPU的连接方式和带宽。
代价模型： 估算不同并行策略下的计算、通信和内存开销。
搜索算法： 在巨大的并行策略空间中搜索最优解（或次优解）。

优点：

极大降低用户门槛： 无需手动切分模型或配置复杂的并行策略。
可能找到更优的策略： 自动化算法可能发现人类工程师难以想到的优化方案。

缺点：

技术仍在发展中： 效果可能不如专家手动调优，有时生成的策略不一定最优。
“黑盒”问题： 用户可能不清楚内部具体的并行方式，调试困难。
框架锁定： 通常与特定框架绑定。

实现： Alpa, Colossal-AI Auto Parallelism, OneFlow, Google的Pathways等（虽然不直接对外提供，但体现了这种思想）。

七、MoE并行 / 专家并行：为“稀疏”而生

最后，我们简单提一下与MoE（Mixture of Experts）架构相关的并行方式。

MoE回顾： MoE模型包含多个“专家”（通常是FFN/MLP网络），和一个“门控网络（Gate）”。对于每个输入Token，门控网络会计算一个权重分布，决定将这个Token主要发送给哪些专家进行处理，最后将选定专家的输出加权组合起来。这样，模型参数量可以很大（有很多专家），但实际计算量（每个Token只激活少数几个专家）可以保持相对较低。

MoE并行（专家并行）：

核心思想： 将不同的专家分布到不同的GPU上。
通信模式： 这涉及到一种All-to-All的通信。因为一个Batch中的不同Token可能会选择不同的专家，而这些专家可能位于不同的GPU上。需要将Token路由（Route）到持有对应专家的GPU上，计算完成后，再将结果收集回来。
本质： 也是一种模型并行，但其并行模式和通信模式与前面讨论的TP、PP有所不同，是专门针对MoE架构的。它通常会和数据并行结合使用（即，将整个MoE层/模型在多个GPU上做数据并行复制，同时每个副本内部的专家分布在不同的GPU上）。

图11：MoE专家并行示意图

优点：

允许训练参数量巨大但计算相对稀疏的MoE模型。

缺点：

All-to-All通信开销： 是主要的瓶颈，对网络带宽要求很高。
负载均衡： 如果门控网络分配不均，可能导致某些GPU上的专家很忙，而另一些很闲。

实现： DeepSpeed-MoE, FairScale MoE, Google的一些内部框架。

总结与展望

好了，今天我们把大模型分布式训练中常见的几大并行技术——数据并行（DP）、张量并行（TP）、流水线并行（PP）、优化器并行（ZeRO）、异构并行、混合并行、自动并行以及MoE并行——都过了一遍。

给大家画个重点：

数据并行 (DP)：简单有效，但有内存冗余。
张量并行 (TP)：解决层内计算和内存瓶颈，但通信大，适合节点内。
流水线并行 (PP)：解决模型深度瓶颈，减少内存，但有气泡，需要微批次。
ZeRO：专治DP的内存冗余，分阶段优化，是DP的强力升级版。
异构并行：利用CPU/NVMe，突破显存极限，牺牲速度。
混合并行：组合拳，威力大，复杂度高，是训练超大模型的标配。
自动并行：让框架代劳，降低门槛，未来趋势。
MoE并行：为MoE架构定制，核心是专家分布和All-to-All通信。

理解这些技术的核心思想、解决的问题、优缺点以及典型的通信模式，对于我们设计、实施和优化大模型训练至关重要，选择哪种或哪几种并行策略，并没有标准答案，需要根据你的模型大小、硬件资源（GPU数量、显存、网络带宽）、可接受的实现复杂度以及对训练速度的要求来综合权衡。

这只是分布式训练世界的第一站。后面我们还可以继续深入聊聊具体的框架实现（比如DeepSpeed、Megatron-LM、FairScale、Colossal-AI等），以及分布式训练中的通信原语、性能优化技巧、容错处理等等。

希望今天的分享能给大家带来一些启发。如果你有任何问题或者想深入讨论某个技术点，欢迎在评论区留言！我们下次再见！

参考：
https://zhuanlan.zhihu.com/p/598714869