近日，字節(jié)跳動(dòng)人工智能實(shí)驗(yàn)室宣布開源一款高性能分布式深度學(xué)習(xí)訓(xùn)練框架 BytePS，在性能上顛覆了過去幾年 allreduce 流派一直占據(jù)上風(fēng)的局面，超出目前其他分布式訓(xùn)練框架一倍以上的性能，且同時(shí)能夠支持 Tensorflow、PyTorch、MXNet 等開源庫。

BytePS 結(jié)合了字節(jié)跳動(dòng)人工智能實(shí)驗(yàn)室?guī)讉€(gè)月來對分布式訓(xùn)練通信的多個(gè)研究與優(yōu)化成果，包含通信優(yōu)先級調(diào)度、PS 的 RDMA 實(shí)現(xiàn)、針對 PCIe switch 與 NUMA 的優(yōu)化，以及 BytePS 本身構(gòu)架的創(chuàng)新等。

深度學(xué)習(xí)的效果取決于模型與數(shù)據(jù)，目前行業(yè)內(nèi)不斷刷新深度學(xué)習(xí)準(zhǔn)確率的最新研究，大多都基于更大的模型以及更大的數(shù)據(jù)集。然而，大模型與大數(shù)據(jù)對訓(xùn)練時(shí)的計(jì)算能力提出了極高要求，單張 GPU 卡，或者單臺(tái)服務(wù)器上的 GPU 卡，已經(jīng)遠(yuǎn)遠(yuǎn)不能夠滿足內(nèi)部訓(xùn)練任務(wù)的需求。因此，分布式訓(xùn)練的效率，即使用多臺(tái)服務(wù)器協(xié)同進(jìn)行訓(xùn)練，現(xiàn)在成為了深度學(xué)習(xí)系統(tǒng)的核心競爭力。

一直以來，在分布式訓(xùn)練中有兩大流派，分別是 allreduce 和 PS（Parameter Server）。過去三年中，尤其是百度提出 allreduce，以及 Uber 開源基于 allreduce 的 Horovod 之后，行業(yè)內(nèi)的認(rèn)知中，allreduce 是最好的分布式訓(xùn)練通信方式，而過去的 PS 實(shí)現(xiàn)的性能也確實(shí)與 allreduce 存在一定差距。

BytePS 則顛覆了 allreduce 長期領(lǐng)先的局面，BytePS 擁有著超出目前其他所有分布式訓(xùn)練框架一倍以上的性能，包括 NVIDIA 開源的 NCCL，Uber 開源的 Horovod，以及 Tensorflow、PyTorch、MXNet 自帶的分布式訓(xùn)練方案等。

BytePS 開發(fā)團(tuán)隊(duì)表示，在公有云或者私有云這類共享集群中，經(jīng)過精巧設(shè)計(jì)和高質(zhì)量實(shí)現(xiàn)的 PS，PS 架構(gòu)不僅不比 allreduce 差，而且在一些環(huán)境還能得到比 allreduce 還高一倍的速度。

測試中，BytePS 團(tuán)隊(duì)使用了公有云上的虛擬機(jī)，每個(gè)虛擬機(jī)有 8 張 Tesla V100 16GB GPU，GPU 之間通過 NVLink 進(jìn)行高速互連。每個(gè) GPU 上的 batch size 選取為 64。虛擬機(jī)之間通過 20Gbps 的 TCP/IP 網(wǎng)絡(luò)進(jìn)行連接。在這種情況下，由于機(jī)器之內(nèi)帶寬足夠大，TCP/IP 的網(wǎng)絡(luò)帶寬則成為了主要瓶頸。

BytePS 選擇了 Resnet50 和 VGG16 兩個(gè)模型進(jìn)行評測，其中 Resnet50 是計(jì)算密集型的模型（對通信要求低，優(yōu)化空間?。?，VGG16 是通信密集型的模型（對通信要求高，優(yōu)化空間大），對照組選擇了目前市面上最流行的通信框架之一 Horovod-NCCL（基于 allreduce 算法實(shí)現(xiàn)），性能指標(biāo)為每秒鐘訓(xùn)練的 ImageNet 圖片數(shù)量，越高代表越好。

通過兩組實(shí)驗(yàn)結(jié)果可以看出，對于計(jì)算密集型的 Resnet50 模型，BytePS 性能超過 Horovod-NCCL 近 44%；而對于通信密集型的 VGG16 模型，BytePS 性能可以超過 Horovod-NCCL 將近 100%。

BytePS 團(tuán)隊(duì)也在配有 100Gbps 的 RDMA 網(wǎng)絡(luò)的私有集群做了測試，BytePS 也有一定的性能提升，具體分析參見 Github（ https://github.com/bytedance/byteps）。

除了在性能上超出目前其他所有分布式訓(xùn)練框架外，BytePS 可以兼容 Tensorflow、PyTorch、MXNet 等訓(xùn)練框架。BytePS 團(tuán)隊(duì)表示，開發(fā)者只需要非常少的改動(dòng)，就可以使用 BytePS 框架進(jìn)行分布式訓(xùn)練，享受 BytePS 帶來的高性能。

此前行業(yè)里的 PS 實(shí)現(xiàn)，都是針對特定通用框架，例如專門為 TensorFlow 實(shí)現(xiàn)的 PS，也有專門為 MXNet 實(shí)現(xiàn)的 PS。

字節(jié)跳動(dòng)人工智能實(shí)驗(yàn)室開源的 BytePS，通過實(shí)現(xiàn)一個(gè)通用的抽象層，抽象層可以被各種通用框架引用，實(shí)現(xiàn)了同時(shí)支持多個(gè)框架的可能性，因此能夠支持 Tensorflow、PyTorch、MXNet 等行業(yè)主流訓(xùn)練框架。

BytePS 提供了 TensorFlow、PyTorch、 MXNet 以及 Keras 的插件，用戶只要在代碼中引用 BytePS 的插件，就可以獲得高性能的分布式訓(xùn)練。BytePS 的核心邏輯，則實(shí)現(xiàn)在 BytePS core 里。具體的通信細(xì)節(jié)，完全由 BytePS 完成，用戶完全不需要操心。

BytePS 團(tuán)隊(duì)表示，深度學(xué)習(xí)領(lǐng)域仍然有非常大的空間和可能性值得行業(yè)同仁們一起探索，開源 BytePS，是希望利用 BytePS 在性能和功能上的先進(jìn)性，降低開發(fā)者和深度學(xué)習(xí)領(lǐng)域參與者們的門檻，幫助更多同道中人一起探索深度學(xué)習(xí)，提升 AI 應(yīng)用效率。