PyTorch DDP多卡训练:从跑通到跑快的三步

第一次把微调任务从单卡迁到4卡DDP,踩了一串坑。从"能跑"到"线性加速"中间隔着三步,记录完整过程。

第一步,跑通。DDP的基本结构:torchrun启动多进程,每个进程持有一份模型副本,DistributedSampler保证各卡数据不重叠,反向传播时梯度自动all-reduce同步。最常见的报错是没设rank判断就所有进程都写checkpoint和打日志,磁盘和日志直接打架,所有IO操作包一层if rank等于0。

第二步,消灭数据倾斜。DistributedSampler自动做数据分片,但它默认pad_to_multiple_of会把数据集补齐到卡数的整数倍,小数据集时重复样本被训练两次,验证指标莫名虚高。shuffle也要靠sampler自带,自己再套DataLoader的shuffle会破坏分片语义。

第三步,性能压榨。起初四卡只有2.3倍加速,排查出三个损耗点。一是梯度同步和计算没有重叠,DDP默认在backward时同步,把小batch的通信占比放大,开梯度累积到等效大batch后通信摊薄;二是每步都调用dist.barrier做同步,纯属多余,删掉;三是各卡负载不均,长短句混排导致某张卡永远最慢,先按长度排序再分桶,配合bucket sampler均衡每步耗时。

sampler = DistributedSampler(dataset, shuffle=True)
loader = DataLoader(dataset, batch_size=32,
                    sampler=sampler, num_workers=8,
                    pin_memory=True)
for epoch in range(epochs):
    sampler.set_epoch(epoch)     # 保证每轮洗牌不同
    for batch in loader:
        loss = model(batch).loss / grad_accum
        loss.backward()
        if step % grad_accum == 0:
            optimizer.step(); optimizer.zero_grad()

三步之后,四卡加速比3.7,基本贴着线性上限。最后提醒一句:验证阶段要用all_gather把各卡的结果聚合再算指标,各自算完取平均在加权指标上是错的,这个隐蔽bug我们上线前一周才发现。

DDP的验证清单再补几条,都是血泪换的。一是确定性检查,固定seed后两次训练的loss曲线应该逐step一致,不一致多半有未被DDP同步的随机性,比如某处用了未加种子的numpy随机,这种问题不影响收敛但影响实验可比性。二是梯度范数监控,DDP下各卡的梯度是all-reduce后的均值,范数应该一致,如果打印各卡本地范数差异巨大,说明数据分片不均匀。三是save和load的map_location,checkpoint在rank0保存后,加载时所有卡都要加载到各自设备,显式指定而不是让框架猜。最后,遇到诡异的多进程挂死,先用py-spy dump看各进程栈,九成是某个rank提前抛异常退出,其他rank在通信点傻等,日志里的第一张报错卡才是真凶。