李宗航/LI Zonghang, 虞紅芳/YU Hongfang,3, 汪漪/WANG Yi

（1. 電子科技大學(xué)，中國 成都，611731;2. 南方科技大學(xué)，中國 深圳，518055;3. 鵬城實驗室，中國 深圳，518055）

(1. University of Electronic Science and Technology of China, Chengdu 611731, China;2. Southern University of Science and Technology, Shenzhen 518055, China;3. Peng Cheng Laboratory, Shenzhen 518055, China)

隨著5G技術(shù)的逐步商用，激增的數(shù)據(jù)為人工智能技術(shù)提供了豐富的數(shù)據(jù)資源，催生了人臉支付、輔助感知等新時代智能應(yīng)用的萌芽。為強化智能應(yīng)用的生命力，加速以深度學(xué)習(xí)為代表的智能模型的訓(xùn)練過程，分布式機器學(xué)習(xí)[1]成為業(yè)界重點關(guān)注的關(guān)鍵技術(shù)。數(shù)據(jù)并行的分布式機器學(xué)習(xí)將完整數(shù)據(jù)均勻分配給數(shù)據(jù)中心內(nèi)的多個處理節(jié)點。這些節(jié)點各自維護一個本地模型副本以并行計算其數(shù)據(jù)切片上的本地梯度，并遞交給其他節(jié)點計算全局梯度，隨后全局梯度在所有節(jié)點間同步，最終被用于更新各節(jié)點的本地模型。

然而，訓(xùn)練所需的海量數(shù)據(jù)分散在多個機構(gòu)的數(shù)據(jù)中心內(nèi)，將這些數(shù)據(jù)碎片匯集于一個數(shù)據(jù)中心將引入巨額通信代價。業(yè)界對數(shù)據(jù)隱私和知識產(chǎn)權(quán)的愈發(fā)重視導(dǎo)致碎片之間更加割裂，使得人工智能技術(shù)從數(shù)據(jù)碎片中挖掘的知識具有明顯的片面性。因此，要發(fā)揮大數(shù)據(jù)的集聚和增值作用，提供高質(zhì)量的智慧化服務(wù)，就需要以數(shù)據(jù)隱私保護為前提，在多機構(gòu)間進行高效的分布式訓(xùn)練，打破機構(gòu)間的數(shù)據(jù)和網(wǎng)絡(luò)壁壘，化解數(shù)據(jù)孤島難題。

地理分布式機器學(xué)習(xí)[2]（GeoML）是一種面向跨地域分布的數(shù)據(jù)中心的分布式訓(xùn)練范式。與聯(lián)邦學(xué)習(xí)[3]類似，GeoML將訓(xùn)練樣本保留在中心內(nèi)，并在中心間交互模型數(shù)據(jù)，從而實現(xiàn)數(shù)據(jù)的隱私保護。需要指出的是，聯(lián)邦學(xué)習(xí)將機構(gòu)整體視為一個參與節(jié)點，而GeoML允許機構(gòu)的中心內(nèi)有多個參與節(jié)點，中心內(nèi)通過大帶寬局域網(wǎng)互聯(lián)，中心間通過小帶寬廣域網(wǎng)互聯(lián)，整體形成多中心跨域互聯(lián)形態(tài)。鑒于數(shù)據(jù)在中心內(nèi)的分布式存儲現(xiàn)狀以及業(yè)界對并行加速模型訓(xùn)練的需求，GeoML更適用于多中心跨域互聯(lián)。

1 地理分布式機器學(xué)習(xí)機遇與挑戰(zhàn)

近年來，人工智能的全球發(fā)展熱潮勢不可擋，并在智慧醫(yī)療、智慧城市、互聯(lián)網(wǎng)金融等關(guān)鍵領(lǐng)域創(chuàng)造了巨大的社會效益。GeoML作為一種多中心跨域互聯(lián)的分布式機器學(xué)習(xí)技術(shù)，為上述關(guān)鍵領(lǐng)域中大型數(shù)據(jù)孤島的融合提供了安全且高效的解決方案，有力推動了人工智能數(shù)據(jù)能源的整合，因而正愈加受到業(yè)界的關(guān)注和重視。此外，聯(lián)邦學(xué)習(xí)和隱私計算等新興領(lǐng)域的安全成果（安全多方計算[4]、差分隱私[5]、同態(tài)加密[6]等）進一步鞏固了GeoML的隱私保護能力，在較大程度上打消了合作機構(gòu)對數(shù)據(jù)安全的顧慮，從而更利于GeoML的發(fā)展與落地。可見，GeoML正處于具有強烈需求且前置技術(shù)完備的時代，擁有無限的機遇和廣闊的應(yīng)用前景。

然而，GeoML同時面臨著基礎(chǔ)框架缺乏與通信開銷龐大的挑戰(zhàn)。一方面，業(yè)界尚缺乏適用于GeoML的通用開源軟件框架。雖然部分開源深度學(xué)習(xí)框架（如MXNET）和聯(lián)邦學(xué)習(xí)框架（如FATE）也能部署于多中心跨域互聯(lián)場景，但這些框架無法區(qū)別和適配中心內(nèi)/外差異化的集群環(huán)境，從而會不可避免地帶來高額通信開銷和資源浪費。另一方面，GeoML擁有比聯(lián)邦學(xué)習(xí)更大的集群規(guī)模和更激烈的競爭環(huán)境，且需要在中心之間反復(fù)交換大量模型數(shù)據(jù)，使其面臨更顯著的通信開銷。

針對上述挑戰(zhàn)，我們在第2章提出一種GeoML通用軟件框架，在第3、4章分別從通信架構(gòu)和壓縮傳輸機制著手優(yōu)化通信，提出一種分層參數(shù)服務(wù)器（HiPS）架構(gòu)減少廣域傳輸梯度流數(shù)量，以及一種雙向稀疏梯度傳輸（BiSparse）技術(shù)減少廣域傳輸梯度流大小。

2 地理分布式機器學(xué)習(xí)軟件框架

本文中，我們基于經(jīng)典深度學(xué)習(xí)框架MXNET[7]開發(fā)了地理分布式機器學(xué)習(xí)框架GeoMX。該框架系統(tǒng)架構(gòu)如圖1所示，包括以下部分：

（1）數(shù)據(jù)源。數(shù)據(jù)源是指機器學(xué)習(xí)訓(xùn)練所需的視頻、音頻、圖像、文本等數(shù)據(jù)的存儲介質(zhì)。根據(jù)數(shù)據(jù)實際存放位置的不同，處理節(jié)點可以從本地文件系統(tǒng)、所在數(shù)據(jù)中心的分布式文件系統(tǒng)（HDFS）以及遠端S3存儲中讀取數(shù)據(jù)。

（2）數(shù)據(jù)處理模塊。該模塊負(fù)責(zé)加載與解析數(shù)據(jù)源中的數(shù)據(jù)，根據(jù)訓(xùn)練配置進行預(yù)處理、切片和采樣，并將數(shù)據(jù)緩存到數(shù)據(jù)隊列以供模型計算使用。數(shù)據(jù)預(yù)處理包括數(shù)據(jù)清洗、歸一化、標(biāo)準(zhǔn)化、結(jié)構(gòu)變換以及隨機翻轉(zhuǎn)、裁剪、色彩變換等數(shù)據(jù)增廣操作，以保證數(shù)據(jù)符合格式規(guī)范，并且擴充數(shù)據(jù)以緩解模型過擬合。數(shù)據(jù)切片是指將存儲于共享數(shù)據(jù)源的完整數(shù)據(jù)集進行切分，其中切片由多個處理節(jié)點分別維護。數(shù)據(jù)采樣是指處理節(jié)點從數(shù)據(jù)切片中采集小批量樣本用于計算梯度，常用采樣方式包括隨機采樣、順序采樣、重要性采樣等。

▲圖1 GeoMX（左）與HiPS（右）系統(tǒng)架構(gòu)圖

（3）存儲與日志模塊。存儲模塊包括內(nèi)存管理模塊和顯存管理模塊，以分別負(fù)責(zé)內(nèi)存和顯存上存儲空間的分配、釋放以及垃圾回收；日志模塊則為模型訓(xùn)練等操作提供便捷美觀的進度跟蹤、關(guān)鍵指標(biāo)記錄、系統(tǒng)運行狀態(tài)記錄等功能。

（4）模型管理模塊。該模塊負(fù)責(zé)模型計算圖的構(gòu)建以及模型參數(shù)的管理。模型計算圖中的張量節(jié)點由模型構(gòu)建模塊產(chǎn)生，節(jié)點間的計算關(guān)系由模型計算、自動梯度計算和模型優(yōu)化器模塊產(chǎn)生。模型計算圖中張量的加載、保存、格式轉(zhuǎn)換等功能則由其他模塊提供。

（5）計算與通信模塊。模型計算圖的計算由依賴圖計算的引擎執(zhí)行，該引擎可分析張量間的依賴關(guān)系和調(diào)度多個張量計算并行執(zhí)行。多個處理節(jié)點間模型計算圖的規(guī)約和同步由HiPS通信模塊提供，該通信模塊提供標(biāo)準(zhǔn)的PUSH和PULL原語，利用層次化通信架構(gòu)和壓縮傳輸機制實現(xiàn)高效通信。

（6）業(yè)務(wù)場景。GeoMX可部署于多種廣域分布式機器學(xué)習(xí)場景。例如，云邊端協(xié)同的移動邊緣計算場景，微云、微數(shù)據(jù)中心和公有云組成的邊緣計算場景，多個公有云跨域互聯(lián)的場景，以及面向機構(gòu)的聯(lián)邦學(xué)習(xí)場景。

（7）業(yè)務(wù)應(yīng)用。GeoMX不僅支持所有兼容梯度優(yōu)化的機器學(xué)習(xí)模型，如神經(jīng)網(wǎng)絡(luò)、梯度支持向量機（GSVM）、梯度提升決策樹（GBDT）等，還可用于支撐大多數(shù)復(fù)雜的智能應(yīng)用，如圖像、語音、視頻、文本等的識別、檢測、預(yù)測和生成任務(wù)。

在HiPS通信模塊中，本文中我們主要關(guān)注以下技術(shù)：

（1）層次化通信架構(gòu)。該架構(gòu)將中心內(nèi)/外網(wǎng)絡(luò)環(huán)境分隔為兩層，第1層在中心內(nèi)實現(xiàn)局部聚合，第2層在中心間實現(xiàn)全局聚合。這樣可有效減少中心間傳輸?shù)臄?shù)據(jù)量，也可降低中心的管理成本和安全風(fēng)險。

（2）混合同步技術(shù)。該技術(shù)包含全同步模式和混合同步模式：全同步模式即中心內(nèi)和中心間都采用同步并行，混合同步模式即中心內(nèi)采用同步并行而中心間采用異步并行。針對這兩種模式，可根據(jù)集群異構(gòu)程度、對模型質(zhì)量與訓(xùn)練效率的需求進行選擇。

（3）壓縮傳輸技術(shù)。該技術(shù)旨在壓縮傳輸模型以減少中心間實際傳輸?shù)臄?shù)據(jù)量，從而降低通信開銷。GeoMX支持的壓縮方法包括半精度量化[8]、2bit量化[9]、深度梯度壓縮（DGC）[10]、DGT[11]，以及第4章提出的BiSparse技術(shù)。

3 HiPS架構(gòu)及流量模型

在通信架構(gòu)優(yōu)化方面，我們基于參數(shù)服務(wù)器架構(gòu)[12]提出HiPS架構(gòu)。以全同步模式為例，假設(shè)有圖2所示的一個中心機構(gòu)和兩個參與機構(gòu)，各機構(gòu)中心內(nèi)的物理節(jié)點通過局域網(wǎng)互聯(lián)，中心間通過廣域網(wǎng)互聯(lián)。

中心機構(gòu)包括全局參數(shù)服務(wù)器、主控工作節(jié)點、本地調(diào)度器和全局調(diào)度器4個節(jié)點，其中前3個節(jié)點彼此互聯(lián)，全局參數(shù)服務(wù)器與全局調(diào)度器互聯(lián)。參與機構(gòu)包括參數(shù)服務(wù)器、本地調(diào)度器和若干工作節(jié)點。參與機構(gòu)內(nèi)的所有節(jié)點彼此互聯(lián)（工作節(jié)點間除外），參數(shù)服務(wù)器與全局調(diào)度器互聯(lián)。參與機構(gòu)內(nèi)部節(jié)點構(gòu)成中心內(nèi)參數(shù)服務(wù)器架構(gòu)，參與機構(gòu)的參數(shù)服務(wù)器與中心機構(gòu)的全局參數(shù)服務(wù)器和全局調(diào)度器構(gòu)成中心間參數(shù)服務(wù)器架構(gòu)。

HiPS流量模型如圖3所示（調(diào)度器因不參與模型數(shù)據(jù)交互而被簡化），其中各類節(jié)點功能如下：

▲圖2 分層參數(shù)服務(wù)器架構(gòu)

▲圖3 分層參數(shù)服務(wù)器流量模型

（1）工作節(jié)點。該類節(jié)點位于參與機構(gòu)內(nèi)，負(fù)責(zé)加載、解析、清洗和預(yù)處理數(shù)據(jù)源中的數(shù)據(jù)并用于計算本地模型梯度。這些梯度由機構(gòu)內(nèi)的參數(shù)服務(wù)器聚合，并被用于更新本地模型。

（2）主控工作節(jié)點。該節(jié)點是位于中心機構(gòu)的一個特殊工作節(jié)點，負(fù)責(zé)構(gòu)建模型計算圖，初始化圖中的張量參數(shù)，并交付給全局參數(shù)服務(wù)器以初始化全局模型。此外，主控工作節(jié)點還負(fù)責(zé)集群的配置，包括設(shè)置同步模式、壓縮策略和優(yōu)化算法等。

（3）參數(shù)服務(wù)器。該類節(jié)點位于參與機構(gòu)內(nèi)，首先負(fù)責(zé)聚合機構(gòu)內(nèi)工作節(jié)點提交的本地模型梯度，并轉(zhuǎn)發(fā)到中心機構(gòu)的全局參數(shù)服務(wù)器。隨后，該類節(jié)點從全局參數(shù)服務(wù)器獲取最新模型數(shù)據(jù)，并同步到機構(gòu)內(nèi)的所有工作節(jié)點。

（4）全局參數(shù)服務(wù)器。該類節(jié)點位于中心機構(gòu)內(nèi)，負(fù)責(zé)對所有參與機構(gòu)提交的局部聚合梯度執(zhí)行全局聚合。本文中，我們默認(rèn)全局參數(shù)服務(wù)器直接返回全局聚合梯度，但若主控工作節(jié)點設(shè)置了全局參數(shù)服務(wù)器的優(yōu)化算法，全局參數(shù)服務(wù)器將用全局聚合梯度更新全局模型，并下發(fā)最新模型參數(shù)到參與機構(gòu)。

（5）本地調(diào)度器。該節(jié)點位于所有機構(gòu)內(nèi)，用于機構(gòu)內(nèi)節(jié)點的注冊、中止以及狀態(tài)管理（節(jié)點地址、集群規(guī)模、心跳等），也負(fù)責(zé)管理機構(gòu)內(nèi)節(jié)點的動態(tài)加入和退出以及異常退出檢測。

（6）全局調(diào)度器。該節(jié)點位于中心機構(gòu)內(nèi)，其作用與本地調(diào)度器類似，負(fù)責(zé)參與機構(gòu)的動態(tài)加入和退出以及異常退出檢測。

我們使用7臺服務(wù)器構(gòu)建計算集群來驗證上述技術(shù)，其中每臺服務(wù)器配有兩張Tesla K40M計算卡，服務(wù)器之間使用千兆以太網(wǎng)互聯(lián)。我們按1∶3∶3的比例將7臺服務(wù)器分為中心機構(gòu)A、參與機構(gòu)B和參與機構(gòu)C，并依照圖2部署各個機構(gòu)，其中B和C各部署6個工作節(jié)點。我們選擇MXNET框架作為對照，其中中心機構(gòu)部署一個參數(shù)服務(wù)器和一個調(diào)度器，兩個參與機構(gòu)各部署6個工作節(jié)點，并且所有工作節(jié)點與參數(shù)服務(wù)器直接通信。上述節(jié)點均部署于Docker容器中。本文中，我們將兩個框架部署于1 Gbit/s帶寬網(wǎng)絡(luò)來模擬數(shù)據(jù)中心內(nèi)的局域網(wǎng)環(huán)境，并使用Wondershaper工具將中心機構(gòu)的上下行網(wǎng)卡帶寬限制為155 Mbit/s以模擬數(shù)據(jù)中心間的廣域網(wǎng)環(huán)境。

我們將圖像分類數(shù)據(jù)集Fashion-MNIST[13]均勻切分到所有工作節(jié)點來訓(xùn)練ResNet-50[14]模型，并配置集群以全同步模式、無壓縮模式運行，工作節(jié)點使用標(biāo)準(zhǔn)隨機梯度下降（SGD）優(yōu)化器更新模型參數(shù)。GeoMX與MXNET在局域和廣域環(huán)境下的測試精度曲線如圖4所示。在相同帶寬條件下，GeoMX明顯提高了MXNET的訓(xùn)練效率并且不損傷收斂精度。此外，多中心跨域互聯(lián)場景下的GeoMX-155 Mbit/s訓(xùn)練效率略優(yōu)于數(shù)據(jù)中心內(nèi)的MXNET-1 Gbit/s。

為消除模型大小對上述結(jié)論的影響，我們用參數(shù)量遞增的ResNet系列模型（ResNet-18/34/50/101）重復(fù)上述實驗，并在表1中給出各組實驗的單輪通信時延。實驗數(shù)據(jù)再次展示了GeoMX趨同局域的通信效率：在所有模型上，GeoMX-155 Mbit/s的通信時延僅為MXNET-1Gbit/s的78%～87% 。

4 BiSparse技術(shù)

為進一步降低廣域通信開銷，在壓縮傳輸方面，我們基于DGC[10]提出BiSparse技術(shù)，通過稀疏傳輸上下行梯度，可以有效減少中心間傳輸?shù)奶荻攘俊閷崿F(xiàn)通信高效且精度無損的稀疏傳輸，BiSparse不僅需要快速且準(zhǔn)確地從大規(guī)模梯度張量中篩選關(guān)鍵梯度，并在稀疏空間中實現(xiàn)高效的傳輸和聚合，還需要解決因丟失非關(guān)鍵梯度而導(dǎo)致的精度損失。如圖5所示，BiSparse具體包含如下技術(shù)：

▲圖4 GeoMX與MXNET在不同帶寬下的測試精度曲線

表1 GeoMX與MXNET在不同模型下的單輪通信時延對比

▲圖5 雙向稀疏梯度傳輸技術(shù)流程示意圖

（1）梯度稀疏化技術(shù)。該技術(shù)旨在從大規(guī)模梯度張量中篩選出少量關(guān)鍵梯度，在保障模型精度幾乎無損的前提下，以減少實際傳輸?shù)奶荻葦?shù)量。假設(shè)某梯度張量包含N個梯度值，且指定壓縮比例為k（0＜k≤100%）。為快速確定梯度張量中絕對值最大的top-k個梯度值，我們按采樣率s從梯度張量中隨機采樣sN個梯度值，并按絕對值從大到小的順序進行排序，其中第skN個值T將被作為梯度張量中top-k個梯度絕對值的近似下界。因此，梯度張量中絕對值大于T的梯度值將被選為關(guān)鍵梯度并參與稀疏聚合。

（2）稀疏聚合與同步技術(shù)。該技術(shù)實現(xiàn)在稀疏空間中對關(guān)鍵梯度進行高效聚合與同步，包括映射、規(guī)約、分發(fā)與重構(gòu)。在映射階段，參與機構(gòu)的參數(shù)服務(wù)器將稀疏梯度張量編碼成元素（索引、梯度值）的向量，其中索引指梯度值在梯度張量中的相對偏移。在規(guī)約階段，中心機構(gòu)的全局參數(shù)服務(wù)器加和索引相同的梯度值，得到聚合稀疏梯度向量。在分發(fā)階段，各參與機構(gòu)的參數(shù)服務(wù)器從全局參數(shù)服務(wù)器獲取聚合稀疏梯度向量，并在重構(gòu)階段轉(zhuǎn)換為張量，隨后該張量被下發(fā)給工作節(jié)點完成梯度同步。

（3）殘余梯度修正技術(shù)。本文中，我們稱梯度稀疏化時未選中的非關(guān)鍵梯度為殘余梯度。為避免丟失殘余梯度而導(dǎo)致精度損失，該技術(shù)緩存殘余梯度并累加到下一輪的聚合梯度中，同時利用動量修正技術(shù)緩解因殘余梯度延遲更新引起的收斂震蕩。假設(shè)當(dāng)前輪次t的聚合梯度張量為Gt，上一輪次緩存的殘余梯度速度和位置張量為vt-1和ut-1，阻尼系數(shù)為m。當(dāng)前速度vt由衰減的歷史速度和當(dāng)前施加的聚合梯度根據(jù)vt=mvt-1+Gt得到，并用于更新位置ut=ut-1+vt，其中ut將代表聚合梯度參與梯度稀疏化。最后，梯度稀疏化篩選出的關(guān)鍵梯度對應(yīng)的速度值和位置值從張量vt和ut中置零，隨后vt和ut作為殘余梯度被緩存并累加到下一輪的聚合梯度中。

本文中，我們設(shè)置k= 1%，s=0.5%，m= 0.9，并在 1∶ 6～1 ∶ 100的廣域與局域帶寬比、參與機構(gòu)B和C各4個工作節(jié)點下對比了GeoMX關(guān)閉/啟用雙向稀疏梯度傳輸GeoMXBiSparse與半精度量化[6]GeoMX-FP16（除稀疏梯度外，梯度均以16位浮點數(shù)傳輸）技術(shù)的收斂精度、時間、輪數(shù)以及傳輸數(shù)據(jù)量。如圖6所示，稀疏化和半精度量化均能有效提高訓(xùn)練效率，其中稀疏化所取得的效率增益最為顯著。此外，稀疏化結(jié)合半精度量化GeoMX-BiSparse-FP16可進一步減少40%的獲得理想模型精度（90.4%）的時間。

為了驗證上述優(yōu)化技術(shù)應(yīng)用于廣域環(huán)境相比基礎(chǔ)GeoMX應(yīng)用于局域環(huán)境的增益，我們選擇GeoMX-1 Gbit/s作為對比。表2數(shù)據(jù)表明，僅采用稀疏化技術(shù)就能取得約20%的訓(xùn)練效率增益（由10.6 h降到8.6 h）、10%的通信效率增益（由7.24 s降到6.58 s）以及90%的傳輸數(shù)據(jù)壓縮率（上下行傳輸數(shù)據(jù)量分別由93.95 MB降到8.15 MB和9.90 MB），且保證收斂精度無損。對于兩種優(yōu)化技術(shù)相結(jié)合的方案，雖然稀疏梯度損失和半精度損失的雙重影響導(dǎo)致了一定的收斂精度下降（由90.8%降到90.6%），但是對于對近50%的訓(xùn)練和通信效率增益（由10.6 h降到5.8 h、由7.24 s降到3.88 s）而言，損失0.2%的模型精度是可容忍的。為模擬帶寬差異更顯著的廣域和局域通信環(huán)境，我們將廣域帶寬進一步限制到100/50/10 Mbit/s。由于廣域帶寬更加緊缺，每輪模型通信需要更長的時延，致使訓(xùn)練效率逐步劣化。在50 Mbit/s廣域帶寬下，GeoMX-BiSparse-FP16仍能取得趨同GeoMX-1 Gbit/s的訓(xùn)練效率（10 h vs.10.6 h）。當(dāng)廣域帶寬持續(xù)減小時，即使近95%的傳輸數(shù)據(jù)壓縮率也難以突破廣域通信瓶頸，這就要求我們必須探索更先進的通信優(yōu)化技術(shù)。

5 結(jié)束語

GeoMX是大數(shù)據(jù)與人工智能結(jié)合的必然產(chǎn)物，它推動著大型數(shù)據(jù)孤島的融合，并反饋大數(shù)據(jù)與人工智能的發(fā)展進程?；A(chǔ)軟件框架的普及與優(yōu)化技術(shù)的全面覆蓋將有力驅(qū)動GeoMX是的應(yīng)用與落地，并與聯(lián)邦學(xué)習(xí)和隱私計算等交叉領(lǐng)域技術(shù)互補，共同構(gòu)建安全且高效的下一代人工智能技術(shù)體系。

▲圖6 不同壓縮方案與廣域帶寬下的測試精度曲線

表2 不同壓縮方案與廣域帶寬下各維度優(yōu)化效果對比

致謝

感謝電子科技大學(xué)信息與通信工程學(xué)院張兆豐碩士對本文第4章技術(shù)與實驗部分的貢獻。