在 AI 浪潮中，無論是企業還是國家，對算力的需求都日益高漲。近期啟動的“東數西算”項目，更是從宏觀層面大力打造 AI 基礎設施。但位于不同地理位置的計算機之間通信延遲較高，如何統籌兼顧、高效利用不同地區的計算能力，是當下亟待解決的重大議題。

“東數西算”項目布局

與此同時，在大數據時代的背景下，如何保護隱私數據也成為社會熱點，國家出臺了數據安全、隱私保護的一系列法規。

針對以上難點，開源項目 Sky Computing 成功利用空間異構分布式計算特性，在保證用戶數據隱私的前提下，可對聯邦學習加速達 55%。

地址：

https://github.com/hpcaitech/SkyComputing

空間異構分布式計算

隨著深度學習的不斷發展，模型的尺寸日益增長，目前的主流模型，例如 BERT 和 GPT-3 都有著數以億計的參數。盡管這些模型在預測精度和性能提升方面有了長足的進步，但同樣也給存儲和運算等帶來了極大的壓力。為了加速AI模型訓練的速度，分布式機器學習得以應運而生，它通常使用大量高速互聯的同類型處理器，如超級計算機。

超級計算機

而空間異構分布式計算則進一步將擁有不同計算能力、通訊能力的計算資源組合在一起，作為一個大的集群完成大型計算任務。其中參與計算的硬件資源可以是大型專業計算服務器，也可以是小型的智能設備。目前，空間異構分布式計算作為一種新形式的異構計算，正在得到越來越多的關注。以我國為例，隨著「東數西算」工作的推行，越來越多的計算資源將廣泛地分布到西部各個地區，如何協調這類混合計算集群聯合高效工作，也將成為高性能計算應用的研究熱點。

近年來，云服務的規模、范圍和對象都被不斷擴展，越來越多的企業選擇將自己的數據存儲和數據計算相關業務部署在云端。然而，將所有服務依托于云端環境的缺點在于數據的遷移成本極高；同時，數據的隱私性和可靠性也難以保證；此外，分布在不同地區的云算力之間高昂的通信成本，也使得他們難以有效聯合完成高算力任務。

云計算

聯邦學習

為保護數據的隱私性，Google 于 2016 年提出聯邦學習，這是一種加密的分布式機器學習技術。顧名思義，它通過搭建一個虛擬的「聯邦」，將大大小小的數據孤島聯合到一起。每一個數據孤島都像是這個「聯邦」中的一個州，既保持一定的獨立自主（比如商業機密，用戶隱私），又能在數據不被對外共享的前提下共同建模，提升 AI 模型效果。目前，聯邦學習廣泛被運用在智能終端的模型訓練中，如各個語音助手例如 Siri、Alex 等等。

聯邦學習

在現有的聯邦學習模型并行中，模型被均勻分配給各個訓練設備。然而，如前文所述，由于聯邦學習的訓練設備往往是用戶的智能終端，性能差異較大，使用均勻分配，往往會造成通信時間瓶頸。

正如我們都知道木桶效應：木桶的盛水量由最短的那塊木板決定。而在傳統的聯邦學習中，存在類似現象：訓練速度由最慢的那個設備決定。

例如，對于處于使用模型并行的同一個聯邦學習任務中的智能手機和樹莓派，它們會被分配相同的任務量。但由于智能手機的運算能力遠超樹莓派，智能手機被迫閑置等待樹莓派的任務完成。

木桶效應

Sky Computing

Sky Computing 針對以上痛點，通過負載均衡，將不同規模和能力的云服務器智能互聯，達到大規模計算的算力需求，同時通過聯邦學習的方式，僅在云服務器內部訪問用戶數據，避免數據遷移和隱私泄露。

負載均衡

要解決負載均衡的問題，首先要了解什么是「負載」。在計算機中，無論進行哪種操作，究其本質，負載都可以理解為「完成任務所需的時間」。由于在聯邦學習中，訓練模型的計算總量是固定的，因此如果我們能通過自適應的方式智能分配計算任務，便能夠使得每個設備完成計算任務的耗時相同，確保整體訓練的時間最優。而為了得到一個好的分配方式，我們需要首先得到模型和設備相關信息，然后再進行實際的適當分配操作。因此，對于訓練模型，我們需要分為兩個階段：基準測試和分配。

訓練過程

基準測試

在基準測試階段，Sky Computing 需要收集來自兩個維度的數據：模型和設備。在模型維度，需要知道模型每一層所需的內存占用和計算量。通過結合模型的預計內存占用和設備的可用內存，可避免內存溢出；而所需計算量越大，同一設備完成該任務的時間就越久。在設備維度，需要知道設備的通訊延時、計算能力和可用內存等，受網絡環境、當前運行負載等因素的影響。對于算力強、通信好但可用內存少的設備，應在內存不溢出的前提下，盡量多分配模型層（計算任務）。由于 Sky Computing 是一個負載均衡的聯邦學習系統，因此我們在基準測試階段只關心設備的機器學習的能力。通過在每個設備運行小型的機器學習測試任務，測探設備的 AI 計算能力。

整體流程

分配

在決定任務分配方式時，經數學分析可知，分配方式本質上是一個 NP-hard 的混合整數線性規劃問題。因此，在多項式時間內，我們無法得到一個最優解。而隨著模型規模的不斷增長，和設備數量的不斷增多，計算最優解的成本顯然是不可接受的。

因此，在實際情況中，我們不會直接計算求得最優解，而是嘗試使用啟發式算法得到近似解。在 Sky Computing 中，我們設計了一個兩階段的啟發式算法：第一階段為預分配，按照設備的實際可用內存大小進行模型的分配，并且計算每個設備實際的工作負載；第二階段為分配調整，根據設備的負載量進行動態的調整，迭代降低整個系統的負載量。同時，為了驗證 Sky Computing 的優越性，我們在實驗中也設置了最優分配作為對比。

實現架構

性能表現

我們在集群環境中，采用控制關鍵因素變量的方式，以聯邦學習 AI 任務的 forward 和 backward 的時間為指標，對 Sky Computing 的性能進行了驗證。

實驗結果

我們測試了三種分配方式（even：均勻分配，heuristic：啟發式算法，optimal：最優分配）。在不同的計算資源數量規模和不同的模型大小下的表現，并記錄了每次完成迭代所花費的時間。可以看到，隨著設備數量的增多和模型深度的增加，我們的啟發式算法的效果十分顯著。在 64 個節點 160 層隱藏層的實驗環境下，Sky Computing 比當前的均勻分配模型并行可加速 55%。

實驗結果

其中，由于最優分配計算成本極高，在 64 節點時已難以計算，不適用于實際應用，僅作為小規模時的參考值。

開源共建

Sky Computing 是我們利用空間異構分布式計算特性加速聯邦學習的一次成功嘗試，獲得了高達 55% 的性能提升。目前該項目仍處于開發階段，未來我們將進行更加充分的實驗，早日部署到實際應用中，并提供動態冗余等功能。

論文地址：https://arxiv.org/abs/2202.11836

項目地址：https://github.com/hpcaitech/SkyComputing