PyTorch分布式训练的启动脚本优化方案
在多机多卡训练场景中，合理的启动脚本配置对训练性能至关重要。本文将分享几个关键优化点。
1. 使用torchrun替代python -m torch.distributed.launch
# 优化前（旧方式）
python -m torch.distributed.launch --nproc_per_node=8 --nnodes=2 --node_rank=$RANK train.py
# 优化后（推荐方式）
torchrun --nproc_per_node=8 --nnodes=2 --node_rank=$RANK train.py
2. 合理设置环境变量
export NCCL_BLOCKING_WAIT=1
export NCCL_IB_DISABLE=0
export NCCL_NET_GDR_LEVEL=3
export OMP_NUM_THREADS=1
3. 完整优化脚本示例
#!/bin/bash
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
export NCCL_BLOCKING_WAIT=1
export NCCL_IB_DISABLE=0
export NCCL_NET_GDR_LEVEL=3
export OMP_NUM_THREADS=1
torchrun \
--nproc_per_node=8 \
--nnodes=2 \
--node_rank=$RANK \
--master_addr="192.168.1.100" \
--master_port=12345 \
train.py
4. 性能调优要点
- 根据GPU内存调整batch size
- 启用混合精度训练
- 避免在启动脚本中设置过多线程
这些优化能显著提升分布式训练的稳定性和效率。
讨论