字節跳動開源高性能分佈式訓練框架BytePS：兼容主流訓練框架

近日，字節跳動人工智能實驗室宣佈開源一款高性能分佈式深度學習訓練框架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應用效率。