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

首先奉上BytePS开源项目地址：https://github.com/bytedance/byteps

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团队重新思考了最佳通信策略，不仅在机器内使用 NCCL，同时也重新部署了机器间的通信方式。

据介绍，在服务器内，GPU是插在不同的PCIe switch上的，相同PCIe switch内的GPU通信带宽较高，跨PCIe switch的通信带宽就较小。NUMA是指服务器上有不止一颗CPU，CPU内存也有类似问题：同CPU的内存访问带宽高，跨CPU的内存访问带宽低。BytePS会根据这些信息，有选择地分配数据在CPU和GPU中的内存位置，以及哪块内存和哪块内存通信，从而最大化通信带宽。

BytePS构架本身也做了一些重要设计，使得PS架构理论上的潜能得以实现，包括：Tensor自动切分、多级灵活流水线处理、网络通信优先级调度、ZeroMQ优化、共享内存zero-copy、RDMA实现和PS端多队列多线程优化。

更详细的实现原理参见这里。

性能表现

测试中，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。

除了在性能上超出目前其他所有分布式训练框架外，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 前，假设你已经安装了以下一种或更多框架：TensorFlow、Keras、PyTorch、MXNet 等。BytePS主要基于 CUDA 和 NCCL。

复制 BytePS 和第三方依赖：

git clone --recurse-submodules https://github.com/bytedance/byteps

进入 BytePS 文件目录，并安装：

python setup.py install

注意：你可能需要设置 BYTEPS_USE_RDMA=1 来安装 RDMA 支持。

现在你可以试试我们已有的示例。假设你使用 MXNet，并想尝试 Resnet50训练基准。

export NVIDIA_VISIBLE_DEVICES=0,1 \
       DMLC_NUM_WORKER=1 \
       DMLC_NUM_SERVER=1 \
       DMLC_WORKER_ID=0 \
       DMLC_ROLE=worker \
       DMLC_PS_ROOT_URI=10.0.0.1 \
       DMLC_PS_ROOT_PORT=1234 \
       DMLC_INTERFACE=eth0

python byteps/launcher/launch.py byteps/example/mxnet/train_imagenet_byteps.py --benchmark 1 --batch-size=32

对于分布式训练，你可能需要建立一个服务器镜像。研发团队提供了 Docker 文件作为例子。你可以将同样的镜像用于调度和服务器。

关于如何启动分布式任务的内容和更多上手教程可参考相关文档。

如何在已有代码中使用BytePS

虽然内核设计有所不同，但BytePS 和 Horovod 接口高度兼容，我们希望通过Horovod接口减少用户测试BytePS的工作量。

如果你的任务只依赖于Horovod 的 allreduce 和广播，你可以在一分钟内切换到 BytePS。只需要用 import byteps.tensorflow as bps 替换 import horovod.tensorflow as hvd，并将代码中所有的 hvd 替换成 bps即可。

BytePS的局限和未来计划

BytePS 目前不支持单纯的 CPU 训练，其中一个原因是BytePS的部分底层逻辑可能无法支持。因此你需要使用 CUDA 或 NCCL 来构建和运行 BytePS。

未来BytePS 计划增加以下特性：

稀疏模型训练
异步训练
容错机制
延迟减缓

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

创作场景

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