RDMA是一个远程通讯技术，它通过Kernel bypass等方式降低数据传输中的延迟和CPU消耗。
在分布式训练中，由于多个Worker之间或者Worker和Paramater Server 之间需要大量传输模型变量。当GPU到达一定数量后，受制于网络带宽以及TCP协议的延迟，通讯往往会成为计算性能的瓶颈，而在分布式训练中使用RDMA技术能够非常明显地提高训练速度。
Tensorflow
Tensorflow是谷歌开源的深度学习框架，它有丰富的平台支持和API，也可以非常轻松地构建分布式模型训练。
Tensorflow 在实现里支持RDMA作为其分布式场景的通讯协议，但是官方镜像默认没有支持RDMA。需要重新构建tensorflow，并开启RDMA相关的构建参数。 Tensorflow 对 RDMA的支持和实现协议参见文档： https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/verbs/README.md
如何build from source 可以见官网
https://www.tensorflow.org/install/source
开始构建
找一台GPU节点，运行develop容器，作为编译环境
docker run -it -w /tensorflow -v $PWD:/mnt -e HOST_PERMS="$(id -u):$(id -g)" tensorflow/tensorflow:nightly-devel-gpu-py3 bash
AI 代码解读
在dev容器中执行build
# 开启RDMA参数
export TF_NEED_GDR=1
export TF_NEED_VERBS=1
tensorflow/tools/ci_build/builds/configured GPU
bazel build --config=opt --config=cuda //tensorflow/tools/pip_package:build_pip_package
./bazel-bin/tensorflow/tools/pip_package/build_pip_package /mnt # create package
chown $HOST_PERMS /mnt/tensorflow-*.whl
AI 代码解读
可以尝试安装， 检查下编译结果：
pip uninstall tensorflow # remove current version
pip install /mnt/tensorflow-*.whl
cd /tmp # don't import from source directory
python -c "import tensorflow as tf; print(tf.contrib.eager.num_gpus())"
AI 代码解读
下载正式构建的dockerfile和所需文件：
curl -OL https://raw.githubusercontent.com/tensorflow/tensorflow/r1.12/tensorflow/tools/dockerfiles/dockerfiles/nvidia.Dockerfile
curl -OL https://raw.githubusercontent.com/tensorflow/tensorflow/r1.12/tensorflow/tools/dockerfiles/bashrc
AI 代码解读
对nvidia.Dockerfile 做一下修改，将之前编译的Tensorflow whl文件添加到镜像并执行安装
ARG TF_PACKAGE=tensorflow-gpu
RUN ${PIP} install ${TF_PACKAGE}
改为
ADD . /tensorflow-rdma-gpu
ARG TF_PACKAGE=/tensorflow-rdma-gpu/*.whl
RUN ${PIP} install ${TF_PACKAGE}
RUN rm -rf /tensorflow-rdma-gpu
AI 代码解读
构建正式镜像
docker build . -f nvidia.Dockerfile -t tensorflow-rdma
AI 代码解读
如果希望在镜像里加入一些mellanox的工具，可以在dockerfile最后再扩展一下。
RUN apt-get update && apt-get install -y --no-install-recommends \
libibverbs-dev \
librdmacm-dev \
wget \
&& \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
RUN wget http://content.mellanox.com/ofed/MLNX_OFED-4.1-1.0.2.0/MLNX_OFED_LINUX-4.1-1.0.2.0-ubuntu16.04-x86_64.tgz && \
tar -xzvf MLNX_OFED_LINUX-4.1-1.0.2.0-ubuntu16.04-x86_64.tgz && \
MLNX_OFED_LINUX-4.1-1.0.2.0-ubuntu16.04-x86_64/mlnxofedinstall --user-space-only --without-fw-update --all -q && \
cd .. && \
rm -rf MLNX_OFED_LINUX-4.1-1.0.2.0-ubuntu16.04-x86_64 && \
rm -rf *.tgz
AI 代码解读
这样你就可以体验RDMA给分布式机器学习带来的极致加速体验。
