硬核拆解：大模型训练的分布式三剑客——数据并行、模型并行与流水线并行

Posted on 2026-05-24 In AI 技术

引言：打破单卡的“物理枷锁”

在过去的几年里，大语言模型（LLM）的参数量经历了从亿级到千亿甚至万亿级的爆炸式增长。从 GPT-3 的 1750 亿参数，到如今的 GLM-5、Llama 3 等模型，我们正在见证一个“暴力美学”的时代。

然而，现实是骨感的。无论 GPU 算力提升得多快，单张显卡的显存（VRAM）始终是有限的。一块顶级的 H100 显存约为 80GB，而训练一个千亿参数的模型，仅模型权重、梯度和优化器状态就需要耗费数 TB 的内存。此外，计算时间也成为了壁垒——如果用单卡训练 GPT-3，可能需要几十年的时间。

面对**“显存墙”和“算力墙”**，分布式训练成为了大模型时代唯一的出路。如何在成百上千张 GPU 上高效地训练一个庞大的模型？业界经过多年的探索，沉淀出了一套成熟的“组合拳”。

今天，我们将深入剖析大模型分布式训练的三大核心策略：数据并行、模型并行（以张量并行为主）与流水线并行。本文不仅会讲透其背后的核心原理，还会结合 PyTorch 代码实战，带你彻底搞懂大模型是怎么“跑”起来的。

一、数据并行：最简单粗暴的“抄作业”

1. 核心思想

数据并行是最基础、最常见的分布式策略。它的核心思想非常简单：复制模型，切分数据。

假设我们有 4 张 GPU，我们就把完整的模型复制 4 份，分别加载到这 4 张 GPU 上。然后，我们把训练数据集切分为 4 份，每张 GPU 拿到不同的数据子集（Mini-batch），独立进行前向传播和反向传播，计算出梯度。最后，将 4 张 GPU 的梯度进行汇总求平均，同步更新各个 GPU 上的模型参数。

2. 从 DP 到 DDP 的进化

在早期，PyTorch 提供了 DataParallel (DP) 模式。DP 采用的是单进程多线程，由于 Python 的 GIL（全局解释器锁）限制，以及参数服务器架构导致的网络通信瓶颈，DP 的效率非常低下。

如今，业界早已全面转向 DistributedDataParallel (DDP)。DDP 采用多进程架构，每个 GPU 对应一个进程，并且引入了环形全规约 通信算法。在反向传播的过程中，各个 GPU 的梯度是异步进行通信的，极大地降低了通信延迟。

3. 进阶：ZeRO (Zero Redundancy Optimizer)

传统的 DDP 有一个致命弱点：显存冗余。每张 GPU 上都保存了完整的模型参数、梯度和优化器状态（如 Adam 的动量）。对于一个 7B（70亿）参数的模型，仅优化器状态就需要占用约 112GB 显存，单卡根本装不下。

微软 DeepSpeed 提出的 ZeRO 技术打破了这一僵局。ZeRO 分为三个阶段：

ZeRO-1：切分优化器状态。显存随 GPU 数量线性下降。
ZeRO-2：切分优化器状态 + 梯度。
ZeRO-3：切分优化器状态 + 梯度 + 模型参数。此时，每张 GPU 上只有模型的一部分，彻底解决了单卡装不下大模型的问题。

4. PyTorch DDP 代码实战

下面是一段使用 PyTorch 原生 DDP 进行数据并行训练的基础代码框架：

import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def setup_ddp():
    # 初始化进程组，通常使用 NCCL 后端进行 GPU 通信
    dist.init_process_method(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def main():
    local_rank = setup_ddp()
    
    # 1. 构建模型并将其移动到当前 GPU
    model = nn.Linear(1024, 1024).to(local_rank)
    
    # 2. 使用 DistributedDataParallel 包装模型
    ddp_model = DDP(model, device_ids=[local_rank])
    
    # 3. 准备数据集与 DistributedSampler
    dataset = torch.randn(10000, 1024) # 简单模拟数据集
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)
    
    optimizer = torch.optim.Adam(ddp_model.parameters(), lr=0.001)
    loss_fn = nn.MSELoss()

    for epoch in range(10):
        # 每次 epoch 必须设置 sampler，保证每轮数据打乱顺序一致
        sampler.set_epoch(epoch)
        for batch_data in dataloader:
            inputs = batch_data.to(local_rank)
            labels = torch.randn_like(inputs).to(local_rank)
            
            optimizer.zero_grad()
            outputs = ddp_model(inputs)
            loss = loss_fn(outputs, labels)
            
            # 反向传播时，DDP 会自动在后台进行梯度同步
            loss.backward()
            
            optimizer.step()

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

运行方式：通过 torchrun --nproc_per_node=4 ddp_script.py 启动。

二、模型并行（张量并行）：解剖模型的“手术刀”

1. 核心思想

数据并行要求单卡必须装得下整个模型（除非使用 ZeRO-3）。当模型大到连一张卡都装不下时，我们就需要模型并行。

广义的模型并行分为两种：张量并行（TP, Tensor Parallelism） 和下一节要讲的 流水线并行（PP, Pipeline Parallelism）。

张量并行是“横向”的切割。以 Transformer 层中的线性层（全连接层）为例，其核心计算为 $Y = XA$ 。张量并行会将权重矩阵 $A$ 按列或按行切分，分配给不同的 GPU。每个 GPU 只计算矩阵的一部分，最后通过通信将结果拼接起来。

2. Megatron-LVLM 的核心魔法：1F1B

在经典的 Transformer 结构中，张量并行主要应用在两个地方：

多头注意力机制（MHA）：天然适合切分。如果有 16 个注意力头和 4 张 GPU，每张 GPU 只需计算 4 个头的输出。
MLP（多层感知机）：包含两个线性层 $Y = GeLU(XA)$ 和 $Z = YB$ 。

列并行：
我们将矩阵 $A$ 按列切分为 $[A_1, A_2]$ 。GPU 1 计算 $Y_1 = XA_1$ ，GPU 2 计算 $Y_2 = XA_2$ 。注意，此时 GPU 1 和 GPU 2 都需要完整的输入 $X$ 。
行并行：
我们将矩阵 $B$ 按行切分为 $\begin{bmatrix} B_1 \\ B_2 \end{bmatrix}$ 。GPU 1 计算 $Z_1 = Y_1B_1$ ，GPU 2 计算 $Z_2 = Y_2B_2$ 。此时，最终结果 $Z = Z_1 + Z_2$ 。

Megatron 的精髓在于：前向传播不需要通信，反向传播不需要通信！
在 MLP 的前向传播中，列并行不需要通信；行并行最后只需要一个 All-Reduce 求和。这种精妙的数学设计使得张量并行的通信量降到了最低。

3. PyTorch 手动实现张量并行示例

这里展示如何将一个巨大的 nn.Linear 层切分到两张 GPU 上：

import torch
import torch.nn as nn

# 假设我们有一个巨大的 4096x4096 线性层
class ColParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        # 按列切分，每张卡只初始化 out_features / world_size 列
        self.weight = nn.Parameter(torch.empty(in_features, out_features // world_size))
        self.bias = nn.Parameter(torch.empty(out_features // world_size))
        
    def forward(self, x):
        # x: [batch, in_features]
        # output: [batch, out_features / world_size]
        return torch.matmul(x, self.weight) + self.bias

class RowParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        # 按行切分，每张卡只初始化 in_features / world_size 行
        self.weight = nn.Parameter(torch.empty(in_features // world_size, out_features))
        self.bias = nn.Parameter(torch.empty(out_features))
        
    def forward(self, x):
        # x 是局部的: [batch, in_features / world_size]
        # output: [batch, out_features]
        return torch.matmul(x, self.weight) # 注意：Bias 需要在 All-Reduce 之后再加

# 模拟在两张卡上的张量并行前向传播
batch = 32
in_f, out_f = 4096, 4096

x = torch.randn(batch, in_f).cuda(0) # 假设输入完整存在于两张卡上

# GPU 1 上的操作 (模拟)
col_linear1 = ColParallelLinear(in_f, out_f, world_size=2).cuda(0)
y1 = col_linear1(x.cuda(0)) # 得到一半的输出特征

# GPU 2 上的操作 (模拟)
col_linear2 = ColParallelLinear(in_f, out_f, world_size=2).cuda(1)
y2 = col_linear2(x.cuda(1)) # 得到另一半的输出特征

# 此处通常需要 All-Gather 或者进入 Row 层计算
# 假设直接拼接 (相当于切分了神经元)
y_parallel = torch.cat([y1, y2], dim=-1)

(注：在实际生产环境中，强烈建议直接使用 Megatron-LM 或 PyTorch 的 device_mesh 及 TP 优化器，无需手写这些底层切分逻辑。)

三、流水线并行：赋予网络“工厂流水线”的节奏

1. 核心思想

如果说张量并行是“横向”切分，那么流水线并行（PP）就是**“纵向”切分**。

流水线并行按照神经网络的层将模型划分为不同的阶段，分配给不同的 GPU。例如，一个 100 层的 Transformer，GPU 0 负责第 1-25 层，GPU 1 负责第 26-50 层，以此类推。

这就像工厂的流水线：数据先进入 GPU 0 完成第一阶段的计算，然后将中间激活值发送给 GPU 1，GPU 1 计算完后发给 GPU 2。

2. 痛点与破局：冒泡与微批次

最原始的流水线并行（朴素 PP）存在一个严重的问题：流水线气泡。

在朴素 PP 中，GPU 0 在计算时，GPU 1、GPU 2、GPU 3 都在闲置等待；当 GPU 1 计算时，GPU 0、2、3 也在闲置。这种串行的等待导致算力利用率（MFU）极低。

为了解决这个问题，研究人员引入了 微批次 技术。
我们将原本较大的一个 Batch 切分成多个 Micro-batch。当 GPU 0 处理完 Micro-batch 1 并将其传给 GPU 1 时，GPU 0 并不闲着，而是立刻开始处理 Micro-batch 2。

GPipe 与 1F1B 策略：

GPipe：先进行所有 Micro-batch 的前向传播，再统一进行反向传播。这种方式虽然减少了气泡，但需要缓存大量前向激活值，对显存消耗极大。
1F1B (One Forward, One Backward)：这是目前最主流的流水线调度策略。在启动期过后，每张 GPU 严格交替执行一次前向传播和一次反向传播。这样不仅减小了气泡，还能及时释放反向传播产生的显存，实现了显存和计算效率的平衡。

3. 流水线并行的逻辑示意

虽然 PyTorch 提供了 torch.distributed.pipeline，但其实现较复杂。这里用一段伪代码来帮助你理解 1F1B 的逻辑调度：

# 假设有 4 张 GPU (Stage 0, 1, 2, 3)，将模型按层切分
# 分配了 8 个 Micro-batches (m0, m1... m7)

def pipeline_1f1b_schedule():
    # 启动阶段
    # Stage 0 执行前向传播，处理 m0, m1, m2
    # Stage 1 处理 m0, m1
    # Stage 2 处理 m0
    # ...
    
    # 稳态阶段 (1F1B)
    # 一旦某个 Stage 收到了来自上一层的激活值，并且其后面还有待处理的数据
    # 它会严格执行：
    # 1. 执行一次前向传播
    # 2. 等待并接收来自下一层的梯度
    # 3. 执行一次反向传播
    # 4. 将梯度发送给上一层
    
    # 冷却阶段
    # 处理最后几个 Micro-batch 的反向传播
    pass

建议在实际工程中，使用 DeepSpeed 或 Megatron-LM 配置文件直接指定 pipeline_parallel_size 即可。

四、终极杀招：3D 混合并行

在训练诸如 GLM-5、GPT-4 级别的万亿参数模型时，单纯依靠上述任何一种策略都无法完美解决问题。因此，现代大模型训练均采用 3D 混合并行 技术。

如何组合这三种策略？

假设我们有一个拥有 1024 张 GPU 的超级集群：

张量并行（TP）：由于 TP 在前向和反向传播中需要频繁进行 All-Reduce 通信（每次前向/反向都要通信），通信量极大，极度依赖网络带宽。因此，TP 通常只在同一个物理节点内（如一台包含 8 张 GPU 的服务器）进行，因为节点内的 NVLink 带宽高达 600GB/s 以上。我们将 TP size 设为 8。
流水线并行（PP）：PP 的通信是点对点（P2P）的，只需要发送相邻层的中间激活值，通信量较小。因此，PP 可以跨节点进行。我们跨 4 个节点组成一个流水线，将 PP size 设为 4。
数据并行（DP）：此时，我们已经用一个规模为 32 (8x4) 的 GPU 组合完全装下了一个庞大的模型。剩下的 1024 / 32 = 32 个模型副本，我们使用数据并行（结合 ZeRO 优化）在全局进行同步。DP 的通信（梯度 All-Reduce）发生在每次迭代结束时，对网络要求适中。

通过这种 3D 混合并行架构，我们可以完美地兼顾显存限制、计算效率与通信带宽，榨干每一块 GPU 的性能。

五、额外的性能优化利器

除了上述三大并行策略，在实际的大模型训练代码库中，你一定还会听到以下两个关键技术：

选择性激活重计算
即使使用了 3D 并行，前向传播产生的激活值依然会占据恐怖的显存。为了“用算力换显存”，研究人员提出在反向传播时，丢弃一部分不重要的前向激活值，需要时重新计算一次。这能节省约 60% 的激活显存，仅仅增加 15% 左右的计算开销，是现代大模型训练的标配。
序列并行
在 Transformer 的 LayerNorm 和 Dropout 阶段，张量并行是无法处理这些操作的（它们需要完整的输入序列）。Megatron-LM 提出了序列并行，沿着序列长度维度继续切分，进一步降低了非 TP 区域的显存占用。

总结

大模型的分布式训练不仅是一门算法艺术，更是一项极具挑战的系统工程。理解并合理配置分布式策略，是突破大模型训练瓶颈的关键所在。让我们用一张表格来做个快速的总结：

策略	切分维度	通信频率与类型	适用场景	典型代表
数据并行 (DP/ZeRO)	数据	每次迭代/ All-Reduce	模型能装进单卡，需要加速训练	PyTorch DDP, DeepSpeed ZeRO
张量并行 (TP)	权重矩阵 (横向)	极高 / All-Reduce	模型某层过大，单卡装不下	Megatron-LM, Tensor Parallel
流水线并行 (PP)	网络层数 (纵向)	较低 / 点对点 (P2P)	模型极深，突破节点显存极限	GPipe, PipeDream (1F1B)

未来，随着 MoE（混合专家模型）架构和长上下文技术的普及，专家并行 和 上下文并行 正在成为新的研究热点。但万变不离其宗，只要你深刻理解了今天所讲的 DP、TP 和 PP 的物理意义，就能游刃有余地应对未来更加复杂的分布式训练架构。

引言：打破单卡的“物理枷锁”

一、 数据并行：最简单粗暴的“抄作业”

1. 核心思想

2. 从 DP 到 DDP 的进化

3. 进阶：ZeRO (Zero Redundancy Optimizer)

4. PyTorch DDP 代码实战

二、 模型并行（张量并行）：解剖模型的“手术刀”

1. 核心思想

2. Megatron-LVLM 的核心魔法：1F1B

3. PyTorch 手动实现张量并行示例

三、 流水线并行：赋予网络“工厂流水线”的节奏

1. 核心思想

2. 痛点与破局：冒泡与微批次

3. 流水线并行的逻辑示意

四、 终极杀招：3D 混合并行

五、 额外的性能优化利器

总结

一、数据并行：最简单粗暴的“抄作业”

二、模型并行（张量并行）：解剖模型的“手术刀”

三、流水线并行：赋予网络“工厂流水线”的节奏

四、终极杀招：3D 混合并行

五、额外的性能优化利器