字節跳動開源分布式訓練框架BytePS，登上GitHub熱榜

問耕 發自 凹非寺

量子位 出品 | 公眾號 QbitAI

字節跳動開源了通用分布式訓練框架BytePS，這個框架支持TensorFlow、Keras、PyTorch、MXNet，可以運行在TCP或RDMA網絡中。

官方介紹稱，BytePS大大優於現有的開源分布式訓練框架。例如，在相同的條件下，BytePS的訓練速度是Horovod+NCCL的兩倍。

BytePS也登上了GitHub趨勢熱榜。

性能表現

為了展示BytePS的性能，官方測試了兩個模型：VGG16（通信密集型）和Resnet50（計算密集型），兩個模型都以fp32精度進行訓練。

訓練使用了公有雲上的虛擬機，每個機器有8個GPU，啟用了NVLink。這些GPU都是Tesla V100 16GB型號GPU，batch size設置為64/每GPU。硬體之間以20Gbps的TCP/IP網絡連接。

測試表明，與Horovod (NCCL)相比，BytePS在VGG16訓練上的速度快100%，Resnet訓練速度快44%。

Horovod是Uber開源的深度學習工具，NCCL是NVIDIA發布的多卡通信框架。

原因解析

BytePS為什麽比Horovod更強？

官方解釋稱，一個主要的原因是BytePS是專門為雲和共享集群而設計，並且拋棄了MPI。MPI是一個跨語言的通訊協議，用於編寫並行計算機。

MPI是為高性能計算機而生，對於使用同類硬體構建的集群以及運行單一任務更有效。但並不是雲（或者內部共享集群）的最佳選擇。

因此字節跳動團隊重新思考了最佳的通信策略。簡而言之，BytePS僅在機器內部使用NCCL，重新構建了機器內部的通信。

BytePS還集成了許多加速技術，例如分層策略、流水線、張量分區、NUMA感知本地通信、基於優先級的調度等等。

傳送門

BytePS現已在GitHub上開源，地址如下：

https://github.com/bytedance/byteps

在GitHub上，還詳細介紹了如何使用BytePS等相關信息。

目前BytePS還不支持純CPU訓練。