近日，快手和蘇黎世理工宣布開源分布式訓(xùn)練框架 Bagua（八卦），相比于 PyTorch、TensorFlow 等現(xiàn)有深度學習開源框架僅針對系統(tǒng)層面進行優(yōu)化，Bagua 突破了這一點，專門針對分布式場景設(shè)計了特定的優(yōu)化算法，實現(xiàn)了算法和系統(tǒng)層面的聯(lián)合優(yōu)化，性能較同類提升 60%。

研發(fā)背景

隨著摩爾定律的失效，單個計算單元的能力已經(jīng)遠遠無法滿足數(shù)據(jù)的指數(shù)級增長。比如，快手每天上傳的新視頻超過千萬條，即便訓(xùn)練簡單的分類模型（比如 ResNet），使用單機單卡的算力，訓(xùn)練快手日內(nèi)新增視頻都需要超過一百天的時間。因此，在數(shù)據(jù)爆炸性增長的互聯(lián)網(wǎng)行業(yè)，多機多卡的并行訓(xùn)練成為了大數(shù)據(jù)時代的必然。隨著深度學習模型功能的日益強大，分布式訓(xùn)練任務(wù)的通信成本和所需算力也隨之急劇增長。

然而，由于多機多卡并行帶來的額外通訊成本，加速比（speedup）經(jīng)常讓大家失望，從而形成了大廠“堆資源”，沒資源的“干瞪眼”的局面。比如，Google 的 Downpour 框架 [1] 使用 80 個 GPU 訓(xùn)練 ImageNet，加速比卻只有 12/80=15%。因此如何提升多機多卡中訓(xùn)練的通訊效率成為了并行訓(xùn)練乃至解決數(shù)據(jù)爆炸性增長的核心問題之一。

項目 GitHub 地址：https://github.com/BaguaSys/bagua

現(xiàn)有的深度學習開源框架（PyTorch，TensorFlow）主要針對系統(tǒng)層面優(yōu)化，把已有的單機單卡優(yōu)化算法擴展到多機多卡的場景。雖然系統(tǒng)層面的優(yōu)化使得并行效率不斷提升，但是邊際效益卻越來越明顯。針對這個問題，快手和蘇黎世理工（ETH Zürich）聯(lián)合開發(fā)了一款名為“Bagua”的分布式訓(xùn)練框架，突破單純的系統(tǒng)層面優(yōu)化，專門針對分布式的場景設(shè)計特定的優(yōu)化算法，實現(xiàn)算法和系統(tǒng)層面的聯(lián)合優(yōu)化，極致化分布式訓(xùn)練的效率。用戶只需要添加幾行代碼，便能把單機單卡訓(xùn)練擴展到多機多卡訓(xùn)練并得到非?？捎^的加速比。

Bagua 設(shè)計思路

從單機單卡的訓(xùn)練到多機多卡訓(xùn)練的核心，是每個卡把自己的計算結(jié)果進行累加和傳播。這個過程好比每個人把自己知道的信息傳遞給他人，然后又從其他人那里獲取信息，最后完成全局的信息同步。如果把計算單元之間的信息同步類比為人與人之間的信息同步，那么社會實踐經(jīng)驗告訴我們，“八卦”可能是消息傳遞最高效的模式?！鞍素浴毕鞑ゾ哂腥ブ行幕?、異步通訊、信息壓縮的特點，這與 Bagua 里面實現(xiàn)的通訊算法剛好一一呼應(yīng)。

為了提升分布式訓(xùn)練效率，Bagua 實現(xiàn)了自研以及前沿的算法，包括去中心化/中心化、同步/異步以及通訊壓縮等基礎(chǔ)通訊組件，通過軟硬結(jié)合的設(shè)計極致優(yōu)化了這些組件的效率，并且靈活支持這些算法的組合，以及更復(fù)雜的算法設(shè)計。

Bagua 將通信過程抽象成了如下的算法選項：

中心化或是去中心化（Centralized or Decentralized）：在中心化的通訊模式中，梯度或模型的同步過程需要所有的工作節(jié)點進行參與，因此，較高的網(wǎng)絡(luò)延時往往會導(dǎo)致訓(xùn)練效率的降低。去中心化的通信模式 [5,6] 往往可以有效的解決這一問題：在該模式下，工作節(jié)點可以被連接成特定的拓撲結(jié)構(gòu)（例如環(huán)），在通信過程中，每一個工作節(jié)點只與和它相鄰的節(jié)點進行通信。

同步或是異步（Synchronous or Asynchronous）：同步模式中，在每一次迭代過程中，所有工作節(jié)點都需要進行通信，并且下一步迭代必須等待當前迭代的通信完成才能開始。反之，異步式分布算法 [2] 則不需要等待時間：當某個節(jié)點完成計算后就可直接傳遞本地梯度，進行模型更新。

完整精度模式或信息壓縮模式（Full-Precision or Low-Precision）：完整精度模式，會使用與本地模型相同的 32 位浮點數(shù)（float32）進行傳輸。另一方面，在通訊存在瓶頸的情況下，基于大量已有研究通過量化 (quantization [3]) 或稀疏化 (sparsification [4]) 等方法壓縮梯度，再用壓縮后的梯度更新參數(shù)。在很多場景下，可以達到和完整精度相同的精度，同時提升通訊效率。

雖然為了提升通訊效率，Bagua 沒有依照傳統(tǒng)的方式同步所有計算節(jié)點的結(jié)果，甚至每次同步的信息還有偏差，但是得益于最新理論上的進展，這幾種通訊策略以及他們的組合最終收斂解的正確性和效率仍然能得到充分保證，而且計算復(fù)雜度跟同步中心化和信息無損的方法相當，但是通訊效率更高 [10]。

值得注意的是，在實踐中，分布式訓(xùn)練算法往往會使用不止一種上述的優(yōu)化方法，從而適配更為極端的網(wǎng)絡(luò)環(huán)境 [7,8,9]。對于分布式算法感興趣的讀者，我們在這里推薦一份最新的完整綜述報告 [10]。Bagua 提供了一套詳盡的通信模式來支持用戶在上述模式中任意選擇組合，我們將這一分布式訓(xùn)練系統(tǒng)對于上述算法選項的支持情況總結(jié)在下表中：

從表格中不難看出，現(xiàn)有框架的優(yōu)化只是針對較為通用的算法（中心化同步完整精度），對于其他的算法組合，這些系統(tǒng)的支持非常有限。對于中心化同步進行信息壓縮，這些系統(tǒng)往往只能支持較為簡單的 float32->float16 壓縮，相較而言，Bagua 則可以支持更為復(fù)雜的 ByteGrad，QAdam 等算法。對于其他的算法組合，現(xiàn)有的框架通常無法支持，而 Bagua 則可以自由支持。

然而，簡單地支持這項算法選項并不能直接在大規(guī)模集群上帶來性能的提升。Bagua 的核心優(yōu)勢在于，為了追求極致化的性能，而實現(xiàn)算法和實現(xiàn)的聯(lián)合優(yōu)化。具體來講，基于上述的通信層抽象，用戶既可以方便得選擇系統(tǒng)提供的各種算法組合從而獲得性能提升，又能靈活得實現(xiàn)新的分布式 SGD 算法 —— Bagua 將自動為這一算法實現(xiàn)提供系統(tǒng)層優(yōu)化。這些系統(tǒng)優(yōu)化包含：

將通訊時間隱藏在計算時間中：為了降低通信開銷，Bagua 能夠?qū)⒉糠滞ㄐ艜r間隱藏在計算時間中。具體來講，在反向梯度的計算過程中，部分已經(jīng)完成的梯度可以在剩余梯度的計算過程中同時進行通信——通過這種流水的處理方式，部分通信時間可以被有效地“隱藏”在反向梯度的計算過程中，從而減小數(shù)據(jù)并行帶來的通信開銷。

l  參數(shù)分桶及其內(nèi)存管理：頻繁的傳輸碎片化數(shù)據(jù)，會降低通信的效率。因此，Bagua 將模型參數(shù)分割成桶，并且分配連續(xù)的內(nèi)存空間來對每一個桶進行管理，這樣通訊的單位就變成了桶，從而能夠更高效地利用通信模型。此外，由于支持了信息壓縮算法，對于壓縮和解壓的函數(shù)，其操作的基本單位也是桶，這樣也能使得這些操作的開銷降低。

l  分層化的通信實現(xiàn)：由于工業(yè)級別的分布式訓(xùn)練往往需要多機多卡，而不同物理連接方式所帶來的延時和帶寬也有較大差異，因此，通訊的有效抽象也對性能的提升至關(guān)重要。Bagua 將涉及多機的通信抽象成：“機內(nèi)”和“機間”，并對于相應(yīng)的通信抽象做了優(yōu)化。例如，對于信息壓縮傳輸，分層化通訊將會把這一算法解讀成“機內(nèi)”完整精度，“機間”信息壓縮，從而為不同的物理鏈接提供最合適的通信算法。

我們想要強調(diào)的是，這些系統(tǒng)實現(xiàn)層面的優(yōu)化是對于各種算法組合廣泛適用，而非局限在某一特定的算法設(shè)置上。因此，所有的系統(tǒng)優(yōu)化都可以被靈活的復(fù)用到各種算法實現(xiàn)中去，這在保證“端到端”的性能提升的同時，也為開發(fā)新的分布式算法提供了良好的平臺。

經(jīng)過實驗，Bagua 的特點如下：

l  并行性能顯著提高：在 ImageNet 數(shù)據(jù)集上，相較當前開源分布式框架（PyTorch-DDP，Horovod，BytePS），當配置同樣的算力（128GPU）與通信網(wǎng)絡(luò)（100Gbps），達到相同的訓(xùn)練精度，Bagua 只需其他框架 80% 左右的時間；

l  對網(wǎng)絡(luò)環(huán)境更魯棒：由于有效的支持了各類算法優(yōu)化（信息壓縮，異步，和去中心化），Bagua 在各類網(wǎng)絡(luò)環(huán)境下（包括不同延時和帶寬）都體現(xiàn)出了良好的適配性。尤其是在高延遲低帶寬的情況下，Bagua 體現(xiàn)出比其他框架更優(yōu)的加速比，比如：在 10Gbps 網(wǎng)絡(luò)帶寬環(huán)境下，同樣的 ImageNet 任務(wù)，Bagua 只需其他框架 50% 左右的訓(xùn)練時間來達到同樣的訓(xùn)練精度；

l  “一鍵式”使用：Bagua 對于端用戶非常友好，現(xiàn)有利用 PyTorch 的模型都可以作為 Bagua 的輸入，Bagua 將自動為其提供豐富的并行方案——只需增加幾行代碼，訓(xùn)練就可以運行在分布式集群上；

l  分布式通訊算法易拓展性：Bagua 提供了針對算法的高拓展性，對于分布式優(yōu)化算法的開發(fā)者，Bagua 提供了有效的通訊抽象，開發(fā)者實現(xiàn)的新算法也可以直接復(fù)用 Bagua 的系統(tǒng)優(yōu)化；

l  可用于工業(yè)級場景大規(guī)模使用：Bagua 為 Kubernetes 實現(xiàn)了定制化的 operator，支持云原生部署，同時考慮機器資源和故障問題，有機結(jié)合 PyTorch Elastic 和 Kubernetes 實現(xiàn)了容災(zāi)功能和動態(tài)訓(xùn)練擴縮容。用戶可以通過使用 Bagua ，在少量機器空閑時就開始訓(xùn)練，在更多機器資源釋放的情況下，訓(xùn)練任務(wù)自動擴容到更多機器。同時機器節(jié)點損壞時，自動剔除壞節(jié)點繼續(xù)訓(xùn)練。方便工業(yè)級訓(xùn)練場景使用，也方便與機器學習平臺結(jié)合使用；

l  安全、故障易排查：Bagua 通訊后端由注重內(nèi)存安全、速度和并發(fā)性的 Rust 語言實現(xiàn)，在編譯期就排除了大量的內(nèi)存安全問題。同時基于 tracing 實現(xiàn)了分模塊、分層級的 log 輸出，使得實際場景中故障排查更加輕松。

此外，Bagua 在快手內(nèi)部也經(jīng)過了工業(yè)級任務(wù)的實踐檢驗，Bagua 已經(jīng)在快手內(nèi)部多個核心業(yè)務(wù)場景投入使用，相較其他開源框架取得了顯著的性能提升：

l  大規(guī)模自然語言處理（GPT2-xl 量級大小的模型），提升效率 65%

l  大規(guī)模圖像識別（10+ 億圖像 / 視頻），提升效率 20%~30%

l  大規(guī)模語音識別（TB 級別語音資料），提升效率 20%~30%

l  大規(guī)模推薦系統(tǒng)（萬億級別參數(shù)模型訓(xùn)練，支持億級別 DAU 的應(yīng)用），提升效率 100% 以上

Bagua 和其他開源方案的性能對比

快手選用了包括圖像，文本，語音，圖像文本混合等多個訓(xùn)練任務(wù)對 Bagua 的性能進行測試，并與 PyTorch-DDP，Horovod，BytePS 進行比較。得益于 Bagua 系統(tǒng)的高效性和算法的多樣性，Bagua 可以在不同任務(wù)中選用相應(yīng)最優(yōu)的算法，從而保證在訓(xùn)練精度和其他系統(tǒng)持平的前提下，訓(xùn)練速度明顯提高。值得注意的是，當網(wǎng)絡(luò)狀況不佳時，Bagua 系統(tǒng)的優(yōu)勢將更加明顯。下面我們選取 GPT2-XL，BERT-Large 和 VGG16 三個典型的通信密集型任務(wù)進行對比說明，更多結(jié)果可在 Bagua 論文和網(wǎng)站中（https://github.com/BaguaSys/bagua）進一步了解。

1. End-to-end 訓(xùn)練時間對比

下圖展示了在 128 個 V100 GPU 上 fine-tune BERT-Large （SQuAD 數(shù)據(jù)集），模型 F1 精度隨訓(xùn)練時間的變化曲線。Bagua 使用 QAdam-1bit 算法加速，機器之間采用 100Gbps TCP/IP 網(wǎng)絡(luò)互聯(lián)。我們可以看到，即使在高速網(wǎng)絡(luò)下，達到相同的訓(xùn)練精度，Bagua 需要的時間僅為其他系統(tǒng)的 60%。

 2. 擴展性對比

下圖展示了各個系統(tǒng)在 VGG16 模型（ImageNet 數(shù)據(jù)集）上的訓(xùn)練速度與 GPU 數(shù)量之間的關(guān)系。在測試中分別使用了 1，8，16，32，64，128 個 V100 GPU 進行測試。該任務(wù)中 Bagua 使用 8bitsGrad 算法加速。可以看出 Bagua 的擴展效率相比其他系統(tǒng)有較明顯的提升。

下圖展示了各個系統(tǒng)在 GPT2-XL 模型上的訓(xùn)練速度與 GPU 數(shù)量之間的關(guān)系。GPT2-XL 模型有 1.5 billion 參數(shù)，在大模型中具有一定代表性。Bagua 使用 8bitsGrad 算法加速。在測試中分別使用了 8，16，32，64，80 個 V100 GPU 進行測試。同樣可以看出 Bagua 的擴展效率相比其他系統(tǒng)有較明顯的提升。

3. 不同網(wǎng)絡(luò)環(huán)境對比

算法是 Bagua 的靈魂。當網(wǎng)絡(luò)環(huán)境變化時，不同的算法會表現(xiàn)出不同的性能特點。在下圖中，我們以 BERT-Large fine-tune 為例，調(diào)整機器之間網(wǎng)絡(luò)的帶寬和延遲，對比 Bagua 中各個算法的 epoch 時間。可以看出，隨著帶寬降低，壓縮算法的優(yōu)勢會越來越明顯，且和壓縮程度相關(guān)；當延遲逐漸升高，去中心化算法逐漸展現(xiàn)出優(yōu)勢。除此之外，當網(wǎng)絡(luò)環(huán)境不佳時，Bagua 對比其他系統(tǒng)的優(yōu)勢也進一步擴大。

Bagua 使用實例

在已有的訓(xùn)練腳本中使用 Bagua 非常簡單，在代碼中算法使用者只需要增加如下幾行代碼對已有模型進行初始化操作即可。以使用 GradientAllReduce 算法為例：

首先，我們需要 import 一下 bagua

import bagua.torch_api as bagua

隨后，我們可以初始化 Bagua 的進程組:

torch.cuda.set_device(bagua.get_local_rank())bagua.init_process_group()

對于數(shù)據(jù)集的初始化，Bagua 完全兼容 PyTorch 的實現(xiàn)：

train_dataset = ...test_dataset = ...train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset,    num_replicas=bagua.get_world_size(), rank=bagua.get_rank())train_loader = torch.utils.data.DataLoader(    train_dataset,    batch_size=batch_size,    shuffle=(train_sampler is None),    sampler=train_sampler,)test_loader = torch.utils.data.DataLoader(test_dataset, ...)

最后，用戶只需要選擇要訓(xùn)練的模型和優(yōu)化器即可以使用 bagua：

# 定義模型model = ...model = model.cuda()# 定義優(yōu)化器optimizer = ...# 選擇 Bagua 算法來使用from bagua.torch_api.algorithms import gradient_allreduce# 實例化 Bagua 算法algorithm = gradient_allreduce.GradientAllReduceAlgorithm()# 對現(xiàn)有模型啟用 Bagua 算法model = model.with_bagua(    [optimizer], algorithm)

這樣，使用 Bagua 的多機多卡訓(xùn)練算法就實現(xiàn)完成了。完整例子和更多場景，歡迎參考 Bagua Tutorial 文檔（https://github.com/BaguaSys/bagua）。

論文: https://arxiv.org/abs/2107.01499

項目 GitHub 地址：https://github.com/BaguaSys/bagua

參考文獻

[1] Dean, Jeffrey, Greg S. Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Quoc V. Le, Mark Z. Mao et al. “Large scale distributed deep networks.” (2012).

[2] Zhengyuan Zhou, Panayotis Mertikopoulos, Nicholas Bambos, Peter Glynn, Yinyu Ye, Li-Jia Li, and Li Fei-Fei. 2018. Distributed asynchronous optimization with unbounded delays: How slow can you go?. In International Conference on Machine Learning. PMLR, 5970–5979.

[3] DanAlistarh, DemjanGrubic, JerryLi, RyotaTomioka, and MilanVojnovic. 2016. QSGD: Communication-efficient SGD via gradient quantization and encoding. arXiv preprint arXiv:1610.02132 (2016).

[4] Dan Alistarh, Torsten Hoefler, Mikael Johansson, Sarit Khirirat, Nikola Konstanti- nov, and Cédric Renggli. 2018. The convergence of sparsified gradient methods. In Proceedings of the 32nd International Conference on Neural Information Processing Systems. 5977–5987.

[5] Anastasia Koloskova, Sebastian Stich, and Martin Jaggi. 2019. Decentralized stochastic optimization and gossip algorithms with compressed communication. In International Conference on Machine Learning. PMLR, 3478–3487.

[6] Xiangru Lian, Ce Zhang, Huan Zhang, Cho-Jui Hsieh, Wei Zhang, and Ji Liu. 2017. Can decentralized algorithms outperform centralized algorithms? a case study for decentralized parallel stochastic gradient descent. In Proceedings of the 31st International Conference on Neural Information Processing Systems. 5336–5346.

[7] Christopher De Sa, Matthew Feldman, Christopher Ré, and Kunle Olukotun. 2017. Understanding and optimizing asynchronous low-precision stochastic gradient descent. In Proceedings of the 44th Annual International Symposium on Computer Architecture. 561–574.

[8] Xiangru Lian, Wei Zhang, Ce Zhang, and Ji Liu. 2018. Asynchronous decentral- ized parallel stochastic gradient descent. In International Conference on Machine Learning. PMLR, 3043–3052.

[9] Hanlin Tang, Shaoduo Gan, Ce Zhang, Tong Zhang, and Ji Liu. 2018. Com- munication compression for decentralized training. In Proceedings of the 32nd International Conference on Neural Information Processing Systems. 7663–7673.

[10] Ji Liu, Ce Zhang, et al. 2020. Distributed Learning Systems with First-Order Methods. Foundations and Trends? in Databases 9, 1 (2020), 1–100.