📚 课程目录第2章数据并行入门

2. 数据并行入门

2.1 数据并行：最朴素的分布式思路

数据并行，说白了就是「把数据切开，分给多张卡一起算」。

我刚开始接触分布式训练时，第一反应就是：一张卡跑不动，那就多上几张卡呗。但怎么分？怎么合？这里头有门道。

数据并行的核心逻辑很简单：

每张卡都有一份完整的模型副本
把训练数据切成若干份，每张卡拿一份
各自算各自的梯度
把梯度汇总，更新模型

你想想看，如果一张卡处理一个 batch 需要 1 秒，那 4 张卡理论上就能在 1 秒内处理 4 个 batch。当然，这是理想情况。实际中因为通信开销，会有一些损耗。

关键点：数据并行不改变模型结构，只改变数据流向。模型有多大，每张卡就得装多大。

2.2 AllReduce 算法：梯度的「快递系统」

数据并行里最核心的问题就是：梯度怎么汇总？

每张卡算完自己的梯度后，需要把梯度「凑」到一起，求个平均。这个「凑」的过程，就是 AllReduce 干的活。

我最早用 PyTorch 时，以为梯度汇总就是简单的「把梯度发到一张卡，算完再广播回去」。后来发现，这种做法叫「参数服务器」，效率其实不高。

AllReduce 有几种经典实现：

算法	思路	通信量	我的评价
Ring AllReduce	卡与卡组成环，数据像接力棒一样传递	2(N-1)数据量	最常用，带宽利用率高
Tree AllReduce	树形结构，逐层归约	2log2(N)数据量	适合小规模集群
Butterfly AllReduce	两两交换，并行度高	log2(N)*数据量	延迟低，但实现复杂

Ring AllReduce 是我用得最多的。它分两步走：

Scatter-Reduce 阶段：每张卡把自己的梯度切成 N 份（N 是卡数），然后沿着环传递。每传一次，累加一次。最后每张卡都有一份完整的累加结果，但只针对其中一块数据。
AllGather 阶段：把累加好的结果广播给所有卡。每张卡最终拿到完整的平均梯度。

嗯，这里要注意：Ring AllReduce 的通信量跟卡数无关，只跟数据量有关。这意味着卡越多，相对效率越高。我在项目中遇到过 64 卡训练，Ring AllReduce 的表现比参数服务器好了一个数量级。

我的经验：如果你用 PyTorch DDP，底层默认就是 Ring AllReduce（NCCL 后端）。你不需要手动实现，但理解它的原理对调优很有帮助。

2.3 PyTorch DDP：开箱即用的数据并行

PyTorch 的 DistributedDataParallel（DDP）是我最推荐的分布式训练工具。它封装了数据并行的所有细节，你只需要改几行代码。

为什么不用 DataParallel（DP）？我刚开始也踩过这个坑。DP 是单机多卡，但它的实现有缺陷：

主卡（GPU:0）要负责汇总和广播，容易成为瓶颈
每张卡的计算结果都要经过主卡，通信开销大
不支持多机

DDP 就不一样了。它用 AllReduce 做梯度同步，每张卡地位平等，没有主卡瓶颈。而且支持多机多卡，扩展性很好。

来看一个最简单的 DDP 示例：

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

def train(rank, world_size):
    # 初始化进程组
    dist.init_process_group(
        backend='nccl',
        init_method='tcp://127.0.0.1:23456',
        rank=rank,
        world_size=world_size
    )
    
    # 创建模型，放到对应 GPU
    model = MyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    # 创建数据加载器（每个进程拿不同数据）
    dataset = MyDataset()
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset, num_replicas=world_size, rank=rank
    )
    dataloader = DataLoader(dataset, sampler=sampler, batch_size=32)
    
    optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.01)
    
    for epoch in range(10):
        sampler.set_epoch(epoch)  # 打乱数据
        for data, target in dataloader:
            data, target = data.to(rank), target.to(rank)
            
            optimizer.zero_grad()
            output = ddp_model(data)
            loss = loss_fn(output, target)
            loss.backward()
            optimizer.step()
    
    dist.destroy_process_group()

if __name__ == '__main__':
    world_size = 4
    mp.spawn(train, args=(world_size,), nprocs=world_size)

这段代码里，我特别想强调几个点：

DistributedSampler：它负责给每个进程分配不同的数据。如果不加这个，所有卡都在算同一批数据，那就白忙活了。
set_epoch：每个 epoch 都要调用一次，否则数据打乱顺序会一样。
device_ids：指定模型放在哪张卡上。DDP 要求每张卡有自己的模型副本。

我曾经踩过的坑：DDP 的 batch_size 是每张卡的 batch_size。如果你设了 batch_size=32，4 张卡实际的总 batch_size 是 128。学习率也要相应调整，一般按卡数线性缩放。

2.4 数据并行的知识体系

下面这张图是我自己梳理的数据并行核心逻辑，你看一眼就能明白整体脉络：

从这张图你能看到，数据并行的流程其实很清晰：数据分片 → 各自前向/反向 → AllReduce 同步梯度 → 更新模型。每一步都有对应的 PyTorch API 帮你搞定。

2.5 启动 DDP 的两种方式

DDP 的启动方式有两种，我分别说一下：

方式一：torchrun（推荐）

# 单机 4 卡
torchrun --nproc_per_node=4 train.py

# 多机 8 卡（每机 4 卡）
torchrun --nnodes=2 --nproc_per_node=4 --master_addr=192.168.1.1 train.py

方式二：手动 spawn

mp.spawn(train, args=(world_size,), nprocs=world_size)

我个人习惯用 torchrun。它自动帮你处理了环境变量、进程管理、故障恢复等琐事。你只需要在代码里读一下环境变量就行：

local_rank = int(os.environ['LOCAL_RANK'])
rank = int(os.environ['RANK'])
world_size = int(os.environ['WORLD_SIZE'])

小技巧：用 torchrun 启动时，代码里不需要写 init_process_group 的 init_method 参数。torchrun 会自动设置好 MASTER_ADDR 和 MASTER_PORT 环境变量。

2.6 避坑指南

最后，分享几个我实际踩过的坑：

BatchNorm 的问题：DDP 默认会对 BatchNorm 做同步，但同步 BN 会引入额外的通信开销。如果模型很大，可以考虑用 SyncBatchNorm 或者干脆用 LayerNorm。
随机种子：每张卡的随机种子要不一样，否则数据打乱结果相同。用 torch.utils.data.distributed.DistributedSampler 会自动处理。
梯度累积：如果用了梯度累积，记得在累积完后再调用 optimizer.step()。DDP 的梯度同步发生在 loss.backward() 时，不是 step() 时。
模型保存：只让 rank 0 保存模型，否则会写冲突。保存时用 model.module.state_dict() 而不是 model.state_dict()。

数据并行是分布式训练的起点，也是最容易上手的方式。搞懂了 DDP，后面再学模型并行、流水线并行就会轻松很多。