1. 从单卡到集群为什么我们需要并行训练几年前当我还在用单张显卡吭哧吭哧跑一个几亿参数的模型时最大的愿望就是训练能快点结束。那时候模型变大和训练变慢之间的矛盾就像是用一根吸管去喝光一个游泳池的水。后来大模型的时代毫无预兆地来了动辄千亿、万亿的参数规模让单卡训练彻底成了天方夜谭。这时候并行训练技术就不再是“锦上添花”的优化选项而是“从零到一”的生存必需品。它本质上是一套“分而治之”的工程哲学既然一张卡装不下、算不动那就把模型和计算任务拆开分给多张卡、多台机器大家协同工作共同完成一次前向传播和反向传播。今天要聊的就是当前大规模深度学习训练中三种最核心的并行范式数据并行Data Parallelism, DP、张量模型并行Tensor Parallelism, TP和流水线并行Pipeline Parallelism, PP。很多人可能听过这些名词但它们的区别、适用场景以及如何组合使用往往是一团迷雾。网上有些资料要么过于学术化满篇公式要么过于简略只说“DP分数据TP分层PP分阶段”看完还是不知道具体怎么用、为什么这么用。这篇文章我就结合这几年在集群里摸爬滚打的实际经验把这三种并行的“里子”和“面子”都掰开揉碎了讲清楚。你会发现它们不是互斥的选择题而是一套可以灵活组合、应对不同瓶颈的“组合拳”。2. 数据并行DP最直观的“人海战术”数据并行这通常是大家接触并行训练的第一个概念也是最容易理解的一种。它的核心思想非常简单我有N张显卡称为Worker每个Worker上都拥有一份完整的模型副本。当一批训练数据比如Batch Size64送来时我把这批数据平均分成N份比如每张卡分到16个样本每张卡用自己那部分数据独立地进行前向传播和反向传播计算出各自的梯度。最后所有Worker把计算出的梯度汇总起来通常取平均同步更新到每一个模型副本上。你可以把它想象成一个老师教多个平行班。每个班Worker的学生模型参数教材和教学进度模型结构完全一样。老师把同一套试卷训练数据拆成几部分分给每个班同时考试前向/反向计算。考试结束后把所有班的成绩梯度收上来分析共同的错题梯度平均然后给所有班统一讲解更正参数更新。这样一次就能完成对多个数据子集的学习理想情况下训练速度可以接近线性提升N张卡速度提升接近N倍。2.1 DP的核心流程与通信开销在实际的代码框架如PyTorch的DistributedDataParallel中DP的流程可以细化为以下几步模型复制主进程将初始模型广播到所有Worker上确保大家起点一致。数据分发每个训练步Step开始时每个Worker从数据加载器中获取一个不同的数据子集。这通常通过给每个Worker分配一个唯一的“排名Rank”并设置不同的随机种子来实现。独立前向与损失计算每个Worker用自己的数据独立完成前向传播计算损失。独立反向传播每个Worker根据损失独立计算相对于自己模型中所有参数的梯度。梯度同步All-Reduce这是DP最关键、也往往是性能瓶颈的一步。所有Worker需要将自己计算出的梯度进行汇总。最常见的操作是All-Reduce求和后广播。具体来说所有卡把同一参数的梯度发送到一个中间节点或通过环状、树状等算法在所有卡间通信求和后再把结果广播回所有卡。这样每张卡最终都得到了基于全局批次Global Batch数据计算出的平均梯度。参数更新每个Worker使用同步后的梯度独立地更新自己副本的模型参数。由于梯度一致更新后的所有模型副本仍然保持一致。注意这里有一个关键点DP的“并行”主要体现在数据和计算上模型本身在每张卡上是完整的。通信开销集中在第5步的梯度同步。梯度的大小等于模型参数的总量。因此模型参数量越大DP的通信开销就越大。对于一个175B参数的模型光是梯度同步的数据量就高达几百GB假设用FP16也有350GB这对网络带宽是巨大的挑战。2.2 DP的优势、局限与实战配置优势实现简单主流框架PyTorch DDP, Horovod都有成熟、高效的实现代码侵入性低。通用性强几乎适用于任何模型结构只要单卡能放下整个模型。扩展性好在模型大小不变的情况下增加Worker数量通常能线性降低训练时间在通信不成为瓶颈时。局限与挑战单卡内存墙这是DP最根本的限制。模型必须能完整装入单张GPU的内存。对于当今的大模型这首先就过不了关。通信瓶颈随着Worker数量N和模型参数量的增加梯度All-Reduce的通信量线性增长。当通信时间接近甚至超过计算时间时增加更多的卡反而会使效率下降加速比曲线变得平缓。全局批次大小Global Batch SizeDP的有效批次大小是Per-GPU Batch Size * N。为了达到同样的效果当使用大量GPU时全局批次大小会变得极其巨大如成千上万。过大的批次大小可能影响模型收敛性和最终精度需要仔细调整学习率等超参数。实战心得什么时候用DP当你的模型单卡装得下但希望用更多数据或更快训练时DP是首选。例如训练一个几亿参数的BERT变体使用4-8张卡进行DP是非常有效的。通信优化使用梯度压缩如FP16混合精度训练本身能减半通信量或更激进的1-bit Adam等、通信计算重叠在PyTorch DDP中当backward()为一个参数计算完梯度后可以立即开始该梯度的All-Reduce而不等所有梯度计算完等技术来缓解通信瓶颈。与模型并行的关系DP无法解决大模型单卡放不下的问题。因此对于大模型DP通常需要与下面要讲的TP或PP结合使用形成混合并行策略。3. 张量模型并行TP对巨型层“动手术”当模型中的单个层比如一个庞大的Transformer FFN层或Attention层的矩阵大到一张卡都放不下时DP就无能为力了。这时就需要张量模型并行。TP的核心思想是将一个层内部的巨大权重张量Tensor在空间上切分分布到不同的设备上。每个设备只持有该张量的一部分共同协作来完成该层的计算。最经典的例子是Megatron-LM论文中提出的对Transformer层中GEMM通用矩阵乘法操作的列并行Column Parallel和行并行Row Parallel切分。3.1 TP的切分方式与计算模式我们以一个简单的全连接层Y XA b为例其中X是输入A是权重矩阵b是偏置Y是输出。假设A的维度是[输入维度, 输出维度]。列并行切分输出维度将权重矩阵A按列切分。假设分到2张卡上那么A [A1, A2]每张卡持有A1或A2。计算时每张卡分别计算Y_part X * A_part。由于矩阵乘法对加法的分配律X * A X * [A1, A2] [X*A1, X*A2]。所以每张卡计算出的Y_part实际上是最终输出Y的一部分列。最后需要通过一个All-Gather通信操作将所有卡上的Y_part收集起来拼接成完整的Y。偏置b也相应地按列切分。行并行切分输入维度将权重矩阵A按行切分。假设分到2张卡上那么A [A1; A2]纵向拼接每张卡持有A1或A2。这时输入X也需要被复制到每张卡上。每张卡计算Y_part X * A_part。注意此时Y_part1 X * A1,Y_part2 X * A2而完整的Y X * A1 X * A2 Y_part1 Y_part2。所以每张卡计算出的Y_part是最终输出的一个“部分和”。最后需要通过一个Reduce-Scatter或All-Reduce通信操作将各部分求和得到完整的Y。在Transformer的实际应用中通常对不同的层采用不同的切分策略以优化通信。例如对FFN层的第一层将隐藏维度放大使用列并行对第二层将维度缩小回原状使用行并行这样中间结果的通信可以部分抵消形成一种更优的模式。3.2 TP的通信模式与性能权衡TP的通信发生在层内。每次前向传播和反向传播都需要在切分边界进行通信如All-Gather或Reduce-Scatter。通信的数据量取决于张量切分的维度和切分数量通常与激活值Activation的大小相关。优势能放下超大层解决了单层参数或激活值超过单卡内存的硬性限制。通信量相对可控相比于DP需要同步所有参数的梯度TP的通信通常只涉及当前层的输入/输出或中间结果通信量可能与批次大小和序列长度相关但不直接与模型总参数量成正比。局限与挑战设备间紧密耦合TP将一个层的计算拆散到多卡上这些卡必须高速互联如NVLink, InfiniBand。如果设备间带宽不足通信延迟会成为主要瓶颈严重拖慢计算。因此TP通常在一个多卡服务器节点内部使用比如一张8卡服务器上做8路TP。计算粒度变细切分后每张卡上的矩阵乘法运算规模变小可能无法充分利用GPU的算力特别是Tensor Core导致计算效率Utilization下降。编程复杂性高需要手动或借助特定框架如Megatron-DeepSpeed来定义模型的切分方式代码侵入性强。实战心得什么时候用TP当模型中有少数极其庞大的层例如MoE模型中的专家层或超宽FFN层导致单卡内存不足时TP是精准的“手术刀”。它通常不用于切分整个模型而是针对瓶颈层。与NVLink是黄金搭档务必在具有高速互联的GPU组内使用TP。跨节点的TP通信开销通常难以承受。组合使用TP很少单独使用。一个典型的模式是在节点内使用TP来切分大层同时在多个这样的节点间使用DP来扩展数据并行规模。这就是TPDP的混合并行。4. 流水线并行PP让计算像工厂流水线一样流水线并行解决的是另一个维度的问题模型层数太多即使每层不大但整个模型纵向堆叠起来仍然远超单卡内存。PP的核心思想是将模型的各层按顺序分组每一组称为一个“阶段”Stage放置到不同的设备上。数据像在工厂流水线上一样依次流过各个阶段。假设一个模型有12层我们使用4张卡进行4阶段PP。那么分配可能是GPU0持有第1-3层Stage 0GPU1持有第4-6层Stage 1GPU2持有第7-9层Stage 2GPU3持有第10-12层Stage 3。4.1 朴素PP与巨大的气泡Bubble问题最直观的PP实现是同步的。在训练时第一个微批次Micro-batch的数据进入GPU0Stage 0完成第1-3层计算后将中间结果激活值发送给GPU1Stage 1然后GPU0开始处理第二个微批次同时GPU1处理第一个微批次的第4-6层以此类推。然而这里有一个严重问题流水线填充和排空的时间开销。在流水线启动时后面的GPU要等待前面的GPU输出在流水线结束时前面的GPU会先空闲下来等待后面的GPU完成。这段设备空闲的时间被称为“流水线气泡”Pipeline Bubble。在朴素的同步PP中气泡可能占据相当大比例的计算时间导致设备利用率很低。4.2 GPipe与1F1B调度优化气泡为了减少气泡研究者提出了多种调度策略。最著名的是Google的GPipe和NVIDIA/Microsoft提出的1F1BOne Forward pass followed by One Backward pass。GPipe它引入了微批次的概念。将一个大的全局批次Global Batch分成许多个小的微批次。首先让所有微批次依次通过流水线完成前向传播流水线填充。然后再让梯度依次反向传播流水线排空。GPipe通过增加微批次数量来“加长”流水线使得气泡时间占比相对减小。但它需要缓存所有微批次的前向激活值以供反向传播这带来了巨大的内存开销与微批次数量成正比。1F1B这是一种更优的调度策略。它让前向和反向传播交错进行。每个设备在为一个微批次完成前向传播后只要收到下一个微批次的反向传播梯度就立即开始反向计算。1F1B的优点是激活值内存占用是常数只与流水线阶段数相关与微批次数量无关并且通常能获得比GPipe更小的气泡和更高的硬件利用率。DeepSpeed和Megatron等框架都实现了1F1B或其变种。4.3 PP的通信与内存特性PP的通信发生在阶段之间传递的是层的输入激活值前向和输出梯度反向。通信是点对点的Point-to-Point数据从一个设备发送到下一个设备。优势能放下超深模型是处理层数极多的模型如千层Transformer的主要手段。通信模式简单阶段间通常是相邻设备通信模式固定。可与DP自然结合不同的流水线阶段可以独立地进行数据并行形成PPDP的混合模式。局限与挑战流水线气泡即使使用1F1B气泡依然存在它直接降低了理论峰值利用率。气泡大小与流水线阶段数成正比。负载均衡需要谨慎地将层划分到各个阶段尽量让每个阶段的计算量均衡。否则最慢的阶段会成为整个流水线的瓶颈木桶效应。复杂性高调度逻辑复杂框架实现难度大。用户需要指定切分点调试起来比DP和TP更麻烦。实战心得什么时候用PP当模型层数极多无法放入单个设备或单个节点时。通常用于模型深度方向的扩展。阶段划分是艺术使用 profiling 工具如PyTorch Profiler, NSight Systems分析每层的计算时间和内存消耗尽量让每个阶段的计算时间相等。自动划分工具如Alpa正在研究但目前实践中仍需手动调整。微批次数量选择微批次数量越多气泡相对越小但也会增加调度复杂度。通常设置为流水线阶段数的整数倍以保证负载均衡。与DP、TP组合这是大模型训练的常态。例如在一个由多个节点组成的集群中可以在节点内使用TP来切分宽层在节点间使用PP来切分深度同时在所有数据副本间使用DP。这就是3D并行DPTPPP。5. 混合并行实战以DeepSpeed/Megatron为例的配置解析理解了三种基本并行方式后我们来看它们如何在实际框架中组合使用。目前业界最主流的两个大规模训练框架是DeepSpeed微软和Megatron-LMNVIDIA它们都支持灵活的混合并行。假设我们有一个巨大的模型需要在由8台服务器节点组成的集群上训练每台服务器有8张A100 GPU共64张卡。我们的目标是最大化内存利用和计算效率。一种可能的3D并行配置如下第一维数据并行DP这是我们扩展训练规模的基础。假设我们设置DP4。这意味着我们将有4个完全相同的模型副本在同时训练。每个副本会处理一部分数据。第二维张量模型并行TP为了放下模型中那些超大的权重矩阵我们在单个节点内部进行TP。因为节点内GPU有NVLink高速互联通信效率高。设置TP8即一个节点内的8张卡组成一个TP组。这样每个模型副本的庞大层会被切分到这8张卡上。第三维流水线并行PP为了放下模型的深度我们在节点之间进行PP。我们有4个模型副本DP4每个副本需要被PP切分。总节点数是8个每个TP组占用1个节点8卡。因此可用于PP的“单元”是TP组共有8个TP组。我们需要将这8个TP组分配给4个模型副本做PP。设置PP2。这意味着每个模型副本被切成2个流水线阶段Stage每个阶段占用1个TP组即1个节点。那么4个模型副本 * 2个阶段 正好需要8个TP组节点与我们的集群匹配。最终并行配置DP4 TP8 PP2。总GPU数验证DP * TP * PP 4 * 8 * 2 64 符合集群总卡数。物理映射集群有8个节点每个节点8卡。每个节点内部8卡通过TP紧密协作形成一个“超级设备”。每2个这样的节点通过PP连接形成一个完整的模型流水线一个模型副本。这样的流水线有4条DP4同时处理4份不同的数据。通信模式分析TP通信发生在节点内的8张卡之间通信频繁数据量中等依赖NVLink高速带宽。PP通信发生在节点间同一个模型副本的两个阶段之间通信是点对点的数据量取决于层间激活值大小对互联带宽有要求但频率低于TP。DP通信发生在所有4个模型副本之间即所有64张卡需要同步梯度。这是通信量最大的一步但频率最低每个微批次或每个梯度累积步一次。通常依赖全局的All-Reduce操作对集群的全局互联如InfiniBand带宽和延迟是终极考验。在DeepSpeed配置文件中的体现{ train_micro_batch_size_per_gpu: 4, gradient_accumulation_steps: 8, zero_optimization: { stage: 3, offload_optimizer: {device: cpu} }, fp16: {enabled: true}, parallelism: { tp: {size: 8}, pp: {size: 2}, dp: {size: 4} } }这里gradient_accumulation_steps用于累积梯度使得per_gpu_batch * accumulation_steps * dp_size达到你想要的全局批次大小。ZeRO-3优化阶段可以进一步将优化器状态、梯度和参数分片与模型并行协同工作极大节省内存。6. 如何为你的任务选择并行策略没有放之四海而皆准的配置。选择哪种或哪几种并行方式取决于你的模型特性和硬件环境。下面是一个决策流程参考评估模型规模参数量估算模型总参数量例如175B。激活值估算训练时前向传播产生的中间激活值内存与批次大小、序列长度强相关。单层大小找出模型中最大的层如FFN中间层。评估硬件环境单卡内存你的GPU有多少显存如80GB A100。节点内互联节点内GPU是否有高速互联NVLink/NVSwitch。节点间互联网络带宽和延迟如何InfiniBand/高速以太网。总卡数你总共可以使用多少张GPU决策流程第一步能否单卡放下如果能直接使用DP扩展到多卡是最简单高效的选择。第二步单卡放不下是因为“宽”还是“深”如果是“宽”少数层参数巨大优先考虑在单个高速节点内使用TP。TP的尺寸tp_size取决于最大层需要被切分成几份才能放入单卡。如果是“深”层数太多考虑使用PP。pp_size取决于你需要将模型切成几段才能放入单个设备或TP组。第三步结合资源确定混合策略。在确定了tp_size和pp_size后你得到了一个“模型并行单元”一个完整的模型需要多少张卡model_parallel_size tp_size * pp_size。用总卡数除以model_parallel_size就得到了可以并行的模型副本数量即dp_size。确保dp_size * tp_size * pp_size 总卡数。第四步微调与优化。负载均衡对于PP手动或利用工具调整阶段划分点使各阶段计算时间相近。通信优化启用梯度压缩、通信计算重叠。对于DP通信考虑使用ZeRO阶段2或3来减少每卡的内存和通信开销。超参数调整特别是学习率需要根据新的全局批次大小per_gpu_batch * dp_size * gradient_accumulation_steps重新调整。在我最近参与的一个千亿参数模型项目中我们最终采用的配置是tp8在8卡节点内切分大层pp4跨4个节点堆叠深度dp16共使用了512张卡。调试初期PP的阶段划分不当导致其中一个节点负载过重成为瓶颈。我们通过细致的性能剖析重新调整了层到阶段的分配使吞吐量提升了近20%。这个经历让我深刻体会到并行训练不仅仅是配置几个数字更是一个需要结合理论、经验和反复调试的系统工程。
1. 从单卡到集群为什么我们需要并行训练几年前当我还在用单张显卡吭哧吭哧跑一个几亿参数的模型时最大的愿望就是训练能快点结束。那时候模型变大和训练变慢之间的矛盾就像是用一根吸管去喝光一个游泳池的水。后来大模型的时代毫无预兆地来了动辄千亿、万亿的参数规模让单卡训练彻底成了天方夜谭。这时候并行训练技术就不再是“锦上添花”的优化选项而是“从零到一”的生存必需品。它本质上是一套“分而治之”的工程哲学既然一张卡装不下、算不动那就把模型和计算任务拆开分给多张卡、多台机器大家协同工作共同完成一次前向传播和反向传播。今天要聊的就是当前大规模深度学习训练中三种最核心的并行范式数据并行Data Parallelism, DP、张量模型并行Tensor Parallelism, TP和流水线并行Pipeline Parallelism, PP。很多人可能听过这些名词但它们的区别、适用场景以及如何组合使用往往是一团迷雾。网上有些资料要么过于学术化满篇公式要么过于简略只说“DP分数据TP分层PP分阶段”看完还是不知道具体怎么用、为什么这么用。这篇文章我就结合这几年在集群里摸爬滚打的实际经验把这三种并行的“里子”和“面子”都掰开揉碎了讲清楚。你会发现它们不是互斥的选择题而是一套可以灵活组合、应对不同瓶颈的“组合拳”。2. 数据并行DP最直观的“人海战术”数据并行这通常是大家接触并行训练的第一个概念也是最容易理解的一种。它的核心思想非常简单我有N张显卡称为Worker每个Worker上都拥有一份完整的模型副本。当一批训练数据比如Batch Size64送来时我把这批数据平均分成N份比如每张卡分到16个样本每张卡用自己那部分数据独立地进行前向传播和反向传播计算出各自的梯度。最后所有Worker把计算出的梯度汇总起来通常取平均同步更新到每一个模型副本上。你可以把它想象成一个老师教多个平行班。每个班Worker的学生模型参数教材和教学进度模型结构完全一样。老师把同一套试卷训练数据拆成几部分分给每个班同时考试前向/反向计算。考试结束后把所有班的成绩梯度收上来分析共同的错题梯度平均然后给所有班统一讲解更正参数更新。这样一次就能完成对多个数据子集的学习理想情况下训练速度可以接近线性提升N张卡速度提升接近N倍。2.1 DP的核心流程与通信开销在实际的代码框架如PyTorch的DistributedDataParallel中DP的流程可以细化为以下几步模型复制主进程将初始模型广播到所有Worker上确保大家起点一致。数据分发每个训练步Step开始时每个Worker从数据加载器中获取一个不同的数据子集。这通常通过给每个Worker分配一个唯一的“排名Rank”并设置不同的随机种子来实现。独立前向与损失计算每个Worker用自己的数据独立完成前向传播计算损失。独立反向传播每个Worker根据损失独立计算相对于自己模型中所有参数的梯度。梯度同步All-Reduce这是DP最关键、也往往是性能瓶颈的一步。所有Worker需要将自己计算出的梯度进行汇总。最常见的操作是All-Reduce求和后广播。具体来说所有卡把同一参数的梯度发送到一个中间节点或通过环状、树状等算法在所有卡间通信求和后再把结果广播回所有卡。这样每张卡最终都得到了基于全局批次Global Batch数据计算出的平均梯度。参数更新每个Worker使用同步后的梯度独立地更新自己副本的模型参数。由于梯度一致更新后的所有模型副本仍然保持一致。注意这里有一个关键点DP的“并行”主要体现在数据和计算上模型本身在每张卡上是完整的。通信开销集中在第5步的梯度同步。梯度的大小等于模型参数的总量。因此模型参数量越大DP的通信开销就越大。对于一个175B参数的模型光是梯度同步的数据量就高达几百GB假设用FP16也有350GB这对网络带宽是巨大的挑战。2.2 DP的优势、局限与实战配置优势实现简单主流框架PyTorch DDP, Horovod都有成熟、高效的实现代码侵入性低。通用性强几乎适用于任何模型结构只要单卡能放下整个模型。扩展性好在模型大小不变的情况下增加Worker数量通常能线性降低训练时间在通信不成为瓶颈时。局限与挑战单卡内存墙这是DP最根本的限制。模型必须能完整装入单张GPU的内存。对于当今的大模型这首先就过不了关。通信瓶颈随着Worker数量N和模型参数量的增加梯度All-Reduce的通信量线性增长。当通信时间接近甚至超过计算时间时增加更多的卡反而会使效率下降加速比曲线变得平缓。全局批次大小Global Batch SizeDP的有效批次大小是Per-GPU Batch Size * N。为了达到同样的效果当使用大量GPU时全局批次大小会变得极其巨大如成千上万。过大的批次大小可能影响模型收敛性和最终精度需要仔细调整学习率等超参数。实战心得什么时候用DP当你的模型单卡装得下但希望用更多数据或更快训练时DP是首选。例如训练一个几亿参数的BERT变体使用4-8张卡进行DP是非常有效的。通信优化使用梯度压缩如FP16混合精度训练本身能减半通信量或更激进的1-bit Adam等、通信计算重叠在PyTorch DDP中当backward()为一个参数计算完梯度后可以立即开始该梯度的All-Reduce而不等所有梯度计算完等技术来缓解通信瓶颈。与模型并行的关系DP无法解决大模型单卡放不下的问题。因此对于大模型DP通常需要与下面要讲的TP或PP结合使用形成混合并行策略。3. 张量模型并行TP对巨型层“动手术”当模型中的单个层比如一个庞大的Transformer FFN层或Attention层的矩阵大到一张卡都放不下时DP就无能为力了。这时就需要张量模型并行。TP的核心思想是将一个层内部的巨大权重张量Tensor在空间上切分分布到不同的设备上。每个设备只持有该张量的一部分共同协作来完成该层的计算。最经典的例子是Megatron-LM论文中提出的对Transformer层中GEMM通用矩阵乘法操作的列并行Column Parallel和行并行Row Parallel切分。3.1 TP的切分方式与计算模式我们以一个简单的全连接层Y XA b为例其中X是输入A是权重矩阵b是偏置Y是输出。假设A的维度是[输入维度, 输出维度]。列并行切分输出维度将权重矩阵A按列切分。假设分到2张卡上那么A [A1, A2]每张卡持有A1或A2。计算时每张卡分别计算Y_part X * A_part。由于矩阵乘法对加法的分配律X * A X * [A1, A2] [X*A1, X*A2]。所以每张卡计算出的Y_part实际上是最终输出Y的一部分列。最后需要通过一个All-Gather通信操作将所有卡上的Y_part收集起来拼接成完整的Y。偏置b也相应地按列切分。行并行切分输入维度将权重矩阵A按行切分。假设分到2张卡上那么A [A1; A2]纵向拼接每张卡持有A1或A2。这时输入X也需要被复制到每张卡上。每张卡计算Y_part X * A_part。注意此时Y_part1 X * A1,Y_part2 X * A2而完整的Y X * A1 X * A2 Y_part1 Y_part2。所以每张卡计算出的Y_part是最终输出的一个“部分和”。最后需要通过一个Reduce-Scatter或All-Reduce通信操作将各部分求和得到完整的Y。在Transformer的实际应用中通常对不同的层采用不同的切分策略以优化通信。例如对FFN层的第一层将隐藏维度放大使用列并行对第二层将维度缩小回原状使用行并行这样中间结果的通信可以部分抵消形成一种更优的模式。3.2 TP的通信模式与性能权衡TP的通信发生在层内。每次前向传播和反向传播都需要在切分边界进行通信如All-Gather或Reduce-Scatter。通信的数据量取决于张量切分的维度和切分数量通常与激活值Activation的大小相关。优势能放下超大层解决了单层参数或激活值超过单卡内存的硬性限制。通信量相对可控相比于DP需要同步所有参数的梯度TP的通信通常只涉及当前层的输入/输出或中间结果通信量可能与批次大小和序列长度相关但不直接与模型总参数量成正比。局限与挑战设备间紧密耦合TP将一个层的计算拆散到多卡上这些卡必须高速互联如NVLink, InfiniBand。如果设备间带宽不足通信延迟会成为主要瓶颈严重拖慢计算。因此TP通常在一个多卡服务器节点内部使用比如一张8卡服务器上做8路TP。计算粒度变细切分后每张卡上的矩阵乘法运算规模变小可能无法充分利用GPU的算力特别是Tensor Core导致计算效率Utilization下降。编程复杂性高需要手动或借助特定框架如Megatron-DeepSpeed来定义模型的切分方式代码侵入性强。实战心得什么时候用TP当模型中有少数极其庞大的层例如MoE模型中的专家层或超宽FFN层导致单卡内存不足时TP是精准的“手术刀”。它通常不用于切分整个模型而是针对瓶颈层。与NVLink是黄金搭档务必在具有高速互联的GPU组内使用TP。跨节点的TP通信开销通常难以承受。组合使用TP很少单独使用。一个典型的模式是在节点内使用TP来切分大层同时在多个这样的节点间使用DP来扩展数据并行规模。这就是TPDP的混合并行。4. 流水线并行PP让计算像工厂流水线一样流水线并行解决的是另一个维度的问题模型层数太多即使每层不大但整个模型纵向堆叠起来仍然远超单卡内存。PP的核心思想是将模型的各层按顺序分组每一组称为一个“阶段”Stage放置到不同的设备上。数据像在工厂流水线上一样依次流过各个阶段。假设一个模型有12层我们使用4张卡进行4阶段PP。那么分配可能是GPU0持有第1-3层Stage 0GPU1持有第4-6层Stage 1GPU2持有第7-9层Stage 2GPU3持有第10-12层Stage 3。4.1 朴素PP与巨大的气泡Bubble问题最直观的PP实现是同步的。在训练时第一个微批次Micro-batch的数据进入GPU0Stage 0完成第1-3层计算后将中间结果激活值发送给GPU1Stage 1然后GPU0开始处理第二个微批次同时GPU1处理第一个微批次的第4-6层以此类推。然而这里有一个严重问题流水线填充和排空的时间开销。在流水线启动时后面的GPU要等待前面的GPU输出在流水线结束时前面的GPU会先空闲下来等待后面的GPU完成。这段设备空闲的时间被称为“流水线气泡”Pipeline Bubble。在朴素的同步PP中气泡可能占据相当大比例的计算时间导致设备利用率很低。4.2 GPipe与1F1B调度优化气泡为了减少气泡研究者提出了多种调度策略。最著名的是Google的GPipe和NVIDIA/Microsoft提出的1F1BOne Forward pass followed by One Backward pass。GPipe它引入了微批次的概念。将一个大的全局批次Global Batch分成许多个小的微批次。首先让所有微批次依次通过流水线完成前向传播流水线填充。然后再让梯度依次反向传播流水线排空。GPipe通过增加微批次数量来“加长”流水线使得气泡时间占比相对减小。但它需要缓存所有微批次的前向激活值以供反向传播这带来了巨大的内存开销与微批次数量成正比。1F1B这是一种更优的调度策略。它让前向和反向传播交错进行。每个设备在为一个微批次完成前向传播后只要收到下一个微批次的反向传播梯度就立即开始反向计算。1F1B的优点是激活值内存占用是常数只与流水线阶段数相关与微批次数量无关并且通常能获得比GPipe更小的气泡和更高的硬件利用率。DeepSpeed和Megatron等框架都实现了1F1B或其变种。4.3 PP的通信与内存特性PP的通信发生在阶段之间传递的是层的输入激活值前向和输出梯度反向。通信是点对点的Point-to-Point数据从一个设备发送到下一个设备。优势能放下超深模型是处理层数极多的模型如千层Transformer的主要手段。通信模式简单阶段间通常是相邻设备通信模式固定。可与DP自然结合不同的流水线阶段可以独立地进行数据并行形成PPDP的混合模式。局限与挑战流水线气泡即使使用1F1B气泡依然存在它直接降低了理论峰值利用率。气泡大小与流水线阶段数成正比。负载均衡需要谨慎地将层划分到各个阶段尽量让每个阶段的计算量均衡。否则最慢的阶段会成为整个流水线的瓶颈木桶效应。复杂性高调度逻辑复杂框架实现难度大。用户需要指定切分点调试起来比DP和TP更麻烦。实战心得什么时候用PP当模型层数极多无法放入单个设备或单个节点时。通常用于模型深度方向的扩展。阶段划分是艺术使用 profiling 工具如PyTorch Profiler, NSight Systems分析每层的计算时间和内存消耗尽量让每个阶段的计算时间相等。自动划分工具如Alpa正在研究但目前实践中仍需手动调整。微批次数量选择微批次数量越多气泡相对越小但也会增加调度复杂度。通常设置为流水线阶段数的整数倍以保证负载均衡。与DP、TP组合这是大模型训练的常态。例如在一个由多个节点组成的集群中可以在节点内使用TP来切分宽层在节点间使用PP来切分深度同时在所有数据副本间使用DP。这就是3D并行DPTPPP。5. 混合并行实战以DeepSpeed/Megatron为例的配置解析理解了三种基本并行方式后我们来看它们如何在实际框架中组合使用。目前业界最主流的两个大规模训练框架是DeepSpeed微软和Megatron-LMNVIDIA它们都支持灵活的混合并行。假设我们有一个巨大的模型需要在由8台服务器节点组成的集群上训练每台服务器有8张A100 GPU共64张卡。我们的目标是最大化内存利用和计算效率。一种可能的3D并行配置如下第一维数据并行DP这是我们扩展训练规模的基础。假设我们设置DP4。这意味着我们将有4个完全相同的模型副本在同时训练。每个副本会处理一部分数据。第二维张量模型并行TP为了放下模型中那些超大的权重矩阵我们在单个节点内部进行TP。因为节点内GPU有NVLink高速互联通信效率高。设置TP8即一个节点内的8张卡组成一个TP组。这样每个模型副本的庞大层会被切分到这8张卡上。第三维流水线并行PP为了放下模型的深度我们在节点之间进行PP。我们有4个模型副本DP4每个副本需要被PP切分。总节点数是8个每个TP组占用1个节点8卡。因此可用于PP的“单元”是TP组共有8个TP组。我们需要将这8个TP组分配给4个模型副本做PP。设置PP2。这意味着每个模型副本被切成2个流水线阶段Stage每个阶段占用1个TP组即1个节点。那么4个模型副本 * 2个阶段 正好需要8个TP组节点与我们的集群匹配。最终并行配置DP4 TP8 PP2。总GPU数验证DP * TP * PP 4 * 8 * 2 64 符合集群总卡数。物理映射集群有8个节点每个节点8卡。每个节点内部8卡通过TP紧密协作形成一个“超级设备”。每2个这样的节点通过PP连接形成一个完整的模型流水线一个模型副本。这样的流水线有4条DP4同时处理4份不同的数据。通信模式分析TP通信发生在节点内的8张卡之间通信频繁数据量中等依赖NVLink高速带宽。PP通信发生在节点间同一个模型副本的两个阶段之间通信是点对点的数据量取决于层间激活值大小对互联带宽有要求但频率低于TP。DP通信发生在所有4个模型副本之间即所有64张卡需要同步梯度。这是通信量最大的一步但频率最低每个微批次或每个梯度累积步一次。通常依赖全局的All-Reduce操作对集群的全局互联如InfiniBand带宽和延迟是终极考验。在DeepSpeed配置文件中的体现{ train_micro_batch_size_per_gpu: 4, gradient_accumulation_steps: 8, zero_optimization: { stage: 3, offload_optimizer: {device: cpu} }, fp16: {enabled: true}, parallelism: { tp: {size: 8}, pp: {size: 2}, dp: {size: 4} } }这里gradient_accumulation_steps用于累积梯度使得per_gpu_batch * accumulation_steps * dp_size达到你想要的全局批次大小。ZeRO-3优化阶段可以进一步将优化器状态、梯度和参数分片与模型并行协同工作极大节省内存。6. 如何为你的任务选择并行策略没有放之四海而皆准的配置。选择哪种或哪几种并行方式取决于你的模型特性和硬件环境。下面是一个决策流程参考评估模型规模参数量估算模型总参数量例如175B。激活值估算训练时前向传播产生的中间激活值内存与批次大小、序列长度强相关。单层大小找出模型中最大的层如FFN中间层。评估硬件环境单卡内存你的GPU有多少显存如80GB A100。节点内互联节点内GPU是否有高速互联NVLink/NVSwitch。节点间互联网络带宽和延迟如何InfiniBand/高速以太网。总卡数你总共可以使用多少张GPU决策流程第一步能否单卡放下如果能直接使用DP扩展到多卡是最简单高效的选择。第二步单卡放不下是因为“宽”还是“深”如果是“宽”少数层参数巨大优先考虑在单个高速节点内使用TP。TP的尺寸tp_size取决于最大层需要被切分成几份才能放入单卡。如果是“深”层数太多考虑使用PP。pp_size取决于你需要将模型切成几段才能放入单个设备或TP组。第三步结合资源确定混合策略。在确定了tp_size和pp_size后你得到了一个“模型并行单元”一个完整的模型需要多少张卡model_parallel_size tp_size * pp_size。用总卡数除以model_parallel_size就得到了可以并行的模型副本数量即dp_size。确保dp_size * tp_size * pp_size 总卡数。第四步微调与优化。负载均衡对于PP手动或利用工具调整阶段划分点使各阶段计算时间相近。通信优化启用梯度压缩、通信计算重叠。对于DP通信考虑使用ZeRO阶段2或3来减少每卡的内存和通信开销。超参数调整特别是学习率需要根据新的全局批次大小per_gpu_batch * dp_size * gradient_accumulation_steps重新调整。在我最近参与的一个千亿参数模型项目中我们最终采用的配置是tp8在8卡节点内切分大层pp4跨4个节点堆叠深度dp16共使用了512张卡。调试初期PP的阶段划分不当导致其中一个节点负载过重成为瓶颈。我们通过细致的性能剖析重新调整了层到阶段的分配使吞吐量提升了近20%。这个经历让我深刻体会到并行训练不仅仅是配置几个数字更是一个需要结合理论、经验和反复调试的系统工程。
