地理分布式機器學習：超越局域的框架與技術

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

（1. 電子科技大學，中國 成都，611731;2. 南方科技大學，中國 深圳，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技術的逐步商用，激增的數據為人工智能技術提供了豐富的數據資源，催生了人臉支付、輔助感知等新時代智能應用的萌芽。為強化智能應用的生命力，加速以深度學習為代表的智能模型的訓練過程，分布式機器學習[1]成為業(yè)界重點關注的關鍵技術。數據并行的分布式機器學習將完整數據均勻分配給數據中心內的多個處理節(jié)點。這些節(jié)點各自維護一個本地模型副本以并行計算其數據切片上的本地梯度，并遞交給其他節(jié)點計算全局梯度，隨后全局梯度在所有節(jié)點間同步，最終被用于更新各節(jié)點的本地模型。

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

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

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

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

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

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

2 地理分布式機器學習軟件框架

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

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

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

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

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

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

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

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

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

在HiPS通信模塊中，本文中我們主要關注以下技術：

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

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

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

3 HiPS架構及流量模型

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

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

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

▲圖2 分層參數服務器架構

▲圖3 分層參數服務器流量模型

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

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

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

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

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

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

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

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

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

4 BiSparse技術

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

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

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

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

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

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

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

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

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

5 結束語

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

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

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

致謝

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