[论文解读] User-transparent Distributed TensorFlow

Abhinav Vishnu, Joseph Manzano+2 more|arXiv (Cornell University)|Apr 15, 2017

Scientific Computing and Data Management参考文献 6被引用 3

一句话总结

本论文提出 MaTEx-TensorFlow，这是一种用户透明的分布式 TensorFlow 实现，可在无需修改用户代码的情况下，在 HPC 系统上实现横向扩展训练。通过将 MPI 直接集成到 TensorFlow 运行时并抽象底层通信，它在使用标准 AlexNet、GoogLeNet、InceptionV3 和 ResNet-50 模型的多节点 CPU 和 GPU 系统上实现了接近理想的扩展性能，且损失曲线与串行训练完全一致。

ABSTRACT

Deep Learning (DL) algorithms have become the {\em de facto} choice for data analysis. Several DL implementations -- primarily limited to a single compute node -- such as Caffe, TensorFlow, Theano and Torch have become readily available. Distributed DL implementations capable of execution on large scale systems are becoming important to address the computational needs of large data produced by scientific simulations and experiments. Yet, the adoption of distributed DL implementations faces significant impediments: 1) most implementations require DL analysts to modify their code significantly -- which is a show-stopper, 2) several distributed DL implementations are geared towards cloud computing systems -- which is inadequate for execution on massively parallel systems such as supercomputers. This work addresses each of these problems. We provide a distributed memory DL implementation by incorporating required changes in the TensorFlow runtime itself. This dramatically reduces the entry barrier for using a distributed TensorFlow implementation. We use Message Passing Interface (MPI) -- which provides performance portability, especially since MPI specific changes are abstracted from users. Lastly -- and arguably most importantly -- we make our implementation available for broader use, under the umbrella of Machine Learning Toolkit for Extreme Scale (MaTEx) at { exttt http://hpc.pnl.gov/matex}. We refer to our implementation as MaTEx-TensorFlow.

研究动机与目标

• 通过消除用户修改其 TensorFlow 代码的需求，降低使用分布式深度学习的入门门槛。
• 在具有高带宽互连（如 InfiniBand）的 HPC 系统上高效执行分布式 TensorFlow，而这类系统对基于 gRPC 的解决方案不适用。
• 提供一个生产就绪、高性能且可移植的分布式 TensorFlow 运行时，支持跨多个节点的数据并行。
• 通过保持与标准 TensorFlow 相同的训练动态和收敛行为，确保功能对等。
• 通过机器学习极大规模工具包（MaTEx）广泛提供该实现，以加速基于 HPC 的深度学习应用。

提出的方法

• 将 MPI 直接集成到 TensorFlow 运行时，以替代基于 gRPC 的通信，从而原生支持 InfiniBand 等高性能互连。
• 通过将 MPI 逻辑封装在运行时内部，从用户处抽象出 MPI 特定实现，使用户无需编写 MPI 调用或修改其现有 TensorFlow 脚本。
• 通过使用 MPI 集体操作同步副本之间的梯度，实现基于同步随机梯度下降（SGD）的数据并行。
• 扩展 TensorFlow 1.0，新增用户操作和通信原语，以支持分布式训练，同时保持向后兼容性。
• 实现一个支持多种数据格式的自定义数据读取接口，简化分布式环境中的数据加载。
• 复用现有的 TensorFlow 计算图和自动微分机制，确保与所有标准模型（如 AlexNet、ResNet-50）的兼容性。

实验结果

研究问题

• RQ1 能否构建一个无需修改用户脚本的分布式 TensorFlow 实现？
• RQ2 能否有效将 MPI 集成到 TensorFlow 运行时中，以在具有 InfiniBand 互连的 HPC 系统上实现高性能、可扩展的训练？
• RQ3 所提出的实现是否与串行 TensorFlow 训练相比保持了相同的收敛行为和损失曲线？
• RQ4 该系统能否在多核 CPU 和多 GPU HPC 平台上实现强扩展性，同时不牺牲性能可移植性？
• RQ5 在 MaTEx 等开源 HPC 软件生态体系下，是否可行提供一个生产就绪、用户透明的分布式深度学习系统？

主要发现

• MaTEx-TensorFlow 与串行 TensorFlow 实现了功能对等，这在使用 AlexNet 在 ImageNet LSVRC12 数据集上训练时，损失曲线完全一致得到验证。
• 该实现可在使用 InfiniBand 互连的 Intel 多核系统和 NVIDIA 多 GPU 系统上，有效扩展至多个计算节点。
• 该系统实现了无需用户修改 TensorFlow 代码的透明分布式训练，显著降低了基于 HPC 的深度学习应用门槛。
• 通过将 MPI 嵌入 TensorFlow 运行时，该方法提供了性能可移植性，并高效利用了超级计算环境中常见的高带宽互连。
• 该解决方案已具备生产就绪状态，并通过 MaTEx 项目在 http://hpc.pnl.gov/matex 公开提供，支持包括 GoogLeNet、InceptionV3 和 ResNet-50 在内的广泛深度学习模型。
• 该方法成功克服了基于 gRPC 的分布式 TensorFlow 的局限性，后者因依赖以太网且缺乏原生 RDMA 支持，而不适用于 HPC 系统。