位元組跳動開源高性能分散式訓練框架 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 應用效率。