作者：倪昊
這篇論文來自 2019 年 SysML 會(huì)議中 Parallel & Distributed Learning 的部分。
數(shù)據(jù)并行訓(xùn)練（Data parallel training） 已經(jīng)廣泛地運(yùn)用在在深度神經(jīng)網(wǎng)絡(luò)的分布式計(jì)算中，但是，分布式計(jì)算帶來的性能提升經(jīng)常受限于參數(shù)同步性能的瓶頸。作者等人提出了一種新的參數(shù)同步機(jī)制：Priority-based Parameter Propagation （P3），提高了模型的訓(xùn)練集群對(duì)網(wǎng)絡(luò)帶寬的利用率，加快了模型的訓(xùn)練速度。
我們首先來回顧一下神經(jīng)網(wǎng)絡(luò)的分布式訓(xùn)練。
深度神經(jīng)網(wǎng)絡(luò)（DNN）的分布式訓(xùn)練
神經(jīng)網(wǎng)絡(luò)的分布式訓(xùn)練通常有兩種策略。一種是模型并行，指的是分布式系統(tǒng)中的不同機(jī)器分別負(fù)責(zé)在單個(gè)網(wǎng)絡(luò)的不同部分計(jì)算，即把模型的不同層（layer）放到不同的工作節(jié)點(diǎn)上，常用與模型過大的情況，它的計(jì)算效率不高。另一種是數(shù)據(jù)并行，不同的機(jī)器有著整個(gè)模型的完全拷貝，每個(gè)機(jī)器只獲得整個(gè)數(shù)據(jù)的不同部分，即把訓(xùn)練數(shù)據(jù)分成許多塊，分給不同的 worker 節(jié)點(diǎn)，各個(gè)節(jié)點(diǎn)單獨(dú)進(jìn)行計(jì)算，但是共享一個(gè)模型，計(jì)算的結(jié)果通過某些方法結(jié)合起來。當(dāng)然，這兩種策略并不沖突，完全可以混合使用。
數(shù)據(jù)并行常用的架構(gòu)是參數(shù)服務(wù)器（Parameter Server），即 PS 架構(gòu)。
數(shù)據(jù)并行有兩種不同的并行模式：同步訓(xùn)練和異步訓(xùn)練。同步訓(xùn)練是指所有 worker 節(jié)點(diǎn)的梯度都計(jì)算完成，統(tǒng)一進(jìn)行梯度更新。而異步訓(xùn)練是指 worker 節(jié)點(diǎn)計(jì)算是獨(dú)立地從 PS 節(jié)點(diǎn)獲取模型參數(shù)，獨(dú)立計(jì)算梯度并更新。采用同步隨機(jī)梯度下降算法（SGD）算法的數(shù)據(jù)并行訓(xùn)練是一種非常流行的方法。
每個(gè)節(jié)點(diǎn)的工作大致可以分成三步：
節(jié)點(diǎn)從 PS 節(jié)點(diǎn)獲取最新的模型參數(shù)，對(duì)各自負(fù)責(zé)的訓(xùn)練數(shù)據(jù)進(jìn)行前向傳播。
各個(gè)節(jié)點(diǎn)獨(dú)立地對(duì)不同部分的訓(xùn)練數(shù)據(jù)進(jìn)行反向傳播，得到每個(gè)模型參數(shù)的梯度。
對(duì)各個(gè) worker 節(jié)點(diǎn)的梯度計(jì)算結(jié)果進(jìn)行同步，利用更新模型的參數(shù)
然后不斷的循環(huán)迭代。
節(jié)點(diǎn)的每一次迭代都需要通過網(wǎng)絡(luò)進(jìn)行海量參數(shù)的同步，這對(duì)網(wǎng)絡(luò)帶寬提出了非常高的要求。解決這個(gè)問題的一個(gè)方法是增加網(wǎng)絡(luò)的帶寬，但這樣做意味著極高的成本，同時(shí)隨著模型體量的增大，成本也無法估量。因此，我們應(yīng)該尋求在有限帶寬下的解決方案。
近年來一種流行的方案是梯度壓縮，但是這種方案的缺點(diǎn)在于丟失了一部分信息，可能會(huì)影響到模型的準(zhǔn)確度。
還有一種方法，就是提高網(wǎng)絡(luò)帶寬的利用率。在訓(xùn)練的過程中，如果我們?cè)诿看蔚瓿珊蟛湃ネ教荻龋蜁?huì)導(dǎo)致網(wǎng)絡(luò)流量爆炸性地增高，而在 worker 節(jié)點(diǎn)計(jì)算的過程中，網(wǎng)絡(luò)卻基本處于空閑狀態(tài)。因此，我們可以將已經(jīng)計(jì)算出來的梯度在 worker 節(jié)點(diǎn)計(jì)算其他層梯度的時(shí)候發(fā)給 PS 節(jié)點(diǎn)，即讓 worker 節(jié)點(diǎn)與 PS 節(jié)點(diǎn)之間的通信和 worker 節(jié)點(diǎn)的反向傳播同時(shí)進(jìn)行，這樣，網(wǎng)絡(luò)帶寬就可以得到更有效的利用。一些深度學(xué)習(xí)框架已經(jīng)實(shí)現(xiàn)了這一點(diǎn)，比如 TensorFlow、MXNet 和 Caffe2。
局限性與優(yōu)化
作者等人在工作中發(fā)現(xiàn)了以上方案還存在一定的局限性，通過優(yōu)化，模型對(duì)網(wǎng)絡(luò)帶寬的利用率和計(jì)算速度可以進(jìn)一步提高。
反向傳播中參數(shù)傳遞的優(yōu)先級(jí)
在神經(jīng)網(wǎng)絡(luò)模型的訓(xùn)練中，我們通常是：反向傳播——用反向傳播所得的參數(shù)進(jìn)行前向傳播，更細(xì)致一點(diǎn)，從最后一層開始進(jìn)行反向傳播的計(jì)算，再從第一層開始前向傳播，如下圖所示。
在第 N 次迭代中，我們從輸出層開始根據(jù) Loss 計(jì)算出梯度并更新 L4 層的參數(shù)，然后遞歸地求出輸入層即 L1 層的梯度同時(shí)更新整個(gè)模型的參數(shù)，然后我們從輸入層開始，進(jìn)行 N+1 次的迭代，利用更新后模型的參數(shù)一層一層地計(jì)算，直到獲得輸出層的結(jié)果，然后進(jìn)行再進(jìn)行反向傳播。
我們可以觀察到一點(diǎn)，在迭代的過程中，最先計(jì)算出來的參數(shù)（L4層的參數(shù)）總是最后被使用，而反向傳播最后計(jì)算出來的參數(shù)卻最先被前向傳播所使用。我們發(fā)現(xiàn)在神經(jīng)網(wǎng)絡(luò)的計(jì)算過程中，各個(gè)層的參數(shù)從獲得到使用中間的間隔是不同的，層數(shù)靠近輸出層，這一間隔也越大。在以往的分布式訓(xùn)練中，參數(shù)往往在某一層的反向傳播結(jié)束時(shí)就開始同步，這就有可能導(dǎo)致 L2 層的參數(shù)還沒有同步完，但是 L1 層已經(jīng)反向傳播結(jié)束，不得不等待 L2 層完成同步，L1 層完成參數(shù)同步然后再接受 PS 節(jié)點(diǎn)的參數(shù)更新，才能開始進(jìn)行 L1 層的前向傳播。這樣，就會(huì)導(dǎo)致前向傳播和反向傳播之間的間隔過大。
所以，我們認(rèn)為，在參數(shù)的同步中，層數(shù)低的層相比層數(shù)高的層，應(yīng)該具有更高的優(yōu)先級(jí)。
上圖中（a）代表以往的同步機(jī)制，我們對(duì)其進(jìn)行改進(jìn)，優(yōu)先同步較低的層的梯度，如（b）所示，當(dāng) L1 層的梯度計(jì)算完成，但是更高層的參數(shù)還沒有完成同步時(shí)，優(yōu)先進(jìn)行 L1 層的參數(shù)同步，這樣就縮短了反向傳播和正向傳播之間的間隔，但是并沒有增加網(wǎng)絡(luò)的負(fù)載。
參數(shù)同步中的粒度選取
參數(shù)同步所需的通信時(shí)間主要由三部分組成：
梯度從 worker 節(jié)點(diǎn)傳輸?shù)?PS 節(jié)點(diǎn)的時(shí)間
PS 節(jié)點(diǎn)利用梯度更新模型參數(shù)所需的時(shí)間
PS 節(jié)點(diǎn)發(fā)送更新完的參數(shù)到各個(gè) worker 節(jié)點(diǎn)
就如之前描述的一樣，以往我們以層為粒度去進(jìn)行參數(shù)同步，如圖（a）所示，通過實(shí)現(xiàn)數(shù)據(jù)的傳輸和計(jì)算并行執(zhí)行來盡可能地利用空閑的網(wǎng)絡(luò)帶寬，縮短通信所需要的時(shí)間。
但是在模型中，如果有一層的參數(shù)非常多，比如圖（a）中的 L2 層，它完成參數(shù)同步所需的時(shí)間每一步都是 L1 和 L3 的三倍，我們會(huì)發(fā)現(xiàn)他會(huì)拖慢參數(shù)同步，或者說「阻塞」了參數(shù)的傳遞。比如在時(shí)間為 4 時(shí)，只有梯度的傳輸，卻沒有進(jìn)行任何參數(shù)更新的計(jì)算。這是因?yàn)槲覀兊膮?shù)同步是以層為單位，每一層中參數(shù)的更新必須要等到該層所有的節(jié)點(diǎn)的梯度都傳輸完成，但是我們并不需要等到接收完所有節(jié)點(diǎn)的梯度后再開始更新參數(shù)。
所以，我們可以采取更小的參數(shù)同步粒度，比如，將 L2 層的參數(shù)分成 3 份，單獨(dú)地進(jìn)行參數(shù)同步，如（b）所示。通過采取更小的粒度，我們可以盡可能地提升帶寬的利用率，直觀來講就是增大圖中不同層之間的重合度。
總結(jié)
以上兩點(diǎn)就是 P3 采取的優(yōu)化方式，作者在 MXNet 的基礎(chǔ)上實(shí)現(xiàn)了這一機(jī)制，有興趣的同學(xué)可以訪問 GitHub 閱讀相關(guān)源碼。根據(jù)作者的測(cè)試，使用 P3，ResNet-50、Sockeye、VGG-19 等模型的訓(xùn)練效率分別提升了25%、38%、66%。這種同步機(jī)制的優(yōu)點(diǎn)在于它與其他的優(yōu)化方式并不沖突，且不會(huì)影響模型的精度，但是它對(duì)分布式模型訓(xùn)練效率的提升也有限，且并不是所有的模型都有明顯的提升，特別是模型較小或者網(wǎng)絡(luò)帶寬十分受限時(shí)。
參考文獻(xiàn)
