大模型多 GPU 分布式训练并行策略详解与选择指南

• GPU:0 基于收集的输出结果和真实 label 计算总损失 loss，并得到 loss 的梯度。
• GPU:0 将计算出的 loss 梯度分发（scatter）到所有 GPU 上。
• 每个 GPU 根据收到的 loss 梯度反向传播，计算出所有模型参数的梯度。
• 所有 GPU 计算出的参数梯度被汇总回 GPU:0。
• GPU:0 基于汇总的梯度更新模型参数，完成一次迭代的训练。

这种架构存在负载不均衡问题。DataParallel 采用的是 Parameter Server 并行架构，在实现多 GPU 或多节点并行训练时，存在一些固有的局限性：

• 通信开销大：每个「计算节点」在每次迭代中都需要与参数服务器进行多次通信，以获取最新的参数更新并将计算的梯度发送回去。这种频繁的通信会导致网络带宽成为瓶颈，尤其是当模型参数量大且 GPU 数量众多时，通信延迟和带宽消耗会显著影响整体训练速度。
• 负载不均衡：其中一个 GPU 被指定为 Master GPU，负责汇总梯度和广播更新等，Master GPU 可能会承担额外的通信和计算负担，导致负载不均衡。这不仅会影响该 GPU 的计算效率，也可能拖慢整个训练过程的速度。同时导致 GPU 利用率不足。
• 仅支持单机多卡模式，无法实现多机多卡训练。

「DistributedDataParallel（DDP)」

DDP 采用多进程架构，赋予了每个 GPU 更多的自由，支持多机多卡分布式训练。每个进程都拥有自己的优化器 Optimizer，可独立优化所有步骤。每个进程中在自己的 GPU 上计算 loss，反向计算梯度。

在 DDP 中，不存在所谓的 Master GPU，所有 GPU 节点地位平等，共同参与训练过程的每一个环节。每个 GPU 都会计算 loss 和梯度，然后通过高效的通信协议（如 AllReduce）与其它 GPU 同步梯度，确保模型参数的一致性。

实现代码如下：

# 初始化分布式环境
import torch.distributed as dist
import torch.optim as optim
import argparse

# 1) 指定通信后端为 nccl（NVIDIA Collective Communications Library），
#    这是针对 GPU 集群优化的高性能通信库
dist.init_process_group(backend='nccl')

# 2）从命令行接收 local_rank 参数，该参数表示当前 GPU 在本地机器上的编号，用于后续的设备设置
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1, type=int)
args = parser.parse_args()

# 3) 设置 cuda
# 根据 local_rank 设置当前进程使用的 GPU 设备，创建对应的 device 对象
torch.cuda.set_device(args.local_rank)
device = torch.device("cuda", args.local_rank)

# 模型设置
# 将模型封装进 DistributedDataParallel，
# 指定模型运行在 local_rank 对应的 GPU 上，同时将模型移动到相应的设备
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
model.to(device)

# Dataset 设置
from torch.utils.data import DataLoader, DistributedSampler
Test_data = FunDataset(args.input)
# 创建 DistributedSampler，用于在分布式环境中对数据集进行采样，确保每个进程处理不同的数据子集
test_sample = DistributedSampler(Test_data)
# 使用 DataLoader 加载数据，指定 sampler 为 DistributedSampler，确保数据的分布式加载和处理
test_data_dataset = DataLoader(dataset=Test_data, batch_size=args.batch_size, shuffle=False,
                               collate_fn=Test_data.collate__fn,
                               drop_last=False, sampler=test_sample)

# 优化器初始化（应在循环外）
optimizer = optim.SGD(model.parameters(), lr=0.001)

# 运行的时候需要设置
for epoch in range(num_epochs):
    # 在每个 epoch 开始时，更新 DistributedSampler 的 epoch，确保数据的随机重排
    test_data_dataset.sampler.set_epoch(epoch)
    # 遍历数据集，前向传播计算预测值，计算损失，执行反向传播和参数更新
    for data, label in trainloader:
        prediction = model(data)
        loss = loss_fn(prediction, label)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

DDP 解决了 DP 模式下存在的效率瓶颈和资源分配不均等问题，每个 GPU 节点都具备独立的数据加载能力，无需等待主 GPU 的数据分发。并且，可执行模型训练的每一环节，包括前向传播、损失计算和反向传播，实现了真正的并行计算。

引入了 DistributedSampler，用于在分布式环境中均匀分割数据集，确保每个 GPU 处理的数据互不重叠，避免了数据冗余和计算浪费。

采用了高效的 Ring All-Reduce 算法作为通信后端，用于梯度的聚合和参数的同步。在每个训练迭代结束时，每个 GPU 计算出自己的梯度后，通过环形网络结构与其他 GPU 进行梯度交换，最终每个 GPU 都能获取到所有 GPU 的梯度信息。

针对 DP 来说，DataLoader 的 batch_size 是针对所有卡训练 batch_size 的和，例如 10 卡，每张卡 batch_size 是 20，那么就要指定 batch_size 为 200。针对 DDP 来说，batch_size 就是每张卡所训练使用的 batch_size 为 20。

2、模型并行 MP

模型并行（Model Parallelism）通常用于解决单节点内存不足的问题。以 GPT-3 为例，该模型拥有 1750 亿参数，如果每个参数都使用 32 位浮点数表示，那么模型需要占用 700GB（即 175G× 4 Bytes）内存。即使使用 16 位浮点数表示，每个模型副本也需要 350GB 的内存。单个加速卡（如 NVIDIA H100）的显存容量（80GB）显然不足以容纳整个模型。

模型并行从计算图的切分角度，可以分为以下两种：

1. 按模型的 「layer 层切分」 到不同设备，即 「层间并行或算子间并行」（Inter-operator Parallelism），也称之为 「流水线并行」（Pipeline Parallelism，PP）。
2. 将计算图层内的 「参数切分」到不同设备，即「层内并行或算子内并行」（Intra-operator Parallelism），也称之为 「张量并行」（Tensor Parallelism，TP）。

1）张量并行

张量并行（Tensor Parallelism，TP）旨在通过将模型参数和中间张量在多个设备（如 GPU）之间进行切分，以克服单个设备内存限制的问题。

张量并行要解决的两个核心问题：如何合理地将参数切分到不同设备，以及如何确保切分后的计算结果在数学上保持一致。

大语言模型都是以 Transformer 结构为基础，Transformer 结构主要由以下三种算子构成：

• 嵌入式表示（Embedding）
• 矩阵乘法（MatMul）
• 交叉熵损失（Cross Entropy Loss）

这三种类型的算子均需要设计对应的张量并行策略，才可以将参数切分到不同设备。

「嵌入式表示（Embedding）」

对于 Embedding 算子，总词表数很大，将导致单计算设备显存无法容纳 Embedding 层参数。

例如，词表数为 64000，嵌入表示维度为 5120，使用 32 位浮点数表示，整层参数需要的显存大约为 6400x5120x4/1024/1024=1250MB，加上反向传播时的梯度存储，总计近 2.5GB，这对于显存有限的设备而言是一个严峻挑战。

为了解决这一问题，可以采用张量并行策略，将 Embedding 层的参数按词维度切分，每个计算设备只存储部分词向量，然后通过汇总各个设备上的部分词向量，从而得到完整的词向量。

在两节点上，可以将 Embedding 层参数按词维度切分为两半，每台设备存储一半的词向量，即参数大小为 [word_size/2, hidden_size]，分别存储在两个设备上。在前向传播过程中，每个设备根据输入的词汇 ID 查询自身的词向量。如果无法查到，则该词的表示为 0。各设备计算得到的词向量结果形状为 [bz, hidden_size]，由于词汇可能被分割在不同设备上，需要通过跨设备的 All Reduce Sum 操作，将所有设备上的词向量结果进行汇总求和，以得到完整的词向量表示。可以看出，这里的输出结果和单节点执行的结果一致。

「矩阵乘法（MatMul）」

矩阵乘的张量模型并行充分利用矩阵分块乘法的原理。

例如，要实现矩阵乘法 Y=X*A。其中，X 是维度为 MxN 的输入矩阵，A 是维度为 NxK 的参数矩阵，Y 是输出，维度为 MxK。

当参数矩阵 A 的尺寸过大，以至于单个卡无法容纳时，可以将参数矩阵 A 切分到多张卡上，并通过集合通信汇集结果，保证最终结果在数学计算上等价于单卡计算结果。

这里 A 切分方式包括按列切块和按行切块两种。

• 「按列切块」。如下图所示，将参数矩阵 A 按列方向切分为 A1 和 A2。将子矩阵 A1 和 A2 分配到两张卡上。在计算过程中，每张卡将执行独立的矩阵乘法，即卡 1 计算 Y1=XA1，卡 2 计算 Y2=XA2。计算完成后，通过 All Gather 操作，每张卡将获取另一张卡上的计算结果。在收集到所有计算结果后，每张卡将拼接它们收到的片段，形成完整的 Y 矩阵。最终，无论在哪张卡上查看，Y 都将是一个完整的 MxK 矩阵，与单卡计算的结果完全等价。

• 「按行切块」。如下图所示，将参数矩阵 A 按列方向切分为 A1 和 A2。为了与按行切块的 A 矩阵相乘，输入矩阵 X（尺寸为 MxN）也需要按列方向切分为 X1 和 X2。将子矩阵 A1 和 A2 分别分配到两张卡上。同时，将 X1 和 X2 也分别分配到对应的卡上。每张卡将执行独立的矩阵乘法，即卡 1 计算 Y1=X1A1，卡 2 计算 Y2=X2A2。计算完成后，通过 All Reduce Sum 操作，每张卡将汇总另一张卡上的计算结果。在收集到所有计算结果后，每张卡将整合它们收到的片段，形成完整的 Y 矩阵的行。同理，参数矩阵 A 按行切块的张量模型并行策略，通过巧妙地切分矩阵和利用多卡的计算能力，有效地解决了单卡显存限制的问题，同时确保了计算结果的数学等价性。

「交叉熵 Loss 计算（CrossEntropyLoss）」

分类网络中，最后一层通常使用 softmax 函数结合交叉熵损失（CrossEntropyLoss）来评估模型的预测结果与真实标签之间的差距。然而，当类别数量非常大时，单个 GPU 卡可能无法存储和计算 logit 矩阵，导致显存溢出。为了解决这一问题，可以采用张量并行策略，将 logit 矩阵和标签按类别数维度切分，通过中间结果的通信，计算出全局的交叉熵损失。具体的计算步骤如下：

首先计算 softmax 值。其中，p 表示张量并行的设备号。

得到 Softmax 结果之后，同时对标签 Target 按类别切分，每个设备得到部分损失，最后再进行一次通信，得到所有类别的损失。

具体的计算步骤如下图所示。

2）流水线并行

所谓流水线并行，就是由于模型太大，无法将整个模型放置到单张 GPU 卡中，因此，将模型的不同层放置到不同的计算设备，降低单个计算设备的显存消耗，从而实现超大规模模型训练。

流水线并行 PP（Pipeline Parallelism），是一种最常用的并行方式，将模型的各个层分段处理，并将每个段分布在不同的计算设备上，使得前后阶段能够流水式、分批进行工作。

前向计算过程中，输入数据首先在 Device 0 上通过 Layer 1 的计算得到中间结果，并将中间结果传输到 Device 1，再继续在 Device 1 上计算得到 Layer 2 和 Layer 3 的输出，并将模型 Layer 3 的输出结果传输到 Device 2。在 Device 2 上，数据经过 Layer 4 的计算，得到最终的前向计算结果。反向传播过程类似。

最后，各个设备上的网络层会使用反向传播过程计算得到的梯度更新参数。由于各个设备间传输的仅是相邻设备间的输出张量，而不是梯度信息，因此通信量较小。

通过将模型切分为多个部分并分布到不同的计算设备上，流水线并行策略有效地扩展了可用于训练的 GPU 显存容量。这意味着原本无法在单一 GPU 上装载的大模型，现在可以通过类似流水线的方式，利用更多 GPU 的显存来承载训练中的模型参数、梯度、优化器状态以及激活值等数据，从而实现超大规模模型的高效训练。

「朴素流水线并行」

当模型规模超过单个 GPU 的处理能力时，朴素层并行（Naive Layer Parallelism）是一种直观的并行策略，将模型的不同层分配到不同的 GPU 上，实现模型的并行化训练。如下所示是一个 4 层序列模型：

output=L4(L3(L2(L1(input)))))   

将其划分到两个 GPU 上：

• GPU1 负责计算前两层：intermediate=L2(L1(input))
• GPU2 负责计算后两层：output=L4(L3(intermediate))

整个朴素层并行前向传播和后向传播的过程如上图所示。GPU1 执行 Layer 1 和 Layer 2 的前向传播，并缓存激活（activations），再将 Layer 2 的输出（intermediate）发送给 GPU2。GPU2 接收来自 GPU1 的中间结果，执行 Layer 3 和 Layer 4 的前向传播，然后计算损失后，开始反向传播，计算 Layer 4 和 Layer 3 的梯度，并将 Layer 3 的梯度返回给 GPU1。GPU1 接收梯度，继续完成 Layer 2 和 Layer 1 的反向传播。

「朴素层并行的缺点：」

• 低 GPU 利用率：同一时刻，只有其中一个 GPU 在执行计算，其余 GPU 处于空闲状态 (又称气泡 bubble)，这导致了计算资源的浪费。
• 计算和通信没有重叠：在数据传输期间，无论是前向传播的中间结果还是反向传播的梯度，GPU 都处于等待状态，这进一步降低了计算效率。
• 高显存占用：GPU1 需要保存整个 mini-batch 的所有激活，直到反向传播完成。如果 mini-batch 很大，这将显著增加显存的需求，可能超出单个 GPU 的显存容量。

「GPipe 流水线并行」

GPipe 通过引入微批次（Micro-batch）的概念，显著提高了模型并行训练的效率。将一个大的 mini-batch 进一步细分为多个更小的、相等大小的微批次（microbatches），并在流水线并行的框架下独立地对每个 microbatch 执行前向传播和反向传播。然后将每个 mircobatch 的梯度相加，就能得到整个 batch 的梯度。由于每个层仅在一个 GPU 上，对 mircobatch 的梯度求和仅需要在本地进行即可，不需要通信。

假设我们有 4 个 GPU，并将模型按层划分为 4 部分，每个部分部署在一个 GPU 上。朴素层并行的过程如下所示：

由此可以看出，在每一时刻，仅有一个 1 个 GPU 工作，并且每个 timesep 花费的时间也比较长，因为 GPU 需要跑完整个 minibatch 的前向传播。

GPipe 将 minibatch 划分为 4 个 microbatch，然后依次送入 GPU0。GPU0 前向传播后，再将结果送入 GPU1，以此类推。整个过程如下所示。GPU0 的前向计算被拆解为 F11、F12、F13、F14，在 GPU0 中计算完成 F11 后，会在 GPU1 中开始计算 F21，同时 GPU0 并行计算 F12。

相比于朴素层并行方法，GPipe 流水线方法可以有效降低并行气泡大小。但是 GPipe 只有当一个 Mini-batch(4 个 Microbatch) 中所有的前向传播计算完成后，才能开始执行反向传播计算。因此还是会产生很多并行气泡，从而降低了系统的并行效率。每个 GPU 需要缓存 4 份中间激活值。

「PipeDream 流水线并行」

PipeDream 流水线并行采用 1F1B 策略，即一个前向通道和一个后向通道，采用任务调度机制，使得下游设备能够在等待上游计算的同时执行其他可并行的任务，从而提高设备的利用率。

1F1B 策略分为非交错式和交错式两种方式调度方式。

「1F1B 非交错式调度」 分为三个阶段：

• 热身阶段：在这个阶段，计算设备执行不同数量的前向计算，为后续阶段做准备。
• 前向 - 后向阶段：设备按照顺序执行一次前向计算，紧接着进行一次后向计算。这一阶段循环执行，直到所有 microbatch 被处理完毕。
• 后向阶段：在完成所有前向计算后，设备执行剩余的后向计算，直至所有计算任务完成。

「1F1B 交错式调度」 要求 microbatch 的数量是流水线阶段的整数倍。每个设备不再仅负责连续多个层的计算，而是可以处理多个层的子集，这些子集被称为模型块。例如：

在传统模式下，设备 1 可能负责层 1-4，设备 2 负责层 5-8。

在交错式模式下，设备 1 可以处理层 1、2、9、10，设备 2 处理层 3、4、11、12，以此类推。这样，每个设备在流水线中被分配到多个阶段，可以并行执行不同阶段的计算任务，从而更好地利用流水线并行的优势。

3、混合并行 HP

在进行上百亿/千亿级以上参数规模的超大模型预训练时，通常会组合上述 (数据并行、张量并行、流水线并行) 多种并行技术一起使用。常见的分布式并行技术组合方案如下。

1）DP+PP

DP rank 0 看不到 GPU2，DP rank 1 看不到 GPU3，对于 DP，只有 GPU 0 和 1 提供数据，就好像只有 2 个 GPU 一样。GPU0 使用 PP 将其部分负载卸载到 GPU2，同样的 GPU1 使用 PP 将其部分负载卸载到 GPU3。

由于每个维度至少需要 2 个 GPU，因此这里至少需要 4 个 GPU。

2）3D 并行 (DP+PP+TP)

3D 并行是由数据并行 (DP)、张量并行 (TP) 和流水线并行 (PP) 组成。由于每个维度至少需要 2 个 GPU，因此这里至少需要 8 个 GPU。

四、分布式训练并行策略选择

1、单节点并行化策略

（1）当单个 GPU 可以装入整个模型时

• DDP (Distributed DataParallel)
• ZeRO

（2）当单个 GPU 无法装入整个模型时

• Pipeline Parallel (PP)
• Tensor Parallel (TP)
• ZeRO

（3）当单个 GPU 无法装入模型的最大层时

• 使用 Tensor Parallel（TP），因为仅靠 PipelineParallel（PP）不足以容纳大型层。
• ZeRO

2、多节点并行化策略

（1）具有快速的节点间连接时

• ZeRO
• 组合使用 PP、TP 和 DP

（2）节点间连接速度较慢且 GPU 内存不足时

• 组合使用 PP、TP、DP 和 ZeRO

在实际应用中，选择合适的并行策略需要综合考虑硬件资源、网络带宽、模型规模及训练目标。ZeRO（Zero Redundancy Optimizer）作为一种先进的显存优化技术，常与上述并行策略结合使用，通过分阶段释放优化器状态、梯度和参数，进一步降低显存占用，是实现超大规模模型训练的关键补充技术。