突破显存墙：万字长文带你吃透大模型训练的分布式策略（数据并行、模型并行与流水线并行）

Posted on In

引言：大模型时代的“显存墙”与算力焦虑

自从 ChatGPT 横空出世，大语言模型（LLM）便成为了整个科技界的“皇冠上的明珠”。然而，随着模型参数量从十亿（GPT-2）、千亿（GPT-3）一路狂飙至万亿（GPT-4、GLaMA），开发者们面临着一个极其现实的物理限制——单张 GPU 的显存容量远远装不下这些庞然大物。

哪怕是目前顶级的 NVIDIA H100/A100 显卡，其显存容量通常也只有 80GB。而一个拥有 1750 亿参数的 GPT-3 模型，仅加载权重（FP16 精度）就需要约 350GB 的显存，如果加上训练过程中的激活值、梯度以及优化器状态（如 Adam 的动量和方差），整体显存消耗轻松突破 TB 级别。

在这种情况下，如果还幻想着用单卡训练大模型，无异于痴人说梦。为了跨越这道“显存墙”，分布式训练成为了大模型时代唯一可行的方法论。

提到分布式训练，很多人的第一反应是“多加几台机器”。但分布式并不是简单的“堆卡”，它背后是一套极其精密的并行计算拓扑策略。当前大模型训练的分布式策略主要分为三大核心流派：数据并行、模型并行（特指张量并行）以及流水线并行。

本文将深入浅出地为你剖析这三大策略的底层原理、适用场景以及它们在实际工业级框架（如 PyTorch, DeepSpeed, Megatron-LM）中的代码实现。最后，我们将探讨目前大厂都在使用的“3D 混合并行”策略。准备好了吗？让我们推开大模型底层训练的大门。

一、 数据并行：最经典的“人多力量大”

1.1 核心思想

数据并行是最直观、最容易理解的分布式策略。它的核心思想是：在每个 GPU 上都复制一份完整的模型副本，然后将输入的数据集切分成多个 Batch，分发给不同的 GPU 进行独立计算。

举个例子：如果你有一个 Batch Size 为 32 的数据，且有 8 张 GPU。在数据并行下，每张 GPU 会拿到 Batch Size 为 4 的数据，并各自在自己完整的模型上进行前向传播和反向传播。计算完梯度后，所有 GPU 将梯度进行同步聚合（All-Reduce 操作），然后各自更新本地的模型权重，确保模型状态始终保持一致。

1.2 演进：从 DP 到 DDP 再到 FSDP/ZeRO

虽然思想简单，但在工程实现上，数据并行经历了三代演进：

  1. DP (Data Parallelism - Parameter Server 架构)：
    早期的 PyTorch 采用的是参数服务器架构。即有一张 GPU 作为 Master 负责收集梯度、更新权重，然后再把新权重广播给其他 Worker GPU。这会导致 Master 成为通信瓶颈。
  2. DDP (Distributed Data Parallel - Ring-AllReduce 架构)：
    目前最常用的方式。PyTorch 的 DistributedDataParallel 引入了 Ring-AllReduce 算法，取消了中心节点，所有 GPU 连成一个环，梯度在环内流动，使得通信开销与 GPU 数量解耦，效率极高。
  3. FSDP (Fully Sharded Data Parallel) / ZeRO：
    传统 DDP 的致命弱点在于：每张卡都要保存完整的模型参数、梯度和优化器状态。如果模型有 70 亿参数，哪怕用 1000 张卡，单卡依然得装下这 70 亿参数。
    微软 DeepSpeed 提出的 ZeRO (Zero Redundancy Optimizer) 算是彻底打破了这一限制（PyTorch 原生对应实现为 FSDP）。它的核心思想是“分片”：把模型参数、梯度和优化器状态切分到不同的 GPU 上。只有在前向/反向传播需要用到特定层时，才通过通信获取完整的参数，用完即丢。这使得单卡显存消耗与 GPU 数量成反比。

1.3 代码实战：PyTorch DDP 与 FSDP 示例

传统 DDP 实现片段：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_ddp():
    # 初始化进程组
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

def main_ddp():
    setup_ddp()
    
    # 构建模型并将其放到当前 GPU 上
    model = MyLargeModel().cuda()
    
    # 用 DDP 包装模型
    model = DDP(model, device_ids=[local_rank])
    
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    
    for data, target in dataloader:
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        
        # 反向传播时，DDP 会自动进行梯度的 All-Reduce 同步
        loss.backward() 
        optimizer.step()

PyTorch FSDP 实现片段：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

def main_fsdp():
    setup_ddp()
    
    # 构建模型（此时仅放在 CPU 或当前单卡上，无需全量装入 GPU）
    model = MyGiantModel()
    
    # 使用 FSDP 包装模型
    # ShardingStrategy.FULL_SHARD 相当于 ZeRO-Stage 3，全部分片
    model = FSDP(
        model, 
        device_id=local_rank,
        sharding_strategy=ShardingStrategy.FULL_SHARD 
    )
    
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    
    for data, target in dataloader:
        data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        
        # FSDP 在反向传播时自动进行参数的 Gather 和丢弃
        optimizer.step()

1.4 数据并行的局限性

虽然 ZeRO/FSDP 极大地拓展了模型规模的上限，但数据并行本质上依然存在一个软肋：极高的网络通信开销。在计算过程中，GPU 之间需要频繁交换梯度或参数分片。如果使用的是普通万兆以太网，网络带宽会瞬间成为系统的阿喀琉斯之踵。此外，当模型单层的参数量过大（比如巨大的 Embedding 层或超宽的 MLP 层）时，单张卡甚至无法装下单个层进行前向计算，这时仅靠数据并行就无能为力了。

二、 模型并行：拆解庞然大物

当我们说“模型太大，单卡放不下”时，最直接的思路就是把模型切开。这就引出了狭义上的模型并行（通常被称为张量并行 Tensor Parallelism, TP）。

2.1 核心思想

张量并行的核心在于：将模型单层内部的参数矩阵（张量）切分到多张 GPU 上进行计算。 它是一种细粒度的并行方式，发生在模型的前向和反向传播的数学运算内部。

现代深度学习模型的基础模块大多是矩阵乘法（GEMM）。例如在 Transformer 的前馈神经网络（FFN）中，计算公式为 Y=X×AY = X \times A。其中 XX 是输入矩阵，AA 是权重矩阵。

张量并行就是要把这个巨大的 AA 矩阵切开。

2.2 经典实现：Megatron-LM 的 1D 张量并行

NVIDIA 在 Megatron-LM 中提出了一种极其优雅的张量并行方案。以 Transformer 中的 FFN 层为例（包含两个全连接层 AA 和 BB）：

  1. 列切分：
    将第一个权重矩阵 AA 按列切分为 [A1,A2][A_1, A_2]，分别放在 GPU 0 和 GPU 1 上。
    输入 XX 会同时发送给两张 GPU。GPU 0 算出 Y1=X×A1Y_1 = X \times A_1，GPU 1 算出 Y2=X×A2Y_2 = X \times A_2。

  2. 行切分：
    将第二个权重矩阵 BB 按行切分为 [B1B2]\begin{bmatrix} B_1 \\ B_2 \end{bmatrix}。
    GPU 0 拿着刚才算出的 Y1Y_1 乘以 B1B_1（即 Z1=Y1×B1Z_1 = Y_1 \times B_1），GPU 1 同理算出 Z2=Y2×B2Z_2 = Y_2 \times B_2。

  3. All-Reduce 聚合：
    最后，只需要将两张卡的结果相加，即 Z=Z1+Z2Z = Z_1 + Z_2。这个加法在数学上与未切分时的 Z=Y×BZ = Y \times B 结果完全等价。
    这个相加的操作在分布式通信中被称为 All-Reduce（具体为 Sum 操作）。

2.3 代码实战：手写一个简单的张量并行线性层

通过下面的 PyTorch 伪代码，你可以直观感受到矩阵是如何被切分和合并的：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
import torch
import torch.nn as nn
import torch.distributed as dist

class ColumnParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        self.world_size = world_size
        self.rank = dist.get_rank()
        
        # 假设 out_features 可以被 world_size 整除
        self.out_features_per_partition = out_features // world_size
        
        # 每张卡只初始化真实权重 1/world_size 的大小（按列切分）
        self.weight = nn.Parameter(torch.empty(self.out_features_per_partition, in_features))
        # 初始化省略...

    def forward(self, x):
        # x: [batch, in_features]
        # output_parallel: [batch, out_features_per_partition]
        output_parallel = torch.matmul(x, self.weight.t())
        return output_parallel

class RowParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        self.world_size = world_size
        self.rank = dist.get_rank()
        
        self.in_features_per_partition = in_features // world_size
        
        # 按行切分权重
        self.weight = nn.Parameter(torch.empty(out_features, self.in_features_per_partition))
        # 初始化省略...

    def forward(self, x):
        # x 为上一层（列切分层）的输出，无需通信直接本地计算
        output_parallel = torch.matmul(x, self.weight.t())
        
        # All-Reduce 聚合所有 GPU 的结果
        dist.all_reduce(output_parallel, op=dist.ReduceOp.SUM)
        return output_parallel

# 模拟 Transformer 的 FFN: Y = GeLU(X * A) * B
class TPPFFN(nn.Module):
    def __init__(self, hidden_dim, ffn_dim, world_size):
        super().__init__()
        self.fc1 = ColumnParallelLinear(hidden_dim, ffn_dim, world_size)
        self.fc2 = RowParallelLinear(ffn_dim, hidden_dim, world_size)
        self.relu = nn.ReLU()

    def forward(self, x):
        # [Batch, hidden] -> [Batch, ffn_dim/world_size]
        x = self.fc1(x)
        x = self.relu(x)
        # [Batch, ffn_dim/world_size] -> All-Reduce -> [Batch, hidden]
        x = self.fc2(x)
        return x

2.4 张量并行的优缺点

  • 优点：极大地降低了单卡的显存占用，使得超宽的 Transformer 层得以计算；且由于 GPU 之间是高速互联（通常在同一个物理节点内），计算与通信可以做到一定程度的重叠。
  • 缺点：通信极其密集。在每次前向和反向传播时，都需要进行一次 All-Reduce 操作。如果节点之间不是 NVLink 这种超高带宽连接，张量并行的效率会暴跌。因此，张量并行通常只在单个机台内部（通常 8 张卡）使用。

三、 流水线并行：深度的接力赛

如果说张量并行是把模型“横向”切开，那么流水线并行就是把模型“纵向”切开。

3.1 核心思想

流水线并行将模型的不同层按顺序分配给不同的 GPU。就像工厂里的流水线一样：GPU 0 负责前 10 层的计算，GPU 1 负责第 11-20 层，GPU 2 负责第 21-30 层……

这种切分方式非常符合直觉，且能彻底解决模型深度扩展的问题。比如 GPT-3 有 96 层 Transformer，我们可以把它均分到 8 个节点上，每个节点只需处理 12 层。

3.2 致命问题：“气泡”

流水线并行听起来很完美，但实际操作中却面临着一个极其尴尬的问题：GPU 空闲（气泡现象）。

在最朴素的 GPipe 流水线策略中：

  1. 阶段一：GPU 0 输入数据（Micro-batch 1），算完前向传播，把激活值发给 GPU 1。此时 GPU 0 就停工了。
  2. 阶段二：GPU 1 拿到数据算完，发给 GPU 2。此时 GPU 0、GPU 1 都停工了。
  3. 阶段三：反向传播也是一样，必须等最后一层算完，梯度才能一层层往回传。

这就导致：在大部分时间里，只有一张 GPU 在疯狂计算，其他 GPU 全在“摸鱼”。这种由于等待造成的空闲时间，被称为流水线气泡。这造成了昂贵的 GPU 算力资源的极大浪费。

3.3 破局之道：微批次

为了减少气泡，工程师们借鉴了 CPU 指令流水线的思路，引入了微批次技术。

核心逻辑：不要一次性把一个巨大的 Mini-batch 喂给模型，而是把它拆分成多个更小的 Micro-batch（比如把 Batch Size 64 拆成 8 个 Micro-batch，每个 Size 为 8）。

当 GPU 0 处理完 Micro-batch 1 并将激活值发送给 GPU 1 后，GPU 0 并不空闲，而是立刻开始处理 Micro-batch 2 的前向传播。
这样，随着时间推移，GPU 0, 1, 2… 就像流水线上的工人一样，接连不断处理不同的子任务，完美衔接。

目前主流的流水线调度策略主要有两种：

  • GPipe（同步模式）：先执行所有 Micro-batch 的前向传播，再统一执行反向传播。这种方式实现简单，且在数学上与单卡训练完全等价，但仍然存在一定的气泡。
  • PipeDream（1F1B 模式）：即 One Forward, One Backward。一张 GPU 做完一次前向传播，如果刚好有前方的 GPU 传回来的梯度，它就会立刻做一次反向传播。这种方式极大降低了显存峰值（不需要保存所有 Micro-batch 的激活值），并将气泡降到了最低，是当前工业界的主流。

3.4 代码实战：PyTorch Pipeline 并行示例

PyTorch 在 1.8 版本后引入了对流水线并行的原生支持（基于 torch.distributed.pipeline，较新版本结合 PiPPy 项目）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
import torch
import torch.nn as nn
import torch.distributed as dist
import os
from torch.distributed.pipeline.sync import Pipe

# 1. 初始化进程组
dist.init_process_group(backend="nccl")
rank = dist.get_rank()
world_size = dist.get_rank()
torch.cuda.set_device(rank)

# 2. 定义一个巨大的模型（例如由 8 个 Transformer 层组成）
class GiantTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([TransformerLayer() for _ in range(8)])
        
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

model = GiantTransformer()

# 3. 切割模型！
# 假设我们有 2 张 GPU (world_size = 2)
# 我们将前 4 层放在 GPU 0，后 4 层放在 GPU 1
# 注意：实际生产中多使用 balanced partition 策略按显存占用均分
model = nn.Sequential(*list(model.layers.children()))
if rank == 0:
    model = model[:4].cuda() # GPU 0 获取前 4 层
else:
    model = model[4:].cuda() # GPU 1 获取后 4 层

# 4. 转换为流水线模型
# chunks=8 意味着我们将一个 Batch 切分为 8 个 Micro-batch (微批次)
pipe_model = Pipe(model, chunks=8)

# 5. 模拟训练
optimizer = torch.optim.Adam(pipe_model.parameters(), lr=1e-4)

for data, target in dataloader:
    data, target = data.cuda(), target.cuda()
    
    # 流水线前向传播（只在不同阶段间传递激活值）
    output = pipe_model(data)
    
    # 计算损失并反向传播
    loss = loss_fn(output, target)
    
    # 流水线内部会处理反向传播的梯度传递
    pipe_model.backward(loss)
    
    optimizer.step()
    optimizer.zero_grad()

3.5 流水线并行的优缺点

  • 优点：可以训练无限深度的模型；层与层之间的通信只在相邻的 GPU 上发生，不需要全局同步（如 All-Reduce），对网络带宽的要求相对较低，非常适合跨节点（跨服务器）训练。
  • 缺点：无论怎么优化，流水线气泡始终存在，无法做到 100% 的计算资源利用率；并且它要求模型的深度必须能被合理划分，如果某些层特别宽，还会导致各个 GPU 间的负载不均衡。

四、 3D 混合并行：大模型的终极形态

读到这里，你可能已经发现了：数据并行（DP）、张量并行（TP）、流水线并行（PP）各有千秋，但单独使用都无法完美驾驭成百上千台服务器和万亿参数的模型。

于是，工业界（如 Megatron-Deepspeed 联合团队）提出了终极武器：3D 混合并行。

4.1 组合的艺术：如何分配 GPU？

假设我们的计算集群有 64 张 GPU（8 台服务器，每台 8 张 GPU），我们要训练一个拥有 5000 亿参数的巨兽。我们应该怎么组合这三种策略呢？

核心原则是根据硬件的物理拓扑结构（网络带宽特性）来分配：

  1. 张量并行（TP）—— 极速的节点内通信
    由于张量并行需要极其频繁地传递激活值（每次前向/反向都要 All-Reduce），它要求极高的通信带宽。因此，TP 必须被限制在单台服务器（节点）内部，利用服务器内的 NVLink/NVSwitch 进行通信。
    例如：我们将每台服务器内的 8 张 GPU 组成一个 TP 组（TP Degree = 8）。

  2. 流水线并行（PP）—— 跨节点的层间接力
    当模型大到单机放不下时，我们需要把模型的不同深度阶段分配给不同的服务器。流水线并行只需要相邻阶段（服务器）之间点对点（P2P）传递激活值和梯度，对带宽要求中等。
    例如：我们将模型切分为 4 个 Stage。这样我们需要 4 个节点来组成一个完整的流水线（PP Degree = 4）。此时共消耗了 4 * 8 = 32 张 GPU，组成了一条完整的流水线（一个完整的模型副本）。

  3. 数据并行（DP / ZeRO）—— 全局的数据分片
    当单个模型副本训练速度太慢，或者想要进一步通过 ZeRO 降低显存时，我们就横向扩展流水线的数量。数据并行需要全局的梯度同步（All-Reduce），可以在集群的交换机网络上进行。
    例如：我们有 64 张 GPU，一个模型副本消耗 32 张 GPU，那么我们可以开 2 个数据并行组（DP Degree = 2）。

4.2 3D 并行的通信与计算拓扑

在上述配置下（TP=8, PP=4, DP=2）：

  • GPU 0 到 7 负责模型第 1 阶段的前向/反向传播，它们通过节点内 NVLink 进行张量并行通信。
  • GPU 0-7（Stage 1） 计算完后，通过以太网/InfiniBand 将数据发送给 GPU 8-15（Stage 2），形成流水线。
  • 上述 32 张 GPU 组成了模型副本 1。与此同时，另外 32 张 GPU（GPU 32-63）组成了完全相同的模型副本 2。它们在每次迭代后，通过全局网络进行梯度的同步。

4.3 为什么这是大模型训练的“唯一解”？

  1. 突破内存限制：ZeRO 切碎了优化器状态和参数，TP 切碎了单层权重，PP 切碎了模型深度。三管齐下，无论是多宽、多深的模型，都可以被装入集群。
  2. 计算效率最大化：将最耗时的 TP 限制在机箱内部（NVLink 带宽可达 900GB/s），将带宽要求低的 PP 和 DP 放在机箱外部（InfiniBand 带宽通常为 400Gb/s），充分利用了异构网络的物理特性，把计算与通信的比例压榨到了极致。
  3. 弹性扩展：只要网络交换机足够强，理论上你可以无限增加 DP 的数量来加快训练速度，或者增加 PP 的深度来扩大模型参数规模。

五、 总结与展望：通往大模型圣杯的基石

大模型的军备竞赛已经从单纯拼模型架构（Transformer, MoE），演变成了底层系统工程的较量。理解分布式训练的策略，是每一个想深入大模型底层、甚至做 LLM 基建工程师的必修课。

让我们用最简短的话总结本文的核心：

  • 数据并行 (DP/ZeRO/FSDP)：最容易上手，每张卡一份模型（或分片），切分数据。适合参数规模中等，但想要加速训练的场景。
  • 张量并行 (TP)：切分单层权重矩阵。通信极高频，必须在单机内部（NVLink）使用。解决“单层太宽”的问题。
  • 流水线并行 (PP)：切分模型深度。通过微批次减少空闲气泡。解决“模型太深”的问题，适合跨节点使用。
  • 3D 混合并行：工业界的事实标准。TP 榨干单机算力，PP 拓展模型深度，DP (ZeRO) 扩展数据规模并降低显存。

随着技术的演进，我们还在看到诸如 序列并行（Sequence Parallelism，针对超长上下文）、多维混合专家 的分布式策略不断涌现。此外，硬件层面上，NVIDIA GB200 NVL72 机架的推出，正试图用极度变态的机架内 NVLink 连接（72 张 GPU 共享显存池）来掩盖分布式通信的复杂性。

但万变不离其宗，无论硬件如何发展，只要模型的规模还在呈指数级增长，理解这些分布式切分和通信的数学与工程逻辑，将永远是我们把握大模型底层脉搏的核心锚点。

希望这篇文章能为你揭开大模型训练黑盒的一角，让你在面对动辄成百上千卡集群的参数配置时，不再迷茫，胸有成竹。