姜 珊, 程志超, 曾榮飛, 黃 敏, 王興偉*

(東北大學(xué) a.計算機(jī)科學(xué)與工程學(xué)院； b.軟件學(xué)院； c.信息科學(xué)與工程學(xué)院, 遼寧 沈陽 110819)

由于機(jī)器學(xué)習(xí)(Machine Learning，ML)框架和平臺在開發(fā)上具有很大的靈活性，同時，越來越豐富的數(shù)據(jù)集以及逐漸完善的高性能計算體系，也使得人工智能領(lǐng)域得以迅猛發(fā)展。隨著機(jī)器學(xué)習(xí)模型愈發(fā)復(fù)雜，訓(xùn)練深度逐漸加深，訓(xùn)練模型的計算成本超出了單個機(jī)器的負(fù)荷能力，因此，分布式機(jī)器學(xué)習(xí)(Distributed Machine Learning，DML)成為短時間內(nèi)完成大規(guī)模模型訓(xùn)練的有效解決方案之一。卷積神經(jīng)網(wǎng)絡(luò)(Convolutional Neural Networks，CNN)作為ML的一種代表性算法，主要用于圖像特征提取，已經(jīng)成功應(yīng)用在圖像識別和自然語言處理領(lǐng)域。隨著計算機(jī)軟硬件的不斷發(fā)展，網(wǎng)絡(luò)模型日益復(fù)雜，數(shù)據(jù)類型愈加豐富，數(shù)據(jù)量急速增長，單一機(jī)器已經(jīng)無法滿足需求，分布式CNN成為新的解決方法。

在分布式CNN模型訓(xùn)練中，每次迭代都是計算節(jié)點(diǎn)先接收到更新后的參數(shù)，然后根據(jù)各自的訓(xùn)練數(shù)據(jù)子集進(jìn)行計算梯度，最后，聚集不同節(jié)點(diǎn)的梯度以更新模型參數(shù)。因此，隨著DML集群規(guī)模的增長，通信可能成為制約分布式CNN模型訓(xùn)練速度的瓶頸之一。同時，隨著計算機(jī)硬件加速器的快速發(fā)展，如GPU和FPGAs，頻繁的參數(shù)/梯度交換很容易使網(wǎng)絡(luò)通信成為瓶頸，從而降低DML的訓(xùn)練性能[1-2]。

當(dāng)前系統(tǒng)中的數(shù)據(jù)傳輸時間都是在執(zhí)行過程中確定的，在這期間并沒有考慮參數(shù)計算和通信重疊帶來的影響。TensorFlow是一種用圖來表示計算的平臺。在數(shù)據(jù)并行模式下，每個計算節(jié)點(diǎn)都有一套完整的模型副本和部分?jǐn)?shù)據(jù)子集，即數(shù)據(jù)集被分割成多個子集。每次迭代由參與的計算節(jié)點(diǎn)使用相同的計算圖進(jìn)行處理，每次迭代通常持續(xù)幾毫秒到幾秒鐘。在每次迭代結(jié)束時，參數(shù)服務(wù)器會聚合所有節(jié)點(diǎn)的梯度來更新模型參數(shù)，在此過程中服務(wù)器會交換大量數(shù)據(jù)，這種通信開銷對系統(tǒng)的吞吐量有很大影響，也限制了模型的可擴(kuò)展性[3-4]。

分布式CNN的迭代時間取決于計算時間和通信時間。每次迭代開始時，計算節(jié)點(diǎn)會從參數(shù)服務(wù)器中獲取參數(shù)，但是接收到的參數(shù)并不會同時被使用，而是根據(jù)底層計算圖中的依賴關(guān)系進(jìn)行使用[5]。由于計算節(jié)點(diǎn)之間接收參數(shù)的順序是隨機(jī)的，數(shù)據(jù)流模型可能有多個可行的遍歷，可能造成因參數(shù)未接收完畢而帶來的計算阻塞。因此，如果可以找到一個最優(yōu)的參數(shù)執(zhí)行順序，會有助于改善計算和通信之間的重疊度，進(jìn)而縮短迭代時間，加快模型訓(xùn)練。

1 相關(guān)工作

隨著DML集群規(guī)模的不斷增長，多個節(jié)點(diǎn)和服務(wù)器之間的通信可能會制約模型的訓(xùn)練速度。在典型的隨機(jī)梯度下降方法中，直接獲得解析解是十分困難且不易實(shí)現(xiàn)的，需通過迭代精化不斷逼近最優(yōu)解。因此，在每次迭代過程中，每個計算節(jié)點(diǎn)首先從服務(wù)器端獲取參數(shù)，根據(jù)各自的訓(xùn)練數(shù)據(jù)來計算梯度，然后將梯度推送至服務(wù)器，參數(shù)服務(wù)器將聚集所有計算節(jié)點(diǎn)的梯度以更新模型參數(shù)[6]。隨后計算節(jié)點(diǎn)再一次從服務(wù)器端獲取更新后的參數(shù)進(jìn)行下次迭代。

在分布式CNN中，大量且頻繁的參數(shù)/梯度交換很容易造成網(wǎng)絡(luò)阻塞，從而降低跨機(jī)器分布式模型訓(xùn)練的性能。為了應(yīng)對昂貴的參數(shù)同步成本，計算機(jī)相關(guān)硬件加速器逐漸發(fā)展，比如RDMA這種高速且低延遲的網(wǎng)絡(luò)技術(shù)以及其他一些計算機(jī)硬件加速器被不斷開發(fā)，用來加速神經(jīng)網(wǎng)絡(luò)模型訓(xùn)練。目前，有很多研究工作是針對如何加速分布式CNN模型訓(xùn)練，進(jìn)而提高訓(xùn)練性能展開的。從多個不同的角度進(jìn)行分析，研究出了許多有效的解決方案。除了一些硬件加速的研究，在模型和算法方面，模型壓縮和參數(shù)壓縮等是較為普遍的做法，比如網(wǎng)絡(luò)剪枝、模型量化、低秩估計以及模型蒸餾。本文重點(diǎn)研究在網(wǎng)絡(luò)層面分布式CNN模型訓(xùn)練中，如何通過改善計算和通信之間的重疊，即實(shí)現(xiàn)計算和通信的細(xì)粒度并行，將通信開銷盡可能地隱藏在計算之后。

計算圖是CNN模型訓(xùn)練的前提，無論是靜態(tài)構(gòu)建還是動態(tài)構(gòu)建，在常用的模型副本(即數(shù)據(jù)并行訓(xùn)練的模式)中，輸入數(shù)據(jù)被分割，并由參與計算的節(jié)點(diǎn)使用相同的計算圖進(jìn)行處理，每次迭代通常持續(xù)幾毫秒到幾秒鐘。在每次迭代結(jié)束時，服務(wù)器會和所有計算節(jié)點(diǎn)交換大量與參數(shù)更新相關(guān)的數(shù)據(jù)。這種通信開銷對系統(tǒng)的吞吐量有很大影響，也限制了其可擴(kuò)展性[7]。在耗時較長的模型訓(xùn)練中，即使通信開銷只稍微改善一點(diǎn)點(diǎn)，訓(xùn)練時間也可以縮短幾個小時。之前的研究工作試圖通過在所有計算節(jié)點(diǎn)中執(zhí)行相同的參數(shù)傳輸順序來解決這個問題。然而，這些研究工作是針對早期系統(tǒng)的模型結(jié)構(gòu)，在大規(guī)模機(jī)器學(xué)習(xí)TensorFlow架構(gòu)中[8]，這是一個不小的挑戰(zhàn)。

本文提出的啟發(fā)式算法側(cè)重于深度學(xué)習(xí)框架中的網(wǎng)絡(luò)優(yōu)化，在參數(shù)服務(wù)器架構(gòu)下采用模型副本模式進(jìn)行模型訓(xùn)練。啟發(fā)式算法帶來的性能改進(jìn)主要在2個方面：①提高網(wǎng)絡(luò)吞吐量；②以計算節(jié)點(diǎn)為訓(xùn)練對象，強(qiáng)制排序可以縮短迭代時間。雖然分散聚合方式(如all-reduce和Horovod[9])在高性能網(wǎng)絡(luò)中獲得了越來越多的關(guān)注，但本文沒有針對這類系統(tǒng)，而是基于參數(shù)服務(wù)器架構(gòu)下進(jìn)行研究。

在深度學(xué)習(xí)系統(tǒng)中，至關(guān)重要的通信成本隨著橫向擴(kuò)展而增加。在這種情況下，當(dāng)通信時間小于或等于計算時間，GPU利用率相對較高。此外，通信和計算的有效重疊對于提高吞吐量也很關(guān)鍵。目前已經(jīng)提出了3種技術(shù)來提高系統(tǒng)性能：①增加計算時間。可以通過增加批量來增加計算時間相對于通信時間的比例[10]，但是這種方案會降低模型訓(xùn)練的準(zhǔn)確率[11]，需要額外的校正機(jī)制，在資源限制下可能無法普遍適用[12-13]；②減少通信時間?？梢酝ㄟ^優(yōu)化機(jī)器學(xué)習(xí)算法以降低通信成本[14]，或者通過降低參數(shù)精度[15]來降低通信開銷；③嘗試改善計算和通信之間的重疊情況。例如，文獻(xiàn)[5]將卷積層參數(shù)和全連接層參數(shù)賦予不同優(yōu)先級，在整個網(wǎng)絡(luò)結(jié)構(gòu)中，卷積層參數(shù)始終優(yōu)先于全連接層參數(shù)進(jìn)行傳輸。由于每次CNN模型迭代都是先進(jìn)行卷積操作、激活函數(shù)以及池化操作等，最后經(jīng)過全連接層，因此，通過優(yōu)先發(fā)送卷積層參數(shù)可以改善計算和通信的有效重疊。但這種參數(shù)調(diào)度是粗粒度的，當(dāng)模型較為簡單時，會獲得一定的性能提升；當(dāng)模型深度較大時，存在多個卷積層，無法對參數(shù)進(jìn)行有效調(diào)度。文獻(xiàn)[11]在分布式深度神經(jīng)網(wǎng)絡(luò)中通過對參數(shù)進(jìn)行切片，并根據(jù)深度神經(jīng)網(wǎng)絡(luò)的分層結(jié)構(gòu)為它們指定優(yōu)先級[16]，也就是第一層為最高優(yōu)先級，第二層為次高優(yōu)先級，以此類推。每次都是從待執(zhí)行隊列中選取優(yōu)先級最高的參數(shù)，被首先發(fā)送出去。然而這種解決方案只關(guān)注每個計算節(jié)點(diǎn)中參數(shù)的傳輸順序，忽略了在網(wǎng)絡(luò)結(jié)構(gòu)中的整體情況，從而不能很好地協(xié)調(diào)來自不同計算節(jié)點(diǎn)的參數(shù)。在大規(guī)模DML中會有多個計算節(jié)點(diǎn)向服務(wù)器交換參數(shù)，網(wǎng)絡(luò)中存在大量流量，此時這種終端優(yōu)先級參數(shù)調(diào)度將不再有效。

本文提出一種細(xì)粒度的參數(shù)傳輸調(diào)度機(jī)制，通過對底層計算圖的關(guān)鍵路徑進(jìn)行分析來獲得參數(shù)傳輸?shù)慕谱顑?yōu)調(diào)度。通過一種啟發(fā)式算法完成參數(shù)的優(yōu)先級排序，實(shí)現(xiàn)計算和通信的最大重疊。在網(wǎng)絡(luò)傳輸方面采用基于優(yōu)先級的流調(diào)度方案，進(jìn)一步加速卷積神經(jīng)網(wǎng)絡(luò)的訓(xùn)練。參數(shù)優(yōu)先級會在數(shù)據(jù)包報頭中攜帶，使得參數(shù)在整個網(wǎng)絡(luò)結(jié)構(gòu)中都可被網(wǎng)卡和交換機(jī)等識別。基于計算圖的細(xì)粒度參數(shù)調(diào)度，可以在多個計算節(jié)點(diǎn)之間以不同的順序激活網(wǎng)絡(luò)傳輸，實(shí)現(xiàn)所有計算節(jié)點(diǎn)都遵循相同的最優(yōu)調(diào)度順序。通過細(xì)粒度的參數(shù)調(diào)度機(jī)制，可以改善分布式CNN訓(xùn)練中通信和計算的重疊，此外，該機(jī)制還可以使用所有由ML框架支持的加速器/網(wǎng)絡(luò)結(jié)構(gòu)，可以移植，且不需要對ML框架進(jìn)行修改。

2 卷積神經(jīng)網(wǎng)絡(luò)

2.1 CNN模型訓(xùn)練

在自然語言處理等領(lǐng)域中，文本和圖像無法直接被機(jī)器識別，需要將文本和圖像進(jìn)行處理，轉(zhuǎn)換為數(shù)值型數(shù)據(jù)再進(jìn)行使用。CNN的核心是卷積層，卷積層包含多個大小不同的濾波器，輸入數(shù)據(jù)經(jīng)過卷積處理后完成特征的提取。輸入數(shù)據(jù)經(jīng)歷卷積操作、激活函數(shù)以及池化操作后，最后進(jìn)入全連接層完成輸入數(shù)據(jù)的特征提取。CNN模型由于加入了卷積層和池化層，相比于傳統(tǒng)的神經(jīng)網(wǎng)絡(luò)模型，增強(qiáng)了模型的泛化能力，支持更復(fù)雜的網(wǎng)絡(luò)模型，特征提取的準(zhǔn)確度也較高。

CNN具有分層結(jié)構(gòu)(圖1)，主要包括卷積層、池化層、激活層和全連接層等，每一層都包含大量的神經(jīng)元，相鄰層的這些神經(jīng)元相互連接[5]。一般來說，卷積層包含的參數(shù)較少，但需要進(jìn)行的計算較多。

圖1 CNN結(jié)構(gòu)Fig.1 CNN structure

CNN的訓(xùn)練過程分為2個階段。第一階段是前向傳播，第二階段是反向傳播階段。訓(xùn)練過程(圖2)分為5個步驟：

圖2 CNN訓(xùn)練過程Fig.2 CNN training process

Step 1: 初始化權(quán)重和偏置；

Step 2: 前向傳播；

Step 3: 求出誤差；

Step 4: 當(dāng)誤差大于期望值時，將誤差傳回網(wǎng)絡(luò)中；當(dāng)誤差等于或小于期望值時，結(jié)束訓(xùn)練;

Step 5: 更新權(quán)值，回到Step 2繼續(xù)迭代。

2.2 PS架構(gòu)

隨著CNN模型變得愈發(fā)復(fù)雜，計算工作量也相應(yīng)增加。一方面模型參數(shù)過大，單機(jī)內(nèi)存空間不足，需要采用分布式存儲；另一方面訓(xùn)練數(shù)據(jù)過多，單機(jī)訓(xùn)練太慢，需要增加訓(xùn)練節(jié)點(diǎn)，來提高并發(fā)訓(xùn)練速度。因此，分布式CNN逐漸發(fā)展起來。在實(shí)踐中最廣泛采用的是參數(shù)服務(wù)器(PS)架構(gòu)[6]，并且得到了主流DML框架的良好支持，如TensorFlow、PyTorch[17]等。

本文中設(shè)定分布式CNN以數(shù)據(jù)并行的方式在PS架構(gòu)下訓(xùn)練，同步模式采用批量同步并行(BSP)。在每次迭代過程中，每個計算節(jié)點(diǎn)基于不同的數(shù)據(jù)子集進(jìn)行迭代，最后由參數(shù)服務(wù)器聚合來自不同計算節(jié)點(diǎn)的梯度，更新模型參數(shù)用于下次迭代。也就是說，每次迭代過程包括從服務(wù)器獲取參數(shù)、前向傳播計算、反向傳播計算、向服務(wù)器推送梯度、聚合梯度和更新參數(shù)。其中一些可以進(jìn)行重疊，以達(dá)到更快的訓(xùn)練速度。參數(shù)服務(wù)器架構(gòu)如圖3所示。

圖3 PS架構(gòu)Fig.3 PS framework

參數(shù)服務(wù)器架構(gòu)主要包含Server和Worker 2個部分，訓(xùn)練的具體流程如下：

Step 1: 將訓(xùn)練數(shù)據(jù)分配給不同的Worker;

Step 2: Worker讀取一個mini batch訓(xùn)練數(shù)據(jù)，從Server端獲取最新的參數(shù)、計算梯度和推送梯度給Server;

Step 3: Server接收所有Worker上傳的梯度，聚合后更新參數(shù)。

總迭代時間(T)、通信時間(t)和計算時間(Ct)之間的關(guān)系通常表示為T≤t+Ct，這是因?yàn)橛嬎愫屯ㄐ趴赡苤丿B[18]。通信/計算比ρ、重疊系數(shù)α和GPU利用率U的關(guān)系如下:

(1)

當(dāng)ρ<1時，通信時間小于總計算時間，表示此時GPU利用率較高，然而通信和計算的不良重疊會導(dǎo)致GPU利用率偏低。由于通信成本是隨著橫向擴(kuò)展而增加的[19]，當(dāng)通信時間小于或等于計算時間時，GPU利用率是較高的。因此，更好地優(yōu)化通信和計算的有效重疊，可增加分布式CNN模型訓(xùn)練性能，提高吞吐量縮短迭代延遲。

2.3 計算圖

計算圖通常分為構(gòu)建和執(zhí)行2個階段。計算圖的組成部分見表1。

表1 計算圖的組成

在分布式CNN模型訓(xùn)練過程中，每個Worker都有相同的模型副本。然而在參數(shù)服務(wù)器端有一個不同于Worker的計算圖，該計算圖中的參數(shù)涉及5個操作，分別是參數(shù)聚合、發(fā)送參數(shù)、接收參數(shù)、讀取參數(shù)和更新參數(shù)。通常情況下，在參數(shù)服務(wù)器上的聚合、讀取和更新是輕量級的。此時參數(shù)的傳輸是由Worker驅(qū)動的，每次迭代都是由參數(shù)服務(wù)器激活所有發(fā)送和接收操作，參數(shù)服務(wù)器負(fù)責(zé)控制網(wǎng)絡(luò)傳輸，因此，在參數(shù)服務(wù)器中不會出現(xiàn)通信計算重疊的問題，只需考慮Worker中的計算圖優(yōu)化問題。

在Worker計算圖中，所有獲取參數(shù)操作(recv)都是根操作，推送參數(shù)操作都是葉操作。在某些情況下可以這樣認(rèn)為，recv操作可能會阻礙計算圖中分支計算，造成計算阻塞，延遲下次迭代的開始。本文的研究目標(biāo)是為了達(dá)到更大的吞吐量和更小的延遲，結(jié)合CNN 分層結(jié)構(gòu)特點(diǎn)，通過更好地預(yù)測通信和計算重疊來縮短迭代時間，提高分布式訓(xùn)練性能。

計算圖中的節(jié)點(diǎn)主要涉及2種類型的操作：①計算操作，如乘法、卷積等；②通信操作，如讀取和更新。每個參數(shù)都是獨(dú)立讀取和更新的。同樣,每次迭代也有2個階段：前向傳播階段和反向傳播階段。在前向傳播中，會根據(jù)模型的輸入計算損失函數(shù)；在反向傳播階段，基于計算的損失來更新模型參數(shù)。例如，圖4為一個簡單計算圖，其中存在2種可能的參數(shù)傳輸方案。如果recv1在recv2和recv3操作完成之后執(zhí)行，則緩解了計算阻塞；若首先執(zhí)行recv1操作，會引發(fā)計算阻塞，導(dǎo)致迭代時間增加。因此，在分布式環(huán)境中，可以根據(jù)計算圖中的依賴關(guān)系進(jìn)行參數(shù)傳輸調(diào)度，優(yōu)化計算和通信間的重疊程度，從而加快分布式CNN模型訓(xùn)練。

圖4 不同參數(shù)傳輸順序?qū)τ嬎阕枞挠绊?/p>

在分布式CNN中，隨著模型越來越復(fù)雜，參數(shù)也越來越多，需要將這些參數(shù)分散到多個節(jié)點(diǎn)進(jìn)行存儲和更新，主要包含2種方式：①In-graph模式，即將數(shù)據(jù)分發(fā)到1個節(jié)點(diǎn)上，然后再把數(shù)據(jù)分配給相應(yīng)的節(jié)點(diǎn)，在數(shù)據(jù)量巨大的時候，這種模式會嚴(yán)重制約模型訓(xùn)練速度，但由于其不依賴于外部，因此，更容易訪問，更容易在云環(huán)境中部署;②Between-graph模式,即數(shù)據(jù)分片分別保存在節(jié)點(diǎn)本地，每次迭代只需按照計算圖來訓(xùn)練本地數(shù)據(jù)集，再將計算的梯度推送到參數(shù)服務(wù)器，大數(shù)據(jù)下的深度學(xué)習(xí)更適合這種模式。數(shù)據(jù)轉(zhuǎn)換、傳輸和聚合被定義為標(biāo)準(zhǔn)數(shù)據(jù)流操作，其中，聚合被抽象為數(shù)據(jù)流模型中的單個操作，因此，將參數(shù)聚合實(shí)現(xiàn)為計算圖的一部分。

3 參數(shù)傳輸調(diào)度機(jī)制

參數(shù)傳輸調(diào)度機(jī)制可以描述為如何在計算圖中選擇一個最優(yōu)的操作執(zhí)行順序。其研究目標(biāo)是改善通信和計算重疊，加快模型訓(xùn)練。該算法的輸入包括2個部分：①資源關(guān)聯(lián)圖，即具有與每個操作相關(guān)聯(lián)的資源標(biāo)簽計算圖，其中，計算操作被分配給相應(yīng)的計算資源，通信操作被分配給相應(yīng)的通信信道；②預(yù)測時間表，即時間預(yù)測器預(yù)測給定操作的執(zhí)行時間。對于計算操作表示計算時間，對于通信操作表示傳輸時間。該算法的輸出是使迭代時間最小的可行操作拓?fù)漤樞颉?/p>

3.1 屬性更新算法

分布式CNN模型訓(xùn)練中，Worker在每次迭代開始時從參數(shù)服務(wù)器獲取參數(shù)，所有參數(shù)不會同時被使用，而是根據(jù)底層計算圖中的依賴關(guān)系進(jìn)行使用。然而一個特定的參數(shù)傳輸順序可能有助于更快的計算，因此，確定參數(shù)傳輸?shù)淖罴颜{(diào)度對于減少計算阻塞(由計算圖相關(guān)性決定)，改善計算通信重疊及縮短迭代時間至關(guān)重要。

算法1 操作屬性更新

∥更新當(dāng)前未完成接收參數(shù)操作的屬性

1: function UPDATE(G)

2: for each op∈計算圖Gdo

3:R←{?op∈G;opisrecv}

4: end for

5: for each op∈操作集Rdo

6:op.graph←獲取當(dāng)前op的計算圖

7: 遍歷與當(dāng)前op有數(shù)據(jù)依賴的節(jié)點(diǎn)

8: count←節(jié)點(diǎn)個數(shù)

9:op.count←op.count+count

10: end for

11: end function

本文側(cè)重于深度學(xué)習(xí)框架中的網(wǎng)絡(luò)優(yōu)化，通過對底層計算圖的關(guān)鍵路徑分析來獲得參數(shù)傳輸?shù)慕谱顑?yōu)調(diào)度，利用細(xì)粒度的調(diào)度來解決隨機(jī)參數(shù)的傳輸問題，提高分布式CNN模型訓(xùn)練的性能。

進(jìn)行算法設(shè)計時，優(yōu)先考慮加速計算圖中關(guān)鍵路徑的傳輸，根據(jù)計算圖中的節(jié)點(diǎn)依賴關(guān)系和每個操作的執(zhí)行時間來分配優(yōu)先級，對于計算阻塞更少的參數(shù)傳輸給予更高的優(yōu)先級。通過對計算圖中的所有recv操作節(jié)點(diǎn)進(jìn)行優(yōu)先級排序，選擇通信和計算開銷最小的recv操作，多次迭代直至遍歷所有recv操作節(jié)點(diǎn)。添加recv操作優(yōu)先級屬性，通過遍歷當(dāng)前計算子圖來實(shí)現(xiàn)。

3.2 優(yōu)先級分配算法

優(yōu)先級分配，即為了能夠限制執(zhí)行只接受一個有效集，考慮利用優(yōu)先級來實(shí)現(xiàn)。優(yōu)先級是分配給資源關(guān)聯(lián)圖中一個操作的正整數(shù)，優(yōu)先級較高的操作被賦予較高的優(yōu)先級。如果多個操作的相對順序無關(guān)緊要，則分配相同的優(yōu)先級。優(yōu)先級僅指定在給定資源準(zhǔn)備執(zhí)行的隊列中候選操作之間的相對順序，并且所得到的順序?qū)⑷匀蛔袷赜捎嬎銏D指定的拓?fù)漤樞?。在分布式模型?xùn)練中，當(dāng)一個資源需要從準(zhǔn)備執(zhí)行的隊列中選擇一個新的項目時，它會從包含最高優(yōu)先級的隊列中隨機(jī)選擇。

算法2 優(yōu)先級分配

∥對所有recv op進(jìn)行優(yōu)先級排序

1: function ASSIGNPRIORITY(G,Time)

2: Update(G)

3: ?op∈G;op.time←calculating time

4: for R is not emptry do

5: Find the minimum op.time

6: Remove op from R;∥從R中移除當(dāng)前recv op

7:op.priority←op.count;

8: end for

9: end function

優(yōu)先級分配算法主要是遍歷計算圖中的recv op，對未完成的操作進(jìn)行優(yōu)先級排序，然后獲得與當(dāng)前recv op相關(guān)聯(lián)的參數(shù)，將此時計數(shù)器數(shù)值賦予該參數(shù)為優(yōu)先級數(shù)。該算法綜合考慮計算圖的依賴關(guān)系和實(shí)際執(zhí)行時間，通過估計每個未完成recv op 的未來通信和計算開銷，對其進(jìn)行排序，進(jìn)而完成與之相依賴的參數(shù)優(yōu)先級的分配。該算法由2個步驟組成，首先假設(shè)所有的操作都有相等的開銷，只根據(jù)計算圖中的節(jié)點(diǎn)依賴關(guān)系來分配優(yōu)先級，對計算阻塞更少的傳輸給予更高的優(yōu)先級。其次優(yōu)先考慮加快關(guān)鍵路徑上的傳輸，同時綜合考慮每個操作的實(shí)際執(zhí)行時間，通過時間預(yù)測器來估計每個操作的執(zhí)行時間，結(jié)合計算圖中的依賴關(guān)系，對最大化計算/通信重疊的傳輸進(jìn)行優(yōu)先級排序。

4 具體實(shí)現(xiàn)

4.1 參數(shù)優(yōu)先級分配

首先獲取資源關(guān)聯(lián)圖并評估每個操作的執(zhí)行時間，結(jié)合計算圖中節(jié)點(diǎn)的依賴關(guān)系，執(zhí)行操作分配優(yōu)先級算法，將此時的計數(shù)器數(shù)值賦予與該recv op相依賴的參數(shù)，作為該參數(shù)的優(yōu)先級數(shù)。實(shí)現(xiàn)過程包括以下2個步驟：

(1)確定最佳順序。為了最大化重疊程度，應(yīng)盡可能早地激活通信操作，對recv op 節(jié)點(diǎn)進(jìn)行優(yōu)先級排序，使用貪婪算法來獲得參數(shù)更新的最佳順序。在每一次迭代中，選擇激活所需計算成本最小的參數(shù)，并且它所依賴的參數(shù)被標(biāo)記為相同的優(yōu)先級數(shù)。重復(fù)該過程，直到訪問完所有recv op。

(2)強(qiáng)制執(zhí)行最優(yōu)順序。將與當(dāng)前recv op相關(guān)聯(lián)的參數(shù)賦予相同的優(yōu)先級數(shù)，且始終攜帶在其數(shù)據(jù)包報頭中，實(shí)現(xiàn)在整個網(wǎng)絡(luò)結(jié)構(gòu)中傳輸時，仍然以此優(yōu)先級作為優(yōu)先轉(zhuǎn)發(fā)的依據(jù)。

4.2 網(wǎng)絡(luò)傳輸

在訓(xùn)練開始期間，在不同的節(jié)點(diǎn)之間建立多個具有不同優(yōu)先級的流，并生成通道表，同時為每一個流分配一個唯一的本地標(biāo)識。當(dāng)需要發(fā)送參數(shù)時，首先查找參數(shù)-優(yōu)先級映射表，獲得一個優(yōu)先級編號，然后根據(jù)優(yōu)先級編號查找通道表，獲得一個合適的流進(jìn)行發(fā)送參數(shù)。該參數(shù)被發(fā)送到網(wǎng)絡(luò)結(jié)構(gòu)中，在整個網(wǎng)絡(luò)結(jié)構(gòu)中，優(yōu)先級標(biāo)簽總是在數(shù)據(jù)包報頭中攜帶。當(dāng)多個參數(shù)同時到達(dá)交換機(jī)時，交換機(jī)可以通過標(biāo)簽來區(qū)分它們，并根據(jù)優(yōu)先級標(biāo)簽進(jìn)一步轉(zhuǎn)發(fā)它們。

網(wǎng)絡(luò)參數(shù)調(diào)度可以在較為復(fù)雜的模型訓(xùn)練中使用，終端參數(shù)調(diào)度則不能發(fā)揮優(yōu)勢。例如，在DML集群中有多個計算節(jié)點(diǎn)，為了簡化起見，在圖中只顯示了2個計算節(jié)點(diǎn)和1個參數(shù)服務(wù)器，圖5為終端參數(shù)調(diào)度和網(wǎng)絡(luò)參數(shù)調(diào)度之間的比較。

圖5 終端參數(shù)調(diào)度與網(wǎng)絡(luò)參數(shù)調(diào)度的比較

對于網(wǎng)絡(luò)調(diào)度，流的優(yōu)先級數(shù)和參數(shù)的優(yōu)先級數(shù)相匹配，即較高優(yōu)先級的參數(shù)被分配較高優(yōu)先級的流。當(dāng)Worker1和Worker2分別發(fā)送一個參數(shù)，假設(shè)Worker1發(fā)送參數(shù)的優(yōu)先級較低，Worker2發(fā)送參數(shù)的優(yōu)先級較高，對于網(wǎng)絡(luò)調(diào)度，可以通過識別優(yōu)先級標(biāo)簽來選擇發(fā)送參數(shù)的順序，即高優(yōu)先級參數(shù)將被調(diào)度在低優(yōu)先級參數(shù)之前發(fā)送；但終端調(diào)度只能保證參數(shù)在本地的執(zhí)行順序，無法滿足在網(wǎng)絡(luò)結(jié)構(gòu)中的有效調(diào)度，利用終端調(diào)度的實(shí)際情況是Worker1上的低優(yōu)先級參數(shù)和Worker2上的高優(yōu)先級參數(shù)沒有執(zhí)行調(diào)度(此時終端調(diào)度不起作用)，而是這2個參數(shù)平等地共享有限的帶寬，Worker2上的參數(shù)需要花費(fèi)2倍于網(wǎng)絡(luò)調(diào)度的時間才能到達(dá)服務(wù)器。這種情況還會隨著計算節(jié)點(diǎn)數(shù)量的增加而惡化，因此，縱觀整個網(wǎng)絡(luò)結(jié)構(gòu)，網(wǎng)絡(luò)參數(shù)調(diào)度比終端參數(shù)調(diào)度更有擴(kuò)展性，調(diào)度效率也更高。

4.3 實(shí)驗(yàn)分析

本文通過對VGG16和ResNet32等模型進(jìn)行實(shí)驗(yàn)，分析提出參數(shù)傳輸調(diào)度方法的效率。實(shí)驗(yàn)環(huán)境為Intel(R) Core(TM) i7-7500U CPU @ 2.70 GHz 2.90 GHz、windows10、Tensorflow-gpu 2.2.0、Python 3.8.3及conda 4.9.2，數(shù)據(jù)集為cifar10。

通過分析VGG16和ResNet32的參數(shù)分布(圖6、圖7)，可以發(fā)現(xiàn)很多CNN模型的參數(shù)是偏斜分布的，即最后幾層占據(jù)總參數(shù)的大部分。比如在VGG16模型中，全連接層中的參數(shù)分別占總參數(shù)的80%左右。因此，通過實(shí)現(xiàn)卷積層和全連接層的參數(shù)并行傳輸，優(yōu)化參數(shù)傳輸?shù)募?xì)粒度調(diào)度，才有可能加速分布式CNN的模型訓(xùn)練。

圖6 VGG16 參數(shù)分布Fig.6 VGG16 parameter distribution

圖7 ResNet32 參數(shù)分布Fig.7 ResNet32 parameter distribution

gRPC為每個計算節(jié)點(diǎn)-參數(shù)服務(wù)器對提供一個通道，該對之間的所有傳輸都被發(fā)送到同一隊列。在給定時刻，每個通道只能有一個傳輸處于活動狀態(tài)。網(wǎng)絡(luò)傳輸生命周期如圖8所示，TensorFlow中通過gRPC的網(wǎng)絡(luò)傳輸涉及多個階段。當(dāng)接收方激活recv op時，它會向發(fā)送方發(fā)送一個傳輸請求，如果發(fā)送方的發(fā)送操作也處于活動狀態(tài)，則參數(shù)傳輸由gRPC發(fā)起。此外，還定義一個度量調(diào)度效率，來衡量調(diào)度性能收益，評估啟發(fā)式算法的質(zhì)量。輸入計算圖、資源集和時間表，時間表給出了每個操作的執(zhí)行時間。最長完成時間的上限是假設(shè)在執(zhí)行期間的任何給定時刻僅使用一個資源來計算，即操作是順序執(zhí)行的；最長完成時間的下限是假設(shè)所有的資源總是被利用來計算。

圖8 網(wǎng)絡(luò)傳輸生命周期Fig.8 Lifetime of a network transfer

實(shí)驗(yàn)使用的數(shù)據(jù)集cifar10是一個小型數(shù)據(jù)集，共有60 000張32*32大小的RGB彩色圖片，共10個分類，訓(xùn)練數(shù)據(jù)和測試數(shù)據(jù)分別為50 000張和10 000張。目前的實(shí)驗(yàn)是生成一個In-graph數(shù)據(jù)流模型，其中，所有計算節(jié)點(diǎn)的數(shù)據(jù)流都表示在一個計算圖中，然后對數(shù)據(jù)進(jìn)行分區(qū)。此計算圖的大小隨著計算節(jié)點(diǎn)數(shù)量的增加而增加，這可能會增加大型圖的流處理時間。接下來會使用Between-graph數(shù)據(jù)流模型，其中，每個計算節(jié)點(diǎn)的數(shù)據(jù)流模型都會單獨(dú)生成，使這種細(xì)粒度的參數(shù)調(diào)度機(jī)制更具有擴(kuò)展性，更適合大數(shù)據(jù)下的復(fù)雜模型。針對VGG16和ResNet32進(jìn)行對比實(shí)驗(yàn)，引入啟發(fā)式算法后的細(xì)粒度調(diào)度機(jī)制能夠有效縮短迭代時間，但對于動態(tài)和可變模型，該機(jī)制目前還無法準(zhǔn)確預(yù)測具有動態(tài)控制流的數(shù)據(jù)流模型或具有高度可變輸入大小的模型時序，因此，不準(zhǔn)確的預(yù)測會導(dǎo)致更長的迭代時間，這也是下一步研究需要解決的問題之一。

參數(shù)在計算節(jié)點(diǎn)和服務(wù)器之間進(jìn)行傳輸和聚合，在每一次迭代中，計算節(jié)點(diǎn)從服務(wù)器獲取參數(shù)，計算節(jié)點(diǎn)推送梯度到服務(wù)器，為了提高數(shù)據(jù)在網(wǎng)絡(luò)上傳輸?shù)木W(wǎng)絡(luò)利用率，可以通過參數(shù)計算和參數(shù)傳輸?shù)牟⑿袌?zhí)行及參數(shù)的細(xì)粒度流水線傳輸來實(shí)現(xiàn)。另一種避免這種網(wǎng)絡(luò)利用不足的解決方案是將數(shù)據(jù)分割成幾個部分，并行獨(dú)立傳輸每個塊。組塊數(shù)稱為算法深度。在這種情況下，當(dāng)一個塊在中央處理器上減少時，另一個塊可以通過網(wǎng)絡(luò)發(fā)送，即實(shí)現(xiàn)了跨各種塊的網(wǎng)絡(luò)傳輸和網(wǎng)絡(luò)處理的流水線操作。

在參數(shù)傳輸和聚合階段，需要根據(jù)計算圖的依賴關(guān)系確定recv操作的近似最優(yōu)執(zhí)行順序，且確保在本地模型中只有一個可行的執(zhí)行順序。首先確定最優(yōu)執(zhí)行順序，確定參數(shù)的優(yōu)先等級，盡可能早地激活計算阻塞更少的recv操作。針對每個recv操作盡可能添加少量的操作屬性，保證參數(shù)可以按照最優(yōu)順序進(jìn)行更新/激活。先在一臺機(jī)器上跟蹤迭代執(zhí)行10次，計算每個操作的執(zhí)行時間，接下來，通過使用一個迭代貪心算法來尋找最優(yōu)參數(shù)更新順序。每次迭代中，在計算節(jié)點(diǎn)獲取參數(shù)之前，需要評估完成計算操作所花費(fèi)的總時間，然后選擇激活當(dāng)前時間成本最小的參數(shù)，將和該參數(shù)所依賴的計算操作標(biāo)記為已完成。重復(fù)此過程，直到所有參數(shù)已訪問更新。

最后是執(zhí)行最優(yōu)參數(shù)順序，這是一個迭代過程，參數(shù)會按照前一階段確定的最優(yōu)順序進(jìn)行參數(shù)傳輸。在每一步中，都會找到所選參數(shù)直接或間接依賴的所有操作列表。為計算圖中的每個recv操作維護(hù)一個計數(shù)器，遍歷當(dāng)前計算子圖所有入度為零的操作，每遍歷到一個與之依賴的操作，便將當(dāng)前recv操作的計數(shù)器數(shù)值加1，以此類推，直至遍歷完當(dāng)前計算子圖。選擇發(fā)送參數(shù)時，要先尋找與之相依賴的recv操作的計數(shù)器數(shù)值，將該數(shù)值賦予為參數(shù)的優(yōu)先級數(shù)。這確保了在每個給定時間，只能執(zhí)行目標(biāo)參數(shù)更新所需的操作。但是向數(shù)據(jù)流模型添加額外的控制依賴不會改變計算圖的底層邏輯，強(qiáng)制執(zhí)行的最優(yōu)順序是計算圖中可以執(zhí)行的順序之一。

由于數(shù)據(jù)流模型可能有多個可行的遍歷，即數(shù)據(jù)流模型中計算操作的不同執(zhí)行順序都是有效的。計算圖中的葉節(jié)點(diǎn)代表激活網(wǎng)絡(luò)傳輸?shù)膔ecv操作。在參數(shù)服務(wù)器架構(gòu)下，通常采用同步參數(shù)更新，在每次迭代后產(chǎn)生的突發(fā)流量容易造成網(wǎng)絡(luò)阻塞，從而延遲參數(shù)傳輸。為了解決這個問題，基于計算圖的細(xì)粒度參數(shù)調(diào)度，可以在多個計算節(jié)點(diǎn)之間以不同的順序激活網(wǎng)絡(luò)傳輸，實(shí)現(xiàn)所有計算節(jié)點(diǎn)都遵循最優(yōu)的參數(shù)調(diào)度順序。通過細(xì)粒度的參數(shù)調(diào)度機(jī)制，可以改善分布式CNN訓(xùn)練中通信和計算的重疊。此外,該機(jī)制還可以使用所有由ML框架支持的加速器/網(wǎng)絡(luò)結(jié)構(gòu)，可移植，且不需要對ML框架進(jìn)行額外修改。

筆者在TensorFlow上實(shí)現(xiàn)了細(xì)粒度參數(shù)調(diào)度機(jī)制，其中,VGG16模型訓(xùn)練數(shù)據(jù)集cifar10，學(xué)習(xí)率設(shè)置為0.000 1，batch_size為16，epoch為100的訓(xùn)練結(jié)果如圖9所示。ResNet32模型同樣訓(xùn)練數(shù)據(jù)集cifar10,學(xué)習(xí)率為0.000 1，batch_size為32，epoch為100的訓(xùn)練結(jié)果如圖10所示。

圖9 VGG16模型訓(xùn)練Fig.9 VGG16 model training

圖10 ResNet32模型訓(xùn)練Fig.10 ResNet32 model training

實(shí)驗(yàn)分別在VGG16、ResNet32及CNN+LSTM 3種模型下進(jìn)行對比分析，實(shí)驗(yàn)結(jié)果如表2所示。

表2 CNN模型訓(xùn)練迭代時間

由表2可知，使用參數(shù)傳輸優(yōu)化算法后，3種模型都在不同程度上減少了迭代時延。特別在CNN+LSTM模型中，基準(zhǔn)迭代時間是241 s，引入啟發(fā)式算法后迭代時間縮短為209 s。在上述實(shí)驗(yàn)中，通過細(xì)粒度優(yōu)化參數(shù)的傳輸順序，迭代時間分別減少為之前的93.48%、78.21%和86.72%，加快了CNN模型訓(xùn)練速度。

5 結(jié) 論

為了加快卷積神經(jīng)網(wǎng)絡(luò)的模型訓(xùn)練，本文設(shè)計了一種參數(shù)傳輸調(diào)度機(jī)制，通過一種啟發(fā)式算法來實(shí)現(xiàn)基于計算圖的參數(shù)傳輸機(jī)制，從而避免了由于參數(shù)未接收完畢而造成的計算阻塞。為了縮短迭代延遲，首先對本地計算圖中所有節(jié)點(diǎn)強(qiáng)制執(zhí)行最優(yōu)順序，改善通信和計算的重疊，加快CNN的模型訓(xùn)練，從而提升分布式CNN的訓(xùn)練性能。通過對VGG16和ResNet32進(jìn)行測試對比，實(shí)驗(yàn)結(jié)果表明，使用啟發(fā)式算法后的CNN模型能夠在不同程度上縮短訓(xùn)練迭代時間，迭代時間縮短了7%～22%左右。但對于動態(tài)和可變模型，該機(jī)制目前還無法準(zhǔn)確預(yù)測具有動態(tài)控制流的數(shù)據(jù)流模型或具有高度可變輸入大小的模型時序，因此，不準(zhǔn)確的預(yù)測會導(dǎo)致更長的迭代時間，這也是下一步研究需要解決的問題之一。該算法在實(shí)際參數(shù)服務(wù)器架構(gòu)中應(yīng)用時，網(wǎng)絡(luò)鏈路、參數(shù)同步方式、參數(shù)分片等情況，還有待進(jìn)一步研究和優(yōu)化。