PyTorch梯度如何分布式累加

在 PyTorch 中，“梯度分布式累加”通常出现在以下几种场景，我按常见程度分别说明：

一、同进程内多个 batch / 多路 loss 累加（最常见）

PyTorch 默认就是梯度累加的：

optimizer.zero_grad()
for i in range(accumulation_steps):
outputs = model(inputs[i])
loss = criterion(outputs, targets[i])
loss = loss / accumulation_steps
loss.backward()   # 梯度自动累加
optimizer.step()
optimizer.zero_grad()

✅ 关键点：

• loss.backward() 不会清除梯度
• 多次 backward() 的梯度会自动相加
• 除以 accumulation_steps 是为了等效大 batch

二、DataParallel（单机多卡）中的梯度累加

nn.DataParallel 自动做梯度求和：

model = nn.DataParallel(model)
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()  # 各 GPU 梯度自动 sum 到主卡
optimizer.step()

✅ 内部机制：

• 每张卡计算局部梯度
• 主卡 all_reduce(sum) 梯度
• 再更新参数

⚠️ 注意：

• 已经是“分布式累加”
• 不需要手动处理

三、DistributedDataParallel（DDP，多机多卡）

在 DistributedDataParallel 中：

model = DistributedDataParallel(model, device_ids=[local_rank])

梯度是如何累加的？

• 每个进程独立前向 + 反向
• loss.backward() 时：
• 梯度自动 all_reduce(SUM)
• 所有 rank 得到相同梯度

✅ 等价于：

global_grad = grad_rank0 + grad_rank1 + ... + grad_rankN

⚠️ 重要细节：

• DDP 默认是 SUM，不是 MEAN
• 如果你想等效大 batch：

loss = loss / world_size

或

for p in model.parameters():
p.grad /= world_size

四、手动跨进程梯度累加（不常见）

如果你不用 DDP，想自己控制：

# rank 0 ~ N
grad = p.grad.clone()
dist.all_reduce(grad, op=dist.ReduceOp.SUM)
p.grad = grad

或：

dist.all_reduce(p.grad, op=dist.ReduceOp.SUM)

✅ 常用于：

• 自定义并行策略
• Pipeline Parallelism
• 混合并行

五、梯度累加 + DDP 的正确写法

optimizer.zero_grad()
for i in range(accum_steps):
out = model(x[i])
loss = crit(out, y[i]) / (accum_steps * world_size)
loss.backward()
optimizer.step()

六、总结对照表

如果你有具体场景（比如：DDP + 梯度累积 / 自定义并行 / 混合精度），可以直接说，我可以给针对性代码。