在 AI 浪潮中，無(wú)論是企業(yè)還是國(guó)家，對(duì)算力的需求都日益高漲。近期啟動(dòng)的“東數(shù)西算”項(xiàng)目，更是從宏觀層面大力打造 AI 基礎(chǔ)設(shè)施。但位于不同地理位置的計(jì)算機(jī)之間通信延遲較高，如何統(tǒng)籌兼顧、高效利用不同地區(qū)的計(jì)算能力，是當(dāng)下亟待解決的重大議題。

“東數(shù)西算”項(xiàng)目布局

與此同時(shí)，在大數(shù)據(jù)時(shí)代的背景下，如何保護(hù)隱私數(shù)據(jù)也成為社會(huì)熱點(diǎn)，國(guó)家出臺(tái)了數(shù)據(jù)安全、隱私保護(hù)的一系列法規(guī)。

針對(duì)以上難點(diǎn)，開(kāi)源項(xiàng)目 Sky Computing 成功利用空間異構(gòu)分布式計(jì)算特性，在保證用戶數(shù)據(jù)隱私的前提下，可對(duì)聯(lián)邦學(xué)習(xí)加速達(dá) 55%。

地址：

https://github.com/hpcaitech/SkyComputing

空間異構(gòu)分布式計(jì)算

隨著深度學(xué)習(xí)的不斷發(fā)展，模型的尺寸日益增長(zhǎng)，目前的主流模型，例如 BERT 和 GPT-3 都有著數(shù)以億計(jì)的參數(shù)。盡管這些模型在預(yù)測(cè)精度和性能提升方面有了長(zhǎng)足的進(jìn)步，但同樣也給存儲(chǔ)和運(yùn)算等帶來(lái)了極大的壓力。為了加速AI模型訓(xùn)練的速度，分布式機(jī)器學(xué)習(xí)得以應(yīng)運(yùn)而生，它通常使用大量高速互聯(lián)的同類型處理器，如超級(jí)計(jì)算機(jī)。

超級(jí)計(jì)算機(jī)

而空間異構(gòu)分布式計(jì)算則進(jìn)一步將擁有不同計(jì)算能力、通訊能力的計(jì)算資源組合在一起，作為一個(gè)大的集群完成大型計(jì)算任務(wù)。其中參與計(jì)算的硬件資源可以是大型專業(yè)計(jì)算服務(wù)器，也可以是小型的智能設(shè)備。目前，空間異構(gòu)分布式計(jì)算作為一種新形式的異構(gòu)計(jì)算，正在得到越來(lái)越多的關(guān)注。以我國(guó)為例，隨著「東數(shù)西算」工作的推行，越來(lái)越多的計(jì)算資源將廣泛地分布到西部各個(gè)地區(qū)，如何協(xié)調(diào)這類混合計(jì)算集群聯(lián)合高效工作，也將成為高性能計(jì)算應(yīng)用的研究熱點(diǎn)。

近年來(lái)，云服務(wù)的規(guī)模、范圍和對(duì)象都被不斷擴(kuò)展，越來(lái)越多的企業(yè)選擇將自己的數(shù)據(jù)存儲(chǔ)和數(shù)據(jù)計(jì)算相關(guān)業(yè)務(wù)部署在云端。然而，將所有服務(wù)依托于云端環(huán)境的缺點(diǎn)在于數(shù)據(jù)的遷移成本極高；同時(shí)，數(shù)據(jù)的隱私性和可靠性也難以保證；此外，分布在不同地區(qū)的云算力之間高昂的通信成本，也使得他們難以有效聯(lián)合完成高算力任務(wù)。

云計(jì)算

聯(lián)邦學(xué)習(xí)

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

聯(lián)邦學(xué)習(xí)

在現(xiàn)有的聯(lián)邦學(xué)習(xí)模型并行中，模型被均勻分配給各個(gè)訓(xùn)練設(shè)備。然而，如前文所述，由于聯(lián)邦學(xué)習(xí)的訓(xùn)練設(shè)備往往是用戶的智能終端，性能差異較大，使用均勻分配，往往會(huì)造成通信時(shí)間瓶頸。

正如我們都知道木桶效應(yīng)：木桶的盛水量由最短的那塊木板決定。而在傳統(tǒng)的聯(lián)邦學(xué)習(xí)中，存在類似現(xiàn)象：訓(xùn)練速度由最慢的那個(gè)設(shè)備決定。

例如，對(duì)于處于使用模型并行的同一個(gè)聯(lián)邦學(xué)習(xí)任務(wù)中的智能手機(jī)和樹(shù)莓派，它們會(huì)被分配相同的任務(wù)量。但由于智能手機(jī)的運(yùn)算能力遠(yuǎn)超樹(shù)莓派，智能手機(jī)被迫閑置等待樹(shù)莓派的任務(wù)完成。

木桶效應(yīng)

Sky Computing

Sky Computing 針對(duì)以上痛點(diǎn)，通過(guò)負(fù)載均衡，將不同規(guī)模和能力的云服務(wù)器智能互聯(lián)，達(dá)到大規(guī)模計(jì)算的算力需求，同時(shí)通過(guò)聯(lián)邦學(xué)習(xí)的方式，僅在云服務(wù)器內(nèi)部訪問(wèn)用戶數(shù)據(jù)，避免數(shù)據(jù)遷移和隱私泄露。

負(fù)載均衡

要解決負(fù)載均衡的問(wèn)題，首先要了解什么是「負(fù)載」。在計(jì)算機(jī)中，無(wú)論進(jìn)行哪種操作，究其本質(zhì)，負(fù)載都可以理解為「完成任務(wù)所需的時(shí)間」。由于在聯(lián)邦學(xué)習(xí)中，訓(xùn)練模型的計(jì)算總量是固定的，因此如果我們能通過(guò)自適應(yīng)的方式智能分配計(jì)算任務(wù)，便能夠使得每個(gè)設(shè)備完成計(jì)算任務(wù)的耗時(shí)相同，確保整體訓(xùn)練的時(shí)間最優(yōu)。而為了得到一個(gè)好的分配方式，我們需要首先得到模型和設(shè)備相關(guān)信息，然后再進(jìn)行實(shí)際的適當(dāng)分配操作。因此，對(duì)于訓(xùn)練模型，我們需要分為兩個(gè)階段：基準(zhǔn)測(cè)試和分配。

訓(xùn)練過(guò)程

基準(zhǔn)測(cè)試

在基準(zhǔn)測(cè)試階段，Sky Computing 需要收集來(lái)自兩個(gè)維度的數(shù)據(jù)：模型和設(shè)備。在模型維度，需要知道模型每一層所需的內(nèi)存占用和計(jì)算量。通過(guò)結(jié)合模型的預(yù)計(jì)內(nèi)存占用和設(shè)備的可用內(nèi)存，可避免內(nèi)存溢出；而所需計(jì)算量越大，同一設(shè)備完成該任務(wù)的時(shí)間就越久。在設(shè)備維度，需要知道設(shè)備的通訊延時(shí)、計(jì)算能力和可用內(nèi)存等，受網(wǎng)絡(luò)環(huán)境、當(dāng)前運(yùn)行負(fù)載等因素的影響。對(duì)于算力強(qiáng)、通信好但可用內(nèi)存少的設(shè)備，應(yīng)在內(nèi)存不溢出的前提下，盡量多分配模型層（計(jì)算任務(wù)）。由于 Sky Computing 是一個(gè)負(fù)載均衡的聯(lián)邦學(xué)習(xí)系統(tǒng)，因此我們?cè)诨鶞?zhǔn)測(cè)試階段只關(guān)心設(shè)備的機(jī)器學(xué)習(xí)的能力。通過(guò)在每個(gè)設(shè)備運(yùn)行小型的機(jī)器學(xué)習(xí)測(cè)試任務(wù)，測(cè)探設(shè)備的 AI 計(jì)算能力。

整體流程

分配

在決定任務(wù)分配方式時(shí)，經(jīng)數(shù)學(xué)分析可知，分配方式本質(zhì)上是一個(gè) NP-hard 的混合整數(shù)線性規(guī)劃問(wèn)題。因此，在多項(xiàng)式時(shí)間內(nèi)，我們無(wú)法得到一個(gè)最優(yōu)解。而隨著模型規(guī)模的不斷增長(zhǎng)，和設(shè)備數(shù)量的不斷增多，計(jì)算最優(yōu)解的成本顯然是不可接受的。

因此，在實(shí)際情況中，我們不會(huì)直接計(jì)算求得最優(yōu)解，而是嘗試使用啟發(fā)式算法得到近似解。在 Sky Computing 中，我們?cè)O(shè)計(jì)了一個(gè)兩階段的啟發(fā)式算法：第一階段為預(yù)分配，按照設(shè)備的實(shí)際可用內(nèi)存大小進(jìn)行模型的分配，并且計(jì)算每個(gè)設(shè)備實(shí)際的工作負(fù)載；第二階段為分配調(diào)整，根據(jù)設(shè)備的負(fù)載量進(jìn)行動(dòng)態(tài)的調(diào)整，迭代降低整個(gè)系統(tǒng)的負(fù)載量。同時(shí)，為了驗(yàn)證 Sky Computing 的優(yōu)越性，我們?cè)趯?shí)驗(yàn)中也設(shè)置了最優(yōu)分配作為對(duì)比。

實(shí)現(xiàn)架構(gòu)

性能表現(xiàn)

我們?cè)诩涵h(huán)境中，采用控制關(guān)鍵因素變量的方式，以聯(lián)邦學(xué)習(xí) AI 任務(wù)的 forward 和 backward 的時(shí)間為指標(biāo)，對(duì) Sky Computing 的性能進(jìn)行了驗(yàn)證。

實(shí)驗(yàn)結(jié)果

我們測(cè)試了三種分配方式（even：均勻分配，heuristic：?jiǎn)l(fā)式算法，optimal：最優(yōu)分配）。在不同的計(jì)算資源數(shù)量規(guī)模和不同的模型大小下的表現(xiàn)，并記錄了每次完成迭代所花費(fèi)的時(shí)間。可以看到，隨著設(shè)備數(shù)量的增多和模型深度的增加，我們的啟發(fā)式算法的效果十分顯著。在 64 個(gè)節(jié)點(diǎn) 160 層隱藏層的實(shí)驗(yàn)環(huán)境下，Sky Computing 比當(dāng)前的均勻分配模型并行可加速 55%。

實(shí)驗(yàn)結(jié)果

其中，由于最優(yōu)分配計(jì)算成本極高，在 64 節(jié)點(diǎn)時(shí)已難以計(jì)算，不適用于實(shí)際應(yīng)用，僅作為小規(guī)模時(shí)的參考值。

開(kāi)源共建

Sky Computing 是我們利用空間異構(gòu)分布式計(jì)算特性加速聯(lián)邦學(xué)習(xí)的一次成功嘗試，獲得了高達(dá) 55% 的性能提升。目前該項(xiàng)目仍處于開(kāi)發(fā)階段，未來(lái)我們將進(jìn)行更加充分的實(shí)驗(yàn)，早日部署到實(shí)際應(yīng)用中，并提供動(dòng)態(tài)冗余等功能。

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

項(xiàng)目地址：https://github.com/hpcaitech/SkyComputing