系统第 7 讲 / 共 17 讲

第 7 讲：Parallelism（上）

随着语言模型规模的爆发式增长，单个 GPU 已经无法满足训练需求。GPT-3 拥有 1750 亿参数，LLaMA-70B 有 700 亿参数——即便是最先进的 NVIDIA H100（80GB 显存），也无法装下这些模型的全部状态。本讲将系统介绍数据并行（Data Parallelism）的原理与实现，以及如何通过 ZeRO 优化来进一步降低显存占用。

为什么需要并行训练？

在深入并行策略之前，我们需要理解训练一个大模型时 GPU 上到底存了什么。以一个 N 参数的模型为例，使用 mixed precision（FP16/BF16 参数 + FP32 master weights）和 AdamW optimizer：

组件	精度	每参数字节数	说明
Parameters (FP16)	FP16	2 bytes	前向传播使用的参数
Gradients (FP16)	FP16	2 bytes	反向传播产生的梯度
Master Weights (FP32)	FP32	4 bytes	优化器更新用的全精度参数
AdamW First Moment (m)	FP32	4 bytes	梯度的一阶矩估计
AdamW Second Moment (v)	FP32	4 bytes	梯度的二阶矩估计

Total Memory per GPU = 2N + 2N + 4N + 4N + 4N = 16N bytes

也就是说，对于一个 7B 参数的模型，仅模型状态就需要：

16 × 7 × 10⁹ = 112 GB

这还没有算上 activation memory（中间激活值）！一块 80GB 的 H100 根本装不下。因此，我们必须使用多 GPU 并行训练。

核心问题

并行训练的根本挑战有两个：（1）显存限制——单 GPU 装不下模型状态；（2）计算效率——训练时间太长，需要用更多硬件来加速。不同的并行策略解决的侧重点不同。

Data Parallelism 基本原理

Data Parallelism（数据并行）是最简单也是最常用的并行策略。它的核心思想是：

每个 GPU 上放一份完整的模型副本（model replica）
将一个 mini-batch 的数据均匀分割给各个 GPU
每个 GPU 独立进行 forward 和 backward 计算
所有 GPU 通过通信同步梯度（gradient synchronization）
每个 GPU 用相同的聚合梯度更新自己的模型参数

由于每个 GPU 拿到的数据不同，它们计算出的梯度也不同。但通过梯度同步，所有 GPU 的参数在每一步之后都保持一致——从数学上讲，这等价于使用一个更大的 batch size 进行单 GPU 训练。

数学等价性

假设我们有 W 个 worker（GPU），global batch size 为 B，每个 worker 处理 B/W 个样本。对于 worker k，它计算的局部梯度为：

g_k = (1 / (B/W)) × Σ_{i ∈ S_k} ∇L(x_i, θ)

其中 S_k 是分配给 worker k 的数据子集。聚合后的梯度为：

g = (1/W) × Σ_{k=1}^{W} g_k = (1/B) × Σ_{i=1}^{B} ∇L(x_i, θ)

这与单 GPU 在整个 batch 上计算的梯度完全相同。

AllReduce 通信原语

Data Parallelism 的关键在于梯度同步。所有 GPU 需要将各自的梯度汇总，然后让每个 GPU 都拿到相同的聚合结果。这个操作叫做 AllReduce。

AllReduce 可以看作两步操作的组合：

Reduce：将所有 GPU 上的梯度求和，结果存到一个 GPU 上
Broadcast：将结果从那个 GPU 广播到所有其他 GPU

但在实际实现中，这种"先 Reduce 后 Broadcast"的方式效率很低——所有数据都要经过一个节点，形成通信瓶颈。

常见的 Collective Communication 原语

原语	功能	输入	输出
Broadcast	一对多分发	一个 GPU 有数据	所有 GPU 都有相同数据
Reduce	多对一聚合	每个 GPU 有数据	一个 GPU 有聚合结果
AllReduce	多对多聚合	每个 GPU 有数据	每个 GPU 都有聚合结果
Scatter	一对多分块发送	一个 GPU 有完整数据	每个 GPU 有一块
Gather	多对一收集	每个 GPU 有一块	一个 GPU 有完整数据
AllGather	多对多收集	每个 GPU 有一块	每个 GPU 有完整数据
ReduceScatter	聚合后分块	每个 GPU 有完整数据	每个 GPU 有一块聚合结果

Ring AllReduce 算法

Ring AllReduce 是目前使用最广泛的 AllReduce 实现。它的核心思想是把 GPU 排成一个环形拓扑，分两个阶段完成操作：

阶段一：ReduceScatter

将梯度向量分成 W 块（W 为 GPU 数量）。在 W-1 步通信中，每个 GPU 向环中的下一个 GPU 发送一块数据，同时从上一个 GPU 接收一块数据并累加。完成后，每个 GPU 持有一块完整的聚合结果。

阶段二：AllGather

再经过 W-1 步，每个 GPU 将自己持有的聚合结果块传递给环中的下一个 GPU。完成后，每个 GPU 都持有完整的聚合梯度。

通信量分析

假设梯度总大小为 M bytes：

ReduceScatter 阶段：每步每个 GPU 发送 M/W，共 W-1 步，总发送 M(W-1)/W
AllGather 阶段：同样每 GPU 发送 M(W-1)/W
总通信量：每 GPU 发送 2M(W-1)/W ≈ 2M（当 W 较大时）

为什么 Ring AllReduce 高效？

Ring AllReduce 的通信量与 GPU 数量 W 几乎无关（渐近为 2M），而 naive 的"全部发给一个节点再广播回来"方案的瓶颈节点通信量为 W×M。Ring AllReduce 实现了带宽最优——每条链路的利用率都接近 100%。

通信时间估算

Ring AllReduce 的总通信时间可以估算为：

T_comm = 2(W-1) × α + 2 × M(W-1)/(W × β)

其中 α 是 latency（延迟），β 是 bandwidth（每条链路带宽）。当消息足够大时（M >> α×β），latency 项可以忽略，通信时间近似为 2M/β。

Distributed Data Parallel (DDP) 实现

PyTorch 提供了 DistributedDataParallel（DDP）作为 Data Parallelism 的标准实现。DDP 使用 NCCL 后端进行 GPU 间通信，性能非常优秀。

DDP 的核心实现细节

1. Gradient Bucketing（梯度分桶）

DDP 不会等所有梯度都计算完才开始通信。它将参数分成多个 bucket，当一个 bucket 内所有参数的梯度都计算完成时，立即启动该 bucket 的 AllReduce。这样通信和计算可以重叠（overlap），大幅减少总训练时间。

2. Backward Hook

DDP 在每个参数上注册 autograd hook。当一个参数的梯度在 backward pass 中被计算出来时，hook 会通知 DDP 系统，标记该参数所在 bucket 的进度。

3. Bucket 的构建顺序

Bucket 按照参数在 backward pass 中被计算的反序来构建。这样最先计算完梯度的参数在同一个 bucket 中，可以最早开始通信。

DDP 使用示例

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# Initialize process group
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

# Create model and wrap with DDP
model = MyTransformer().cuda(local_rank)
model = DDP(model, device_ids=[local_rank])

# Create distributed sampler
sampler = torch.utils.data.distributed.DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler, batch_size=local_batch_size)

# Training loop - looks identical to single GPU!
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for epoch in range(num_epochs):
    sampler.set_epoch(epoch)  # Ensure different shuffling per epoch
    for batch in dataloader:
        batch = batch.cuda(local_rank)
        loss = model(batch)
        loss.backward()       # AllReduce happens here automatically
        optimizer.step()
        optimizer.zero_grad()

启动 DDP 训练的命令：

# Launch with torchrun (recommended)
torchrun --nproc_per_node=8 train.py

# Multi-node: 2 nodes, each with 8 GPUs
torchrun --nnodes=2 --nproc_per_node=8 \
    --rdzv_backend=c10d --rdzv_endpoint=node0:29400 \
    train.py

Gradient Accumulation

当我们想要使用更大的 effective batch size，但 GPU 显存不够时，可以使用 Gradient Accumulation（梯度累积）。

核心思想：把一个大的 mini-batch 分成多个 micro-batch，依次 forward + backward，把梯度累加起来，最后再做一次参数更新。

accumulation_steps = 4  # Effective batch = local_batch * accumulation_steps * world_size

for i, batch in enumerate(dataloader):
    batch = batch.cuda()
    loss = model(batch) / accumulation_steps  # Scale loss

    # Only sync gradients on the last micro-batch
    if (i + 1) % accumulation_steps != 0:
        with model.no_sync():  # Skip AllReduce
            loss.backward()
    else:
        loss.backward()        # AllReduce happens here
        optimizer.step()
        optimizer.zero_grad()

关键技巧：no_sync()

DDP 的 no_sync() context manager 会跳过 AllReduce。在累积阶段使用它可以避免不必要的通信开销。只有在最后一个 micro-batch 的 backward 时才需要同步梯度。

Gradient Accumulation 与 Data Parallelism 的关系

假设有 W 个 GPU，每个 GPU 的 local batch size 为 b，gradient accumulation steps 为 A，那么：

Effective Batch Size = W × b × A

Gradient Accumulation 是一种用时间换空间的方法——它不需要额外的 GPU，但训练时间会相应延长。相比之下，Data Parallelism 是用空间换时间——用更多 GPU 来加速训练。两者通常结合使用。

ZeRO: Zero Redundancy Optimizer

Data Parallelism 有一个明显的问题：每个 GPU 都持有完整的模型状态。对于一个 7B 模型，每个 GPU 都需要约 112GB 的显存来存储模型参数、梯度和 optimizer states。这是巨大的冗余——这些副本内容完全相同。

ZeRO（Zero Redundancy Optimizer）由 Microsoft DeepSpeed 团队提出，核心思想是：将模型状态分片（shard）到不同 GPU 上，需要时再通过通信聚合。

ZeRO 分为三个阶段，逐步消除更多冗余：

ZeRO Stage 1: Optimizer State Sharding

Optimizer states（AdamW 的 m 和 v 以及 FP32 master weights）占 12N bytes，是内存占用的大头。Stage 1 将 optimizer states 均匀分片到 W 个 GPU 上：

Stage 1 Memory = 2N + 2N + 12N/W = 4N + 12N/W bytes per GPU

当 W = 8 时，每 GPU 内存从 16N 降到 5.5N，节省约 65%。

实现方式：每个 GPU 只负责更新自己分片对应的参数。在每步训练中：

所有 GPU 正常做 forward 和 backward（每个 GPU 有完整参数和梯度）
通过 ReduceScatter 将梯度聚合并分片
每个 GPU 用自己分片的 optimizer state 更新对应的参数
通过 AllGather 将更新后的参数广播到所有 GPU

ZeRO Stage 2: Gradient Sharding

在 Stage 1 的基础上，进一步将梯度也分片。既然每个 GPU 只更新自己分片的参数，那么它也只需要自己分片的梯度。

Stage 2 Memory = 2N + 2N/W + 12N/W = 2N + 14N/W bytes per GPU

当 W = 8 时，每 GPU 内存从 16N 降到 3.75N，节省约 76%。

关键实现：在 backward pass 中，使用 ReduceScatter 替代 AllReduce。这样每个 GPU 只保留自己需要的那份梯度。

ZeRO Stage 3: Parameter Sharding

最激进的方案——连模型参数也分片。每个 GPU 只存储 1/W 的参数。当 forward 或 backward 需要某层的参数时，通过 AllGather 临时从其他 GPU 收集完整参数，用完后立即释放。

Stage 3 Memory = 2N/W + 2N/W + 12N/W = 16N/W bytes per GPU

当 W = 8 时，每 GPU 内存仅需 2N，节省约 87.5%。

ZeRO 各阶段对比

阶段	分片内容	每 GPU 内存	通信量（vs DDP）
DDP (baseline)	无	16N	1× (AllReduce)
ZeRO Stage 1	Optimizer States	4N + 12N/W	1× (等效)
ZeRO Stage 2	+ Gradients	2N + 14N/W	1× (等效)
ZeRO Stage 3	+ Parameters	16N/W	1.5× (额外 AllGather)

ZeRO Stage 选择建议

Stage 1：几乎没有额外通信开销，推荐作为默认选择。Stage 2：通信开销与 DDP 相同，进一步节省显存。Stage 3：通信量增加约 50%，但能训练超大模型。通常在模型参数量超过单 GPU 显存时才使用 Stage 3。

ZeRO 与 FSDP

PyTorch 原生提供的 FSDP（Fully Sharded Data Parallel）本质上就是 ZeRO Stage 3 的实现。它在 PyTorch 2.0+ 中成为推荐的分布式训练方案：

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy

# Wrap model with FSDP
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO Stage 3
    # ShardingStrategy.SHARD_GRAD_OP = ZeRO Stage 2
    # ShardingStrategy.NO_SHARD = DDP
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    ),
    device_id=local_rank,
)

# Training loop is the same as DDP
for batch in dataloader:
    loss = model(batch)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

通信带宽分析与 Overlap

在实际训练中，Data Parallelism 的效率取决于计算时间与通信时间的比值。理想情况下，通信可以完全隐藏在计算中（communication-computation overlap）。

典型互连带宽

互连类型	单向带宽	典型场景
NVLink (H100)	450 GB/s per GPU	同一节点内 GPU 间
PCIe Gen5 x16	64 GB/s	GPU 到 CPU / 跨 PCIe switch
InfiniBand NDR 400G	50 GB/s	跨节点通信
RoCE 100G	12.5 GB/s	数据中心以太网

举个具体例子：一个 7B 模型的梯度大小约为 14GB（FP16）。使用 Ring AllReduce 在 8 个 GPU 之间同步：

节点内（NVLink）：通信时间 ≈ 2 × 14GB / 450 GB/s ≈ 0.06s
跨节点（InfiniBand 400G）：通信时间 ≈ 2 × 14GB / 50 GB/s ≈ 0.56s

如果一步训练的计算时间约为 1 秒，那么节点内通信可以轻松被 overlap 掉，但跨节点通信就需要仔细优化。

Overlap 策略

DDP 的 gradient bucketing 实现了自动 overlap：在 backward pass 计算后面层的梯度时，前面层的梯度已经可以开始通信了。理想情况下，当最后一层的梯度计算完成时，前面所有层的通信也已经完成。

Effective Time = max(T_compute, T_communicate) ≈ T_compute（理想情况）

Data Parallelism 的局限性

Data Parallelism（包括 ZeRO 优化）有以下局限：

模型必须能装入单 GPU（至少在 forward pass 时）：即使是 ZeRO Stage 3，forward 时也需要通过 AllGather 在每个 GPU 上临时重建完整的层参数
Activation Memory 不被分片：每个 GPU 仍然需要存储自己 micro-batch 的全部激活值
通信开销随模型增大：梯度同步的数据量与参数量成正比
Batch Size 限制：增加 GPU 数量就等于增加 global batch size，但过大的 batch size 可能影响收敛

这些局限性促使了更高级的并行策略的发展——Tensor Parallelism 和 Pipeline Parallelism——它们将模型本身拆分到多个 GPU 上。我们将在下一讲详细介绍这些方法。

小结

Data Parallelism 是最基础的并行策略，每个 GPU 持有完整模型，数据均匀分割
AllReduce（特别是 Ring AllReduce）是梯度同步的核心通信原语，通信量约为 2M，与 GPU 数量无关
DDP 通过 gradient bucketing 实现计算与通信的 overlap
Gradient Accumulation 用时间换空间，实现更大的 effective batch size
ZeRO Stage 1/2/3 逐步消除 optimizer state / gradient / parameter 的冗余，最高可节省 87.5% 显存
FSDP 是 PyTorch 原生的 ZeRO Stage 3 实现