多卡训练中模型并行效率测试

多卡训练中模型并行效率测试踩坑记录

最近在进行多卡训练性能优化时，发现模型并行效率存在严重问题。本文记录了从配置到踩坑的完整过程。

环境配置

使用PyTorch 1.9 + CUDA 11.1，4张V100显卡进行测试。基础代码如下：

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

# 初始化分布式环境
def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

# 模型定义
model = MyModel().to(rank)
setup(rank, world_size)
model = DDP(model, device_ids=[rank])

遇到的问题

测试发现，当batch size从32增加到128时，训练速度反而下降了30%。通过nvidia-smi监控发现GPU利用率极低，且通信开销巨大。

解决方案

关键在于合理设置以下参数：

# 设置环境变量优化性能
os.environ['NCCL_IB_DISABLE'] = '0'
os.environ['NCCL_NET_GDR_LEVEL'] = '3'
os.environ['NCCL_P2P_DISABLE'] = '0'

# 调整数据加载器
train_loader = DataLoader(dataset, batch_size=32, num_workers=4, pin_memory=True)

最终将batch size设置为64时，性能达到最优。建议在多卡训练中优先测试不同batch size的性能表现。

性能对比

• 1卡：100%速度
• 2卡：170%速度（加速比0.85）
• 4卡：300%速度（加速比0.75）

这说明在当前配置下，存在明显的通信瓶颈。