近日,字節(jié)跳動人工智能實(shí)驗(yàn)室宣布開源一款高性能分布式深度學(xué)習(xí)訓(xùn)練框架BytePS,在性能上顛覆了過去幾年allreduce流派一直占據(jù)上風(fēng)的局面,超出目前其他所有分布式訓(xùn)練框架一倍以上的性能,且同時(shí)能夠支持Tensorflow、PyTorch、MXNet等開源庫。
BytePS結(jié)合了字節(jié)跳動人工智能實(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卡,或者單臺服務(wù)器上的GPU卡,已經(jīng)遠(yuǎn)遠(yuǎn)不能夠滿足內(nèi)部訓(xùn)練任務(wù)的需求。因此,分布式訓(xùn)練的效率,即使用多臺服務(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ā)者只需要非常少的改動,就可以使用BytePS框架進(jìn)行分布式訓(xùn)練,享受BytePS帶來的高性能。
此前行業(yè)里的PS實(shí)現(xiàn),都是針對特定通用框架,例如專門為TensorFlow實(shí)現(xiàn)的PS,也有專門為MXNet實(shí)現(xiàn)的PS。
字節(jié)跳動人工智能實(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)用效率。
