前言  現有的分布式 DNN 訓練架構無法充分利用異構資源實現高性能訓練。近期，來自字節跳動和清華大學的研究人員提出一種新型分布式 DNN 訓練架構——BytePS，解決了這一問題，實現了大規模訓練性能的顯著提升。這項工作已在國際頂級計算機系統會議 OSDI’20 上發表，其開源代碼在 GitHub 上獲得 2400 stars。

背景

近年來 DNN 為計算機視覺、語音識別與合成、自然語言處理等領域帶來了突破性進展。然而，訓練這些 DNN 模型通常需要大量算力。大型公司一般會建設具備數千塊 GPU 卡的大型訓練集群，其核心資源是 GPU 機器，同時這些 GPU 機器也具備高端 CPU 資源。此外，也有一些僅具備 CPU 的機器，用于訓練數據預處理和生成等任務。這些 GPU/CPU 機器通過高速網絡連接，以加速分布式訓練的通信。目前工業界主流的分布式訓練是基于數據并行方式實現的，其中具有代表性的兩種架構是 All-reduce 和參數服務器（PS）。

研究動機：現有架構的設計缺陷

如下圖所示，目前的 All-reduce 和 PS 架構在訓練性能上距離最優情況都有較大差距。

圖 1：VGG16 訓練性能

究其原因，有如下三個方面的問題。

1、機器間網絡通信問題

在異構集群場景下，All-reduce 和 PS 架構對資源的利用情況如下圖所示。

1. All-reduce 架構中僅用到 GPU 機器，這是因為其設計假定了每個節點都是同構節點。迭代過程中，GPU 獨立計算模型參數的梯度，然后使用 All-reduce 通信聚合梯度；

2. PS 架構則包含 GPU worker 和 CPU server。迭代過程中，GPU worker 將梯度傳輸至 CPU server；后者將接收到的不同 workers 的梯度做聚合，然后執行 DNN 優化器（如 RMSProp 或 Adam 等）并將更新后的參數傳輸回 GPU workers。

圖 2：All-reduce 和 PS 架構示意圖

可以看出，All-reduce 的同構化設計導致其無法充分利用這種異構資源，即只有 GPU workers 之間通信，而無法利用其他 CPU 和帶寬資源。而 PS 雖然能夠利用 CPU 機器作為 server，卻可能在 CPU server 數量較少的時候產生流量熱點（例如形成多對一的情況），從而導致網絡擁塞。 

2、機器內多卡 PCIe 帶寬競爭問題

如今一臺訓練機器通常都具備有多張 GPU 卡（例如 4 或 8 卡）。在做機器間的通信前，機器內部的多 GPU 之間需要首先做一次本地通信，該通信過程一般是基于 PCIe 或 NVLink 鏈路。 我們觀察到，目前的網卡如 Mellanox CX5 帶寬已達 100Gbps，已經很接近 PCIe 3.0 x16 的 128Gbps 鏈路帶寬。而不幸的是，現在流行的機器內聚合方式（例如 8 卡直接做 all-reduce）會使 PCIe 成為瓶頸，導致網卡無法達到其 100Gbps 帶寬上限。即使對含有 NVLink 的拓撲，我們也能發現類似的 PCIe 競爭導致的瓶頸問題。 

3、通信鏈路中的 CPU 瓶頸問題

從前面的問題 1 中可以看出，相比于 All-reduce 而言，PS 架構實際上是存在更大的潛力的，因為它能充分利用異構 GPU/CPU 資源。然而，目前的 PS 甚至比 All-reduce 性能顯著低（圖 1），似乎與之矛盾。這是因為 PS 架構中還存在另外一種瓶頸限制了其性能 – CPU 瓶頸。 顧名思義，PS (參數服務器)需要將參數存儲在 CPU server 上，這就意味著需要將優化器 (如 Adam/RMSProp 等) 放在 Server 上去執行。然而，優化器通常包含復雜的數學運算，將會消耗大量的 CPU 內存帶寬。在 100Gbps 的網絡輸入情況下，CPU 將無法滿足將完整優化器放置在其上運行的需求。

BytePS 主要設計

針對以上三個問題，BytePS 逐一提出了解決方案。

1、機器間網絡通信優化

該優化涉及到對問題的建模，在給出具體的公式前，先從直觀上介紹其思想。

（1）由前文可知，PS 僅利用了 GPU 機器與 CPU 機器之間的帶寬。在 CPU 機器數量較少時，GPU 機器的帶寬 B 無法充分利用。圖 3 給出了一個例子，在這種情況下，GPU 機器僅能達到 2/3 的最大帶寬，剩余 1/3 帶寬未得到利用。

圖 3: PS 帶寬利用示意圖

（2）All-reduce 僅利用了 GPU 機器之間的帶寬。此時，GPU 機器與 CPU 機器之間的帶寬未得到利用。

圖 4：All-reduce 帶寬利用示意圖

（3）BytePS 結合兩者之長，同時利用了 GPU 與 GPU 之間、GPU 與 CPU 之間的帶寬，使得每臺機器的帶寬都能被充分利用。這就是 BytePS 機器間通信的思路。

圖 5: BytePS 從通信層面結合 PS 與 All-reduce

該思路在實現過程中，需要考慮如何分配 GPU 與 GPU 之間（設為 x%）、GPU 與 CPU 之間（設為 y%）的流量比例。經過計算，最優比例如下：

（其中 n 表示 GPU 機器的數量，k 表示 CPU 機器的數量，）

以上即為最優通信策略，對于不同的 n 與 k，采用該策略可使得機器間通信時間最小。

2、機器內多卡通信優化

目前市面上主流有兩種機器拓撲：PCIe-only 型 8 卡機器和 NVLink-based 型 8 卡機器。BytePS 針對這兩類拓撲提出了通用的解決方案和設計原則。

（1）PCIe-only 型拓撲

如圖 6 所示，標記 0-7 的灰框表示 GPU，P0 和 P1 表示 PCIe switch。現實當中， P0-CPU0 以及 P1-CPU1 是帶寬最小的鏈路，因此優化目標是最小化這條鏈路上傳輸的數據量。目前主流的做法是對這 8 卡做一次全局 All-reduce，這樣 P0-CPU0 需要傳輸的數據量是 7M/4（根據 All-reduce 的通信量計算得出），其中 M 是每張卡上的梯度大小。注意到該做法并沒有利用 CPU 的計算能力。

圖 6：PCIe-only 型拓撲

BytePS 的核心思想是利用 CPU 的計算能力減少瓶頸鏈路的傳輸數據量。如圖 7 所示，首先每個 PCIe switch 下的 4 張卡先進行一次 Local Reduce-scatter。該步驟之后，每張卡上已聚合的梯度大小為 M/4。

圖 7：PCIe-only 拓撲解決方案步驟（1）

下一步，每張卡將自身已聚合的 M/4 梯度拷貝到主機的內存上，如圖 8 所示。注意這個過程使得 P0-CPU0 只傳輸了 M/4*4=M 的流量。

圖 8：PCIe-only 拓撲解決方案步驟（2）

此時，兩個 NUMA node 的主機內存上各自有一份大小為 M 的梯度。我們再利用 CPU 將這兩份梯度做一次聚合。但是這個過程的傳輸只發生在帶寬較大的 QPI 上（>300Gbps），并不會產生瓶頸。于是這一系列步驟不但實現了預期中的梯度聚合效果，還使瓶頸鏈路的傳輸量從 7M/4 降低為 M，顯著降低了通信時間。這里的核心設計原則是：盡量避免跨 NUMA GPU 的直接通信，而可以利用 CPU 的聚合能力來間接完成。

（2）NVLink-based 拓撲

圖 9 是 NVLink-based 機型的示意圖。對于這種拓撲，GPU 之間可以通過超高帶寬的 NVLink 鏈路進行通信。由于 NVLink 帶寬顯著大于 PCIe 帶寬，PCIe 瓶頸問題顯得更加嚴重。可以看到，圖 9 中 P0-CPU0 鏈路（標紅色的線段）會同時被以下兩種傳輸同時競爭：（1）CPU0 的內存往 GPU0/GPU1 拷貝數據；（2）CPU0 的內存往網卡 NIC 發送數據。由于 P0-CPU0 鏈路的帶寬與網卡帶寬很接近，這種競爭會導致網卡無法發揮最大帶寬。

圖 9：NVLink-based 拓撲示意圖

為解決這一競爭問題，BytePS 利用了 NVLink 帶寬顯著高于 PCIe 鏈路的事實，利用 Reduce（而非 Reduce-scatter）方式避免 PCIe 競爭。如圖 10 中紅線所示，所有卡先將其梯度通過 NVLink 傳輸至 GPU2 上并做 Reduce，接著 GPU2 將聚合后的梯度拷貝到 CPU0 內存，再經由網卡發送出去。由于 NVLink 帶寬很高，這種做法不會導致 GPU2 產生流量熱點問題，但卻能夠避免在 P0-CPU0 鏈路上發生的競爭。

圖 10：BytePS 對 NVLink-based 拓撲的解決方案

3、CPU 瓶頸優化：Summation Service

前文提到，優化器對于 CPU 而言是比較重的任務，這也是 PS 架構的性能缺陷之一。然而，如何高效利用 CPU 的異構計算能力是 BytePS 的核心訴求之一，這就需要克服數據同步過程中的 CPU 瓶頸。

經分析，優化器可被拆解為兩部分：（1）Sum：將來自其他 GPU workers 的梯度求和并得到一份聚合后的新梯度；（2）Update：利用新梯度對參數進行更新。后者對于 CPU 而言的確是非常消耗內存帶寬的操作，但前者卻能夠在 CPU 上高效實現（例如 AVX 指令集）。如圖 11 所示，求和操作在 CPU 上可以達到遠超網絡帶寬的吞吐率，即不會引入 CPU 瓶頸。

圖 11：求和操作在 CPU 上的吞吐率

受到這個發現的啟發，BytePS 提出了 Summation Service 概念，對傳統 PS 的 CPU 瓶頸問題做了改進。如圖 12 所示，不同于 PS 將完整優化器放置在 CPU 上的設計，Summation Service 只將 Sum 操作放置在 CPU 上，而將 Update 操作交由計算能力更強大、內存帶寬更充足的 GPU 來執行。這種設計能夠避免同步過程中的 CPU 瓶頸。

圖 12：PS 架構與 Summation Service 的對比

總體系統架構

將前述 BytePS 三個設計點結合起來，我們得到 BytePS 完整系統架構，如下圖所示。 

圖 13：BytePS 系統架構

1、每臺 GPU 機器上部署了一個 Communication Service 模塊，負責聚合本地多卡的梯度（即機器內多卡通信），其能充分考慮機器內部復雜的拓撲，避免產生 PCIe 瓶頸。

2、每臺 GPU/CPU 機器上部署了一個 Summation Service 模塊，處理來自其他 GPU 機器的梯度，其能夠高效運行在 CPU 上。

3、模塊之間通過網絡互連，通信策略使用的是前述設計中提到的最優網絡通信方案。經證明，該方案不僅有最佳的性能，且能夠從通信角度統一 All-reduce 和 PS 兩種架構。

系統實現與優化

BytePS 在實現中充分利用了流水線的思想，盡可能地將計算和 PCIe 和網絡傳輸重疊起來。此外，BytePS 還利用了一些 RDMA 實現上的優化技巧（例如內存頁對齊等），解決了實際部署時遇到的 RDMA slow receiver 問題，實現了高性能的網絡傳輸。

對于用戶側，BytePS 提供了豐富的用戶接口，能夠兼容 TensorFlow，PyTorch 和 MXNet 等主流深度學習框架。通常只需要幾行至十幾行代碼的修改，就可將現有基于其他框架（如 Horovod 或 PyTorch DDP 等）的代碼遷移至 BytePS 上運行。

性能評估

BytePS 對多種 CV 類（包括 ResNet-50，VGG-16 ，GAN）和 NLP 類（Transformer，BERT-Large，GPT-2）模型都做了分布式性能評測，規模從 8 卡 - 256 卡。所使用的硬件是 V100 GPU 和 100Gbps RDMA 網絡。對照組為目前廣泛使用的 All-reduce 和原生 PS 實現。

圖 14 和圖 15 分別展示了 CV 和 NLP 模型上的評估結果。總體而言，BytePS 在各類模型上都取得了正向收益，且相比于 All-reduce 和 PS 能夠達到的最大提升幅度達 84% 和 245%。

圖 14：CV 類模型性能評估

圖 15：NLP 類模型性能評估

／End.