字节跳动开源高性能分布式训练框架 BytePS：兼容 TensorFlow 等主流框架

近日，字节跳动人工智能实验室宣布开源一款高性能分布式深度学习训练框架 BytePS，在性能上颠覆了过去几年 allreduce 流派一直占据上风的局面，超出目前其他所有分布式训练框架一倍以上的性能，且同时能够支持 Tensorflow、PyTorch、MXNet 等开源库。

BytePS 结合了字节跳动人工智能实验室几个月来对分布式训练通信的多个研究与优化成果，包含通信优先级调度、PS 的 RDMA 实现、针对 PCIe switch 与 NUMA 的优化，以及 BytePS 本身构架的创新等。

深度学习的效果取决于模型与数据，目前行业内不断刷新深度学习准确率的最新研究，大多都基于更大的模型以及更大的数据集。然而，大模型与大数据对训练时的计算能力提出了极高要求，单张 GPU 卡，或者单台服务器上的 GPU 卡，已经远远不能够满足内部训练任务的需求。因此，分布式训练的效率，即使用多台服务器协同进行训练，现在成为了深度学习系统的核心竞争力。

一直以来，在分布式训练中有两大流派，分别是 allreduce 和 PS（Parameter Server）。过去三年中，尤其是百度提出 allreduce，以及 Uber 开源基于 allreduce 的 Horovod 之后，行业内的认知中，allreduce 是最好的分布式训练通信方式，而过去的 PS 实现的性能也确实与 allreduce 存在一定差距。

BytePS 则颠覆了 allreduce 长期领先的局面，BytePS 拥有着超出目前其他所有分布式训练框架一倍以上的性能，包括 NVIDIA 开源的 NCCL，Uber 开源的 Horovod，以及 Tensorflow、PyTorch、MXNet 自带的分布式训练方案等。

BytePS 开发团队表示，在公有云或者私有云这类共享集群中，经过精巧设计和高质量实现的 PS，PS 架构不仅不比 allreduce 差，而且在一些环境还能得到比 allreduce 还高一倍的速度。

测试中，BytePS 团队使用了公有云上的虚拟机，每个虚拟机有 8 张 Tesla V100 16GB GPU，GPU 之间通过 NVLink 进行高速互连。每个 GPU 上的 batch size 选取为 64。虚拟机之间通过 20Gbps 的 TCP/IP 网络进行连接。在这种情况下，由于机器之内带宽足够大，TCP/IP 的网络带宽则成为了主要瓶颈。

BytePS 选择了 Resnet50 和 VGG16 两个模型进行评测，其中 Resnet50 是计算密集型的模型（对通信要求低，优化空间小），VGG16 是通信密集型的模型（对通信要求高，优化空间大），对照组选择了目前市面上最流行的通信框架之一 Horovod-NCCL（基于 allreduce 算法实现），性能指标为每秒钟训练的 ImageNet 图片数量，越高代表越好。

通过两组实验结果可以看出，对于计算密集型的 Resnet50 模型，BytePS 性能超过 Horovod-NCCL 近 44%；而对于通信密集型的 VGG16 模型，BytePS 性能可以超过 Horovod-NCCL 将近 100%。

BytePS 团队也在配有 100Gbps 的 RDMA 网络的私有集群做了测试，BytePS 也有一定的性能提升，具体分析参见GitHub（https://github.com/bytedance/byteps）。

除了在性能上超出目前其他所有分布式训练框架外，BytePS 可以兼容 Tensorflow、PyTorch、MXNet 等训练框架。BytePS 团队表示，开发者只需要非常少的改动，就可以使用 BytePS 框架进行分布式训练，享受 BytePS 带来的高性能。

此前行业里的 PS 实现，都是针对特定通用框架，例如专门为 TensorFlow 实现的 PS，也有专门为 MXNet 实现的 PS。

字节跳动人工智能实验室开源的 BytePS，通过实现一个通用的抽象层，抽象层可以被各种通用框架引用，实现了同时支持多个框架的可能性，因此能够支持 Tensorflow、PyTorch、MXNet 等行业主流训练框架。

BytePS 提供了 TensorFlow、PyTorch、 MXNet 以及 Keras 的插件，用户只要在代码中引用 BytePS 的插件，就可以获得高性能的分布式训练。BytePS 的核心逻辑则实现在 BytePS core 里。具体的通信细节完全由 BytePS 完成，用户完全不需要操心。

BytePS 团队表示，深度学习领域仍然有非常大的空间和可能性值得行业同仁们一起探索，开源 BytePS，是希望利用 BytePS 在性能和功能上的先进性，降低开发者和深度学习领域参与者们的门槛，帮助更多同道中人一起探索深度学习，提升 AI 应用效率。