前提:环境说明
显卡:Nvidia 3090Ti
CUDA 11.1
Cudnn 8.0
OS:Ubuntu 20.x
Tensorflow:nvidia-tensorflow(后面有安装过程)
tensorflow1.x版本不提供对A100/3090等新版显卡的支持,因此使用nvidia-tensorflow(如果不安装此版本tf,就算检测到显卡也不会调用)。另外,nvidia-tensorflow 只支持ubuntu20.04。
horovod介绍
horovod是Uber 团队开发的分布式训练框架,他可以满足让你尽量少的修改代码即可将在单卡训练的脚本横行扩展为多卡并行训练,同时又兼顾训练的加速。目前支持tensorflow/keras/pytorch/mxnet.底层通信主要依赖NCCL/Gloo(测试后NCCL是最快的),支持MPI(CPU 训练更快)。由于其训练加速效果比tensorflow 原生的distributedStrategy 快很多,所以在分布式训练时,推荐使用。
下面主要针对tensorflow1.x 下做分布式训练进行说明。
环境主要依赖tensorflow1.x/horovod/nccl/mpi ,有两种方式搭环境:local/docker,本文采用local本地安装。
1. NCCL安装
https://github.com/NVIDIA/nccl
安装NCCL,注意版本,为了兼顾A100,推荐使用v2.8.3-1 这个版本。
# 编译nccl
git clone https://github.com/NVIDIA/nccl.git
cd nccl && git checkout v2.8.3-1
make -j src.build
# 如果第一次安装,需要安装一下依赖
# Install tools to create debian packages
sudo apt install build-essential devscripts debhelper fakeroot
# Build NCCL deb package
make pkg.debian.build
ls build/pkg/deb/
# install
sudo make install
安装这里有一步很耗时间,等着就行了
2. mpirun安装
https://www.open-mpi.org/
# 安装依赖
apt-get install libnuma-dev
# 下载安装包
wget https://www.open-mpi.org/software/ompi/v4.0/downloads/openmpi-4.0.0.tar.gz
# 解压
tar zxf openmpi-4.0.0.tar.gz
# 安装
cd openmpi-4.0.0
./configure --enable-orterun-prefix-by-default
make -j $(nproc) all
sudo make install #需要root权限
sudo ldconfig #需要root权限
查看版本
mpirun --version
3. nvidia-tensorflow安装
pip或者conda安装
pip install nvidia-pyindex
pip install nvidia-tensorflow
这里如果下载失败,多换几个网络吧。。
3090欢乐时刻——支持gpu
输入命令检测是否可以使用gpu
import tensorflow as tf
print(tf.test.is_gpu_available())
4. horovod安装
horovod 在安装时,需要安装支持NCCL ,同时建议安装最新版本。
如果安装报错,请安装指定版本,如下命令。
卸载之前的horovod版本
pip show horovod && pip uninstall horovod
安装命令
HOROVOD_WITH_MPI=1 HOROVOD_WITHOUT_GLOO=1 HOROVOD_GPU_OPERATIONS=NCCL HOROVOD_WITH_TENSORFLOW=1 HOROVOD_NCCL_LINK=SHARED pip3 install horovod==0.24.3 --no-cache-dir
验证一下horovod
horovodrun --check-build
前面勾选则代表支持
注:
# 确认horovod链接的nccl版本路径正确
ldd /usr/local/lib/python3.8/dist-packages/horovod/tensorflow/mpi_lib.cpython-38-x86_64-linux-gnu.so
参考资料:
https://www.jianshu.com/p/975f5cca88e4
https://xv44586.github.io/2022/05/25/horovod/index.html#huan-jing-zhong-cai-guo-de-keng