一、数据并行：分布式训练的基石

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

我刚开始接触分布式训练时，觉得这概念挺简单的——不就是多卡跑同一个模型嘛。但真正上手才发现，里面的门道比想象中多得多。今天咱们就把数据并行聊透。

1.1 核心思想：分而治之

假设你有一个batch size为64的训练任务，单卡要算64个样本。换成4卡数据并行，每卡只算16个样本。每张卡都持有完整的模型副本，但处理的数据不同。

流程是这样的：

1. 每张卡拿到自己的数据切片
2. 各自前向传播，计算loss
3. 各自反向传播，得到梯度
4. 所有卡的梯度做同步（AllReduce）
5. 每卡用同步后的梯度更新自己的模型

嗯，这里要注意：模型参数是冗余存储的。每张卡都有一份完整的参数副本。这既是优点也是缺点——优点是实现简单，缺点是显存开销大。

1.2 梯度同步机制：AllReduce

梯度同步是数据并行的核心。我见过不少新手在这里翻车。

AllReduce 不是单一算法，而是一族算法的统称。最常用的有：

算法	通信量	特点
Ring AllReduce	2(N-1) * 梯度大小 / N	带宽最优，NVIDIA默认方案
Tree AllReduce	2 * 梯度大小 * logN	延迟敏感场景更优
Recursive Halving-Doubling	2 * 梯度大小 * (1 - 1/N)	兼顾带宽和延迟

Ring AllReduce 是目前的主流。它分两步走：

1. Scatter-Reduce：每张卡把梯度切成N份，环形传递并累加
2. AllGather：把累加后的结果广播给所有卡

我曾经在8卡机器上调试一个模型，发现通信开销占了总时间的40%。后来发现是Ring的块大小没调好——默认的4MB在某些场景下不是最优的。调成1MB后，通信时间降到了25%。

# PyTorch 中的 AllReduce 示例
import torch
import torch.distributed as dist

# 初始化进程组
dist.init_process_group(backend='nccl')

# 每张卡计算自己的梯度
loss = model(data)
loss.backward()

# 同步梯度：所有卡的梯度求和后平均
for param in model.parameters():
    dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
    param.grad /= dist.get_world_size()

# 更新参数
optimizer.step()

你想想看，如果不用AllReduce，每张卡各自为政，那模型就乱套了。同步机制保证了所有卡看到的是同一个「全局梯度」。

1.3 梯度累积：小batch也能跑大模型

梯度累积是个很实用的技巧。说白了就是：攒几次梯度再更新。

为什么要这么做？因为显存有限。比如你想用batch size=1024，但单卡只能塞下batch size=128。那就每卡算4次，每次梯度都累加起来，第4次结束后才更新参数。

我习惯这样实现：

accumulation_steps = 4  # 累积4次
optimizer.zero_grad()

for i, (data, target) in enumerate(dataloader):
    output = model(data)
    loss = criterion(output, target)
    loss = loss / accumulation_steps  # 重要！要除以累积步数
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

梯度累积和AllReduce可以配合使用。每张卡在自己的数据上累积梯度，累积到指定步数后再做一次AllReduce。这样通信次数从「每步一次」降到了「每N步一次」，通信开销大幅降低。

1.4 学习率调整：Linear Scaling Rule

这个规则很简单：batch size翻倍，学习率也翻倍。

为什么？因为梯度是batch内样本梯度的平均。batch size越大，梯度方差越小，方向越稳定。所以可以用更大的学习率来加速收敛。

公式长这样：

lr_new = lr_base * (batch_size_new / batch_size_base)

举个例子：单卡batch size=32，lr=0.001。换成8卡后，全局batch size=256，lr应该设为0.008。

但这里有个坑——不是所有层都适用。我遇到过Batch Normalization层的问题。BN的统计量是逐卡计算的，batch size翻倍后，BN的统计量反而更不准了。这时候需要调整BN的momentum参数。

另外，学习率调整不是线性的那么简单。有些场景下，平方根缩放效果更好：

lr_new = lr_base * sqrt(batch_size_new / batch_size_base)

我个人习惯先试线性缩放，如果loss震荡严重，再换平方根缩放。没有银弹，得根据具体任务调。

知识体系总览

下面这张图把数据并行的核心逻辑串起来了：

这张图把四个核心知识点串在了一起。你可以看到，数据并行不是孤立的技术，它需要同步机制、累积策略和学习率调整三者配合。

想要更多资料，联系可可