分布式训练中批处理大小选择

分布式训练中批处理大小选择踩坑指南

在分布式训练中，批处理大小（batch size）的选择直接影响训练效率和模型性能。本文记录了在Horovod和PyTorch Distributed环境下的实际踩坑经验。

核心问题

在多机多卡训练中，过小的batch size会导致：

• 梯度估计不准确
• 训练不稳定
• GPU利用率低

而过大的batch size则可能：

• 内存溢出（OOM）
• 降低模型泛化能力
• 增加通信开销

PyTorch Distributed实战配置

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

# 初始化分布式环境
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
rank = int(os.environ['RANK'])
world_size = int(os.environ['WORLD_SIZE'])

# 设置设备
device = torch.device(f'cuda:{rank}')

# 根据GPU数量调整batch size
# 建议：每个GPU 8-32个样本
per_gpu_batch_size = 16
actual_batch_size = per_gpu_batch_size * world_size

# 创建数据加载器
train_loader = DataLoader(
    dataset, 
    batch_size=per_gpu_batch_size, 
    shuffle=True,
    num_workers=4,
    pin_memory=True
)

Horovod配置示例

import horovod.torch as hvd
import torch.nn.functional as F

# 初始化Horovod
hvd.init()

# 设置GPU
torch.cuda.set_device(hvd.local_rank())

# 根据worker数量调整batch size
global_batch_size = 64  # 总batch size
per_gpu_batch_size = global_batch_size // hvd.size()

# 数据加载器配置
train_loader = DataLoader(
    dataset,
    batch_size=per_gpu_batch_size,
    shuffle=True,
    num_workers=2
)

调试建议

1. 从较小的batch size开始（如8-16）
2. 逐步增加直到出现OOM
3. 监控GPU内存使用率
4. 记录不同batch size下的训练速度和loss曲线

实际经验：在8卡环境下，推荐每个GPU 16-32个样本作为起始点。

注意：不同模型架构对batch size敏感度不同，建议根据具体模型调整参数。