构建运行 NCCL-Tests 和 SSH 的 PyTorch Docker 镜像

前言

NCCL-Tests 工具是 NVIDIA 开源的一项用于测试 NCCL 集合通信的工具。可以用于检测集合通信是否正常、压测集合通信速率。官方开源地址：https://github.com/NVIDIA/nccl-tests。

在工作中想要测试 NCCL ，并且尽量不影响宿主机，方便的切换 PyTorch 版本，可以使用 Docker 来测试。

一般情况下，操作步骤如下：

拉取 PyTorch 镜像，运行容器。在构建好的 NCCL-Tests 可执行文件拷贝到容器。
为容器安装 SSHD ，并修改配置文件为后续互信做准备。
在容器运行测试。

因为 PyTorch 容器默认的 APT 源是国外的，所以要替换国内源，再安装一些 VIM 工具方便更改配置。这些步骤大部分都和执行 NCCL-Tests 无关，属于准备工作，所以我写了一个 Dockerfile ，去完成上述操作。

项目地址

build-nccl-tests-with-pytorch

快速开始

你可以直接拉取我上传到 dockerhub 的镜像。

docker pull mayooot/nccl-tests-with-pytorch:v0.0.2

或者拉取源码，手动构建镜像。

git clone https://github.com/mayooot/build-nccl-tests-with-pytorch
cd build-nccl-tests-with-pytorch

docker build -t nccl-tests-with-pytorch:latest .

启动容器

需要注意的是，PORT 和 PASS 都是可以替换的，如果你不指定这两个环境变量，它默认都是 12345 ，并将主机的公钥挂载到容器中。然后需要将你的公钥挂载进容器，方便配置互信。这里最好将 ~/.ssh 下面的 id_rsa 和 id_rsa.pub 复制到其他文件夹，然后挂载。避免出现修改了容器的公钥文件，影响宿主机。

docker run --name foo \
  -d -it \
  --network=host \
  -e PORT=1998 -e PASS=P@88w0rd \
  -v /tmp/id_rsa:/root/.ssh/id_rsa \
  -v /tmp/id_rsa.pub:/root/.ssh/id_rsa.pub \
  --gpus all --shm-size=1g \
  --cap-add=IPC_LOCK --device=/dev/infiniband \
  mayooot/nccl-tests-with-pytorch:v0.0.2

下面以all_reduce_perf 为例，演示一下如何使用容器。

首先选出一台机器用来执行测试，然后进入容器。

docker exec -it foo bash

将自己的公钥上传到集群中其他主机上。

ssh-copy-id -p 1998 root@all_cluster_ip

NCCL-Tests 已经被编译好，并放在了 /nccl-tests/build 目录下，我们只需要进入到该目录，然后执行要测试的文件即可。

需要注意的是 --host cluster_ip1,cluster_ip2,... 需要替换成集群的具体 IP ，以,分割。

cd /nccl-tests

mpirun --allow-run-as-root \
  -mca plm_rsh_args "-p 1998" \
  -x NCCL_DEBUG=INFO \
  -x NCCL_IB_HCA=mlx5_10,mlx5_11,mlx5_12,mlx5_13,mlx5_14,mlx5_15,mlx5_16,mlx5_17 \
  --host cluster_ip1,cluster_ip2,... \
  ./build/all_reduce_perf \
  -b 1G -e 4G -f 2 -g 8

最后

如果对你有帮助的话，欢迎到 GitHub 点一个 ⭐。