多GPU内存分配策略调优踩坑记录
最近在优化PyTorch分布式训练时，遇到了一个经典的多GPU内存分配问题。项目使用Horovod进行多机多卡训练，但在训练过程中频繁出现CUDA out of memory错误。
问题现象
使用4卡GPU训练时，单卡内存占用达到15GB，远超预期的8GB。查看日志发现，即使设置了torch.cuda.set_per_process_memory_fraction(0.5)，内存分配依然不合理。
踩坑过程
第一步：基础配置失败
import torch
import horovod.torch as hvd
torch.cuda.set_per_process_memory_fraction(0.4)
hvd.init()
结果：内存分配依然不均，部分GPU占用过高。
第二步：强制内存限制
import os
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
# 限制单卡内存增长
for i in range(4):
torch.cuda.set_per_process_memory_fraction(0.3, device=i)
结果：训练启动失败，报错RuntimeError: CUDA out of memory。
第三步：正确配置策略 最终通过以下方式解决：
import torch
torch.cuda.set_per_process_memory_fraction(0.25) # 每卡分配25%内存
# 设置环境变量
os.environ['HOROVOD_FUSION_THRESHOLD'] = '0'
os.environ['NCCL_IB_DISABLE'] = '0' # 启用IB通信
os.environ['NCCL_SOCKET_IFNAME'] = 'eth0' # 指定网络接口
优化建议
- 使用
torch.cuda.set_per_process_memory_fraction()控制单卡内存分配 - 设置合理的
HOROVOD_FUSION_THRESHOLD避免过度融合 - 根据实际GPU显存大小调整内存分配比例
- 配置正确的网络接口以提升通信效率
配置完成后，训练稳定性显著提升，内存利用率达到90%以上。
讨论