PythonAI算法
大模型多 GPU 分布式训练并行策略详解与选择指南
大模型在多 GPU 环境下的分布式训练并行策略。涵盖数据并行（DP/DDP）、模型并行（张量并行 TP/流水线并行 PP）及混合并行（HP）。分析了各策略的原理、优缺点及适用场景，如 DP 适合单机多卡，TP 解决显存不足，PP 处理层间通信。最后提供了单节点与多节点环境下的策略选择建议，帮助开发者根据硬件资源和模型规模优化训练效率。
大模型在多 GPU 环境下的分布式训练并行策略。涵盖数据并行（DP/DDP）、模型并行（张量并行 TP/流水线并行 PP）及混合并行（HP）。分析了各策略的原理、优缺点及适用场景，如 DP 适合单机多卡，TP 解决显存不足，PP 处理层间通信。最后提供了单节点与多节点环境下的策略选择建议，帮助开发者根据硬件资源和模型规模优化训练效率。
分布式训练系统的核心目标是将原本在单一计算节点上进行的模型训练过程，转化为能在多个计算节点上并行执行，以加速训练速度并支持更大规模的模型和数据集。
在单节点模型训练中，系统结构主要由两大部分组成：数据和模型。训练过程由多个数据小批次（Mini-batch）完成。数据表示一个数据小批次。训练系统会利用数据小批次根据损失函数和优化算法生成梯度，从而对模型参数进行修正。
针对大语言模型多层神经网络的执行过程，模型训练过程可以抽象为一个计算图（Computational Graph）。这个图由多个相互连接的算子（Operator）构成，每个算子对应神经网络中的一个层（Neural Network Layer），如卷积层、全连接层等。参数（Weights）则是这些层在训练过程中不断更新的权重。
计算图的执行过程可以分为前向传播和反向传播两个阶段。
「前向计算（Forward Pass）」
「反向计算（Backward Pass）」
根据单设备模型训练流程，可以看出，如果进行并行加速，可以从数据和模型两个维度考虑：
数据并行是最常用的并行训练方式，主要分为 DataParallel(DP) 和 DistributedDataParallel(DDP) 两种。
「DP」
DP 是早期使用的数据并行方案，通过 torch.nn.DataParallel() 来调用，代码如下:
# 设置可见的 GPU
import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0,1,2,3"
# 将模型放到 GPU 0 上，必须先把模型放在 GPU 上，后面才可以调用 DP
model.cuda()
# 构建 DataParallel 数据并行化
model = torch.nn.DataParallel(model)
DP 核心思想是将一个大的 batch 数据分割成多个子 batch，并将子 batch 分配给不同的 GPU 进行并行计算。
「前向传播：」
「反向传播：」
这种架构存在负载不均衡问题。DataParallel 采用的是 Parameter Server 并行架构，在实现多 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)
data, label 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。
模型并行（Model Parallelism）通常用于解决单节点内存不足的问题。以 GPT-3 为例，该模型拥有 1750 亿参数，如果每个参数都使用 32 位浮点数表示，那么模型需要占用 700GB（即 175G× 4 Bytes）内存。即使使用 16 位浮点数表示，每个模型副本也需要 350GB 的内存。单个加速卡（如 NVIDIA H100）的显存容量（80GB）显然不足以容纳整个模型。
模型并行从计算图的切分角度，可以分为以下两种：
张量并行（Tensor Parallelism，TP）旨在通过将模型参数和中间张量在多个设备（如 GPU）之间进行切分，以克服单个设备内存限制的问题。
张量并行要解决的两个核心问题：如何合理地将参数切分到不同设备，以及如何确保切分后的计算结果在数学上保持一致。
大语言模型都是以 Transformer 结构为基础，Transformer 结构主要由以下三种算子构成：
这三种类型的算子均需要设计对应的张量并行策略，才可以将参数切分到不同设备。
「嵌入式表示（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 按类别切分，每个设备得到部分损失，最后再进行一次通信，得到所有类别的损失。
具体的计算步骤如下图所示。
所谓流水线并行，就是由于模型太大，无法将整个模型放置到单张 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 执行 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 的反向传播。
「朴素层并行的缺点：」
「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 非交错式调度」 分为三个阶段：
「1F1B 交错式调度」 要求 microbatch 的数量是流水线阶段的整数倍。每个设备不再仅负责连续多个层的计算，而是可以处理多个层的子集，这些子集被称为模型块。例如：
在传统模式下，设备 1 可能负责层 1-4，设备 2 负责层 5-8。
在交错式模式下，设备 1 可以处理层 1、2、9、10，设备 2 处理层 3、4、11、12，以此类推。这样，每个设备在流水线中被分配到多个阶段，可以并行执行不同阶段的计算任务，从而更好地利用流水线并行的优势。
在进行上百亿/千亿级以上参数规模的超大模型预训练时，通常会组合上述 (数据并行、张量并行、流水线并行) 多种并行技术一起使用。常见的分布式并行技术组合方案如下。
DP rank 0 看不到 GPU2，DP rank 1 看不到 GPU3，对于 DP，只有 GPU 0 和 1 提供数据，就好像只有 2 个 GPU 一样。GPU0 使用 PP 将其部分负载卸载到 GPU2，同样的 GPU1 使用 PP 将其部分负载卸载到 GPU3。
由于每个维度至少需要 2 个 GPU，因此这里至少需要 4 个 GPU。
3D 并行是由数据并行 (DP)、张量并行 (TP) 和流水线并行 (PP) 组成。由于每个维度至少需要 2 个 GPU，因此这里至少需要 8 个 GPU。
（1）当单个 GPU 可以装入整个模型时
（2）当单个 GPU 无法装入整个模型时
（3）当单个 GPU 无法装入模型的最大层时
（1）具有快速的节点间连接时
（2）节点间连接速度较慢且 GPU 内存不足时
在实际应用中，选择合适的并行策略需要综合考虑硬件资源、网络带宽、模型规模及训练目标。ZeRO（Zero Redundancy Optimizer）作为一种先进的显存优化技术，常与上述并行策略结合使用，通过分阶段释放优化器状态、梯度和参数，进一步降低显存占用，是实现超大规模模型训练的关键补充技术。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online