近日,字节跳动人工智能实验室宣布开源一款高性能分布式深度学习训练框架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应用效率。
