基于InfiniBand网络的分布式训练性能调优踩坑记录
最近在部署大规模分布式训练时，遇到了一个典型的性能瓶颈问题：虽然模型规模达到了万亿参数级别，但训练速度始终无法突破预期。经过一周的排查和优化，终于找到了问题根源并成功调优。
问题现象
使用4台服务器进行分布式训练，每台配2张GPU，通过InfiniBand网络连接。在训练初期，网络带宽利用率只有30%左右，训练效率远低于预期。
排查过程
首先，通过ibstat命令确认了InfiniBand链路状态正常，然后使用perf_test工具测试网络延迟和带宽：
# 测试网络性能
perf_test -t ud -a 1048576
结果显示网络带宽可达20GB/s，理论上应该能够满足训练需求。但实际训练中发现：
核心问题定位
通过nvidia-smi和ibstat同时监控发现，在模型参数同步阶段，网络利用率突然下降，推测是网络缓冲区设置不当导致。
解决方案
- 调整InfiniBand驱动参数：
# 增加RDMA缓冲区大小
echo 4096 > /sys/class/infiniband/mlx5_0/rdma_device/rdma_write_max_sge
- 优化NCCL配置参数：
import os
os.environ['NCCL_IB_DISABLE'] = '0'
os.environ['NCCL_IB_HCA'] = 'mlx5_0'
os.environ['NCCL_IB_GID_INDEX'] = '3'
os.environ['NCCL_IB_TIMEOUT'] = '22'
os.environ['NCCL_IB_RETRY_CNT'] = '1'
- 调整GPU驱动参数：
# 优化GPU内存分配策略
nvidia-smi -pl 400
最终效果
优化后，训练效率提升了约35%，网络利用率稳定在85%以上。建议在大规模分布式训练前，务必先进行网络性能测试。
注意： 此方案适用于InfiniBand环境，普通以太网环境不适用。
讨论