字節(jié)跳動開源高性能分布式訓(xùn)練框架BytePS:兼容TensorFlow等主流訓(xùn)練框架

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

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

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

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

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

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

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

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

通過兩組實驗結(jié)果可以看出,對于計算密集型的Resnet50模型,BytePS性能超過Horovod-NCCL近44%;而對于通信密集型的VGG16模型,BytePS性能可以超過Horovod-NCCL將近100%。

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

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

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

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

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

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

生成海報

免責聲明：本網(wǎng)站內(nèi)容主要來自原創(chuàng)、合作伙伴供稿和第三方自媒體作者投稿，凡在本網(wǎng)站出現(xiàn)的信息，均僅供參考。本網(wǎng)站將盡力確保所提供信息的準確性及可靠性，但不保證有關(guān)資料的準確性及可靠性，讀者在使用前請進一步核實，并對任何自主決定的行為負責。本網(wǎng)站對有關(guān)資料所引致的錯誤、不確或遺漏，概不負任何法律責任。任何單位或個人認為本網(wǎng)站中的網(wǎng)頁或鏈接內(nèi)容可能涉嫌侵犯其知識產(chǎn)權(quán)或存在不實內(nèi)容時，應(yīng)及時向本網(wǎng)站提出書面權(quán)利通知或不實情況說明，并提供身份證明、權(quán)屬證明及詳細侵權(quán)或不實情況證明。本網(wǎng)站在收到上述法律文件后，將會依法盡快聯(lián)系相關(guān)文章源頭核實，溝通刪除相關(guān)內(nèi)容或斷開相關(guān)鏈接。