分布式训练中数据同步机制踩坑记录
最近在搞分布式训练时遇到了一个让人头大的数据同步问题，特来分享一下踩坑经历。
问题背景
使用PyTorch DDP进行多卡训练时，发现模型在不同GPU上的参数更新不一致，导致loss震荡严重。经过排查，定位到是数据同步环节出了问题。
核心问题
最初使用的torch.nn.parallel.DistributedDataParallel并没有正确配置broadcast_buffers参数，在某些情况下会导致各节点的batch norm统计信息不同步。
复现步骤
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# 错误写法
model = nn.Linear(10, 1)
model = DDP(model, device_ids=[0])
# 正确写法
model = nn.Linear(10, 1)
model = DDP(model, device_ids=[0], broadcast_buffers=True) # 关键参数
解决方案
- 确保在DDP初始化时设置
broadcast_buffers=True - 在训练循环中添加
dist.barrier()确保所有节点同步 - 检查各GPU的batch size是否一致
优化建议
- 可以考虑使用
torch.nn.SyncBatchNorm替代普通BN - 对于大模型训练，可以设置
find_unused_parameters=True但要谨慎使用
这个坑踩得有点惨，希望后来者能少走弯路。
讨论