PyTorch DDP训练启动参数调优技巧
在多机多卡训练场景中，PyTorch Distributed Data Parallel (DDP)的启动参数配置直接影响训练性能。本文将分享几个关键调优技巧。
1. 启动脚本优化
使用torchrun替代python -m torch.distributed.launch：
# 推荐方式
python -m torch.distributed.run \
--nproc_per_node=8 \
--nnodes=2 \
--node_rank=0 \
--master_addr="192.168.1.100" \
--master_port=12345 \
main.py
2. 关键参数调优
- nproc_per_node：设置GPU数量，通常等于每台机器的GPU数
- master_addr/master_port：主节点IP和端口，确保网络连通性
- nnodes/node_rank：多机训练时指定节点信息
3. 性能调优建议
使用以下环境变量提升性能：
export TORCH_CUDNN_BENCHMARKS=1
export NCCL_BLOCKING_WAIT=1
export NCCL_MAX_NLINKS=128
4. 验证配置
启动后可通过以下方式验证DDP是否正确初始化：
import torch.distributed as dist
print(f"World size: {dist.get_world_size()}")
print(f"Rank: {dist.get_rank()}")
合理配置这些参数能显著提升分布式训练效率。
讨论