一个用于分布式DNN训练加速的通用通信调度器

论文：https://i.cs.hku.hk/~cwu/papers/yhpeng-sosp19.pdf

译文：一个用于分布式DNN训练加速的通用通信调度器

摘要

我们提出ByteScheduler，一个用于分布式DNN训练加速的通用通信调度器。ByteScheduler基于我们最基本的分析：通过张量切分和重组来进行传输，可以得到理论上的最优结果，即使真实环境中有调度代价，性能表现也很不错。为了让ByteScheduler适用于各种DNN训练框架，我们介绍一种统一的抽象和一种依赖代理机制，使得可以在不破坏框架引擎原有依赖的基础上进行通信调度（注：这是什么东西？）。我们进一步介绍一种贝叶斯优化方法，可以在各种网络环境下，针对不同训练模型，来自动调整张量划分的尺寸，以及其他参数。ByteScheduler目前支持TensorFlow、PyTorch和MXNet，并且不用修改源码，兼容Parameter Server和all-reduce架构的梯度同步，TCP和RDMA网络都行。我们的实验证明，ByteScheduler训练的加速比是2.96。

介绍

深度神经网络（DNN）被广泛运用于各个领域，例如计算机视觉和自然语言处理等等。但是，训练DNN很耗时，主要是因为训练数据多，以及DNN模型的尺寸越来越大。加速DNN训练最常用的方法是数据并行。不幸的是，因为通信开销，它的性能常常比线性加速比差很远。作为一个大型在线服务提供商，在很多我们内部和公开训练任务中，通信通常占总训练时间的很大一部分。因此，很多通信加速方法被提出并且集成到

背景和动机

DNN训练和数据并行：

在深度学习中，DNN模型是通过迭代一个大数据集来训练的（也叫epoch），以最小化损失函数。 

正向和反向传播：

在每个epoch期间，数据集被划分成mini-batch。在每次iteration，一个mini-batch逐层经过DNN并产生一个损失，这个过程叫正向传播（FP）。

张量划分：

 

实现细节 

ByteScheduler是用C++11和Python实现的。核心是一个通用Python库，用于所有ML框架。对于每个框架，都需要一个插件，用于PS或者all-reduce架构。插件通常用C++和Python混编。

ByteScheduler Core：

MXNet plugin：

PyTorch plugin：

TensorFlow plugin：

 

评估

 

讨论和未来方向

 

相关工作 

加速DNN训练的通信：

加速DNN训练的通信：

加速DNN训练的通信： 

结论

ByteScheduler是一个用于分布式DNN训练加速的通用通信调度器。我们的实现，支持很多ML框架，包括MXNet、PyTorch和TensorFlow，兼容PS和all-reduce梯度同步，TCP和RDMA网络都行。ByteScheduler端到端可以加速196%。主要的设计思想包括：通信操作的统一抽象、依赖代理和系统参数自动调整。我们已经开源我们的实现，期待社区可以加入更多对已有和未来框架的支持。