构建Tensorflow RDMA的Docker镜像

简介: RDMA是一个远程通讯技术,它通过Kernel bypass等方式降低数据传输中的延迟和CPU消耗。 在分布式训练中,由于多个Worker之间或者Worker和Paramater Server 之间需要大量传输模型变量。当GPU到达一定数量后,受制于网络带宽以及TCP协议的延迟,通讯往往会成为计算性能的瓶颈,而在分布式训练中使用RDMA技术能够非常明显地提高训练速度。 #### Tenso

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

在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

可以尝试安装, 检查下编译结果:

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())"

下载正式构建的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

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 

构建正式镜像

docker build . -f nvidia.Dockerfile -t tensorflow-rdma

如果希望在镜像里加入一些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

这样你就可以体验RDMA给分布式机器学习带来的极致加速体验。

相关实践学习
在云上部署ChatGLM2-6B大模型(GPU版)
ChatGLM2-6B是由智谱AI及清华KEG实验室于2023年6月发布的中英双语对话开源大模型。通过本实验,可以学习如何配置AIGC开发环境,如何部署ChatGLM2-6B大模型。
目录
相关文章
|
25天前
|
人工智能 前端开发 Docker
从本地到云端:用 Docker Compose 与 Offload 构建可扩展 AI 智能体
在 AI 智能体开发中,开发者常面临本地调试与云端部署的矛盾。本文介绍如何通过 Docker Compose 与 Docker Offload 解决这一难题,实现从本地快速迭代到云端高效扩容的全流程。内容涵盖多服务协同、容器化配置、GPU 支持及实战案例,助你构建高效、一致的 AI 智能体开发环境。
211 1
从本地到云端:用 Docker Compose 与 Offload 构建可扩展 AI 智能体
|
1月前
|
JavaScript Docker 容器
使用Docker多阶段构建优化镜像大小
使用Docker多阶段构建优化镜像大小
260 100
|
1月前
|
缓存 安全 Linux
优化Docker镜像大小的多阶段构建实践
优化Docker镜像大小的多阶段构建实践
206 99
|
1月前
|
缓存 前端开发 Docker
Docker Layer Caching:加速你的容器构建
Docker Layer Caching:加速你的容器构建
|
1月前
|
安全 Go Docker
使用Docker多阶段构建优化镜像大小
使用Docker多阶段构建优化镜像大小
|
3天前
|
NoSQL 算法 Redis
【Docker】(3)学习Docker中 镜像与容器数据卷、映射关系!手把手带你安装 MySql主从同步 和 Redis三主三从集群!并且进行主从切换与扩容操作,还有分析 哈希分区 等知识点!
Union文件系统(UnionFS)是一种**分层、轻量级并且高性能的文件系统**,它支持对文件系统的修改作为一次提交来一层层的叠加,同时可以将不同目录挂载到同一个虚拟文件系统下(unite several directories into a single virtual filesystem) Union 文件系统是 Docker 镜像的基础。 镜像可以通过分层来进行继承,基于基础镜像(没有父镜像),可以制作各种具体的应用镜像。
58 0
|
1月前
|
Java Docker 容器
使用Docker多阶段构建优化镜像大小
使用Docker多阶段构建优化镜像大小
70 8
|
Linux Docker Windows
Docker概述:什么Docker如此特别,它是如何构建的
本文讲的是Docker概述:什么Docker如此特别,它是如何构建的,【编者的话】在上一系列介绍中,我们大体地介绍了CoreOS。在这个系列中我想介绍一下Docker。本文我们看一下为什么Docker如此特别,它是如何构建的。我们会从基本开始介绍,所以读者不必担心自己是刚接触Docker的。
1435 0
|
3月前
|
存储 监控 测试技术
如何将现有的应用程序迁移到Docker容器中?
如何将现有的应用程序迁移到Docker容器中?
290 57
下一篇
日志分析软件