論文標(biāo)題：GPipe: Easy Scaling with Micro-Batch Pipeline Parallelism
論文鏈接：https://arxiv.org/abs/1811.06965
論文來源：Google

一、概述

如下圖所示，近過去十年中竹海，由于開發(fā)了促進(jìn)神經(jīng)網(wǎng)絡(luò)有效容量擴(kuò)大的方法，深度學(xué)習(xí)取得了巨大進(jìn)步。這種趨勢(shì)在圖像分類中表現(xiàn)得尤為明顯卷仑，就像ImageNet的模型容量增加帶來的精度提升一樣孽尽。在自然語言處理的背景下也可以觀察到類似的現(xiàn)象，簡單的淺層句子表示模型被其更深更大的對(duì)應(yīng)模型超越毁葱。

現(xiàn)象

然而，雖然更大的模型為許多領(lǐng)域帶來了顯著的質(zhì)量提升倾剿，但擴(kuò)大神經(jīng)網(wǎng)絡(luò)也帶來了重大的實(shí)踐挑戰(zhàn)蚌成。包括內(nèi)存限制和加速器（GPU或TPU）的通信帶寬在內(nèi)的硬件約束担忧，迫使用戶將更大的模型劃分為多個(gè)部分，并將不同的部分分配給不同的加速器惩猫。然而轧房，高效的模型并行算法非常難以設(shè)計(jì)和實(shí)現(xiàn)，這通常要求實(shí)踐者在擴(kuò)大容量绍绘、靈活性（或特定于特定任務(wù)和架構(gòu)）和訓(xùn)練效率之間做出艱難的選擇奶镶。結(jié)果，大多數(shù)高效的模型并行算法都是針對(duì)特定架構(gòu)和任務(wù)的陪拘。隨著深度學(xué)習(xí)應(yīng)用的不斷增多厂镇，人們對(duì)可靠且靈活的基礎(chǔ)設(shè)施的需求也在不斷增長，這些基礎(chǔ)設(shè)施可以讓研究者輕松地將神經(jīng)網(wǎng)絡(luò)擴(kuò)大到大量的機(jī)器學(xué)習(xí)任務(wù)藻丢。

為了解決這些挑戰(zhàn)剪撬，我們引入了GPipe，這是一個(gè)靈活的庫悠反，可以實(shí)現(xiàn)大型神經(jīng)網(wǎng)絡(luò)的有效訓(xùn)練残黑。通過在不同的加速器上分割模型并在每個(gè)加速器上支持re-materialization，GPipe可以擴(kuò)大任意深度神經(jīng)網(wǎng)絡(luò)架構(gòu)斋否，超越單個(gè)加速器的內(nèi)存限制梨水。在GPipe中，每個(gè)模型都可以被定義為一系列的層茵臭，連續(xù)的層群可以被分割成單元疫诽。每個(gè)單元然后被放置在一個(gè)單獨(dú)的加速器上〉┪基于這種劃分設(shè)置奇徒，我們提出了一個(gè)新的流水線并行算法，通過批量分割缨硝。我們首先將訓(xùn)練樣例的一個(gè)小批量劃分為更小的微批量摩钙，然后將每組微批量的執(zhí)行在各個(gè)單元中流水線化。我們采用同步的小批量梯度下降方法進(jìn)行訓(xùn)練查辩，其中的梯度在一個(gè)小批量的所有微批量中累積胖笛，并在小批量結(jié)束時(shí)應(yīng)用网持。因此，使用GPipe的梯度更新不會(huì)因?yàn)閯澐值臄?shù)量不同而不同长踊，這使得研究者可以通過部署更多的加速器來輕松訓(xùn)練越來越大的模型功舀。GPipe還可以與數(shù)據(jù)并行結(jié)合，進(jìn)一步擴(kuò)大訓(xùn)練規(guī)模身弊。

我們?cè)趫D像分類和機(jī)器翻譯上展示了GPipe的靈活性和效率辟汰。對(duì)于圖像分類，我們?cè)贗mageNet 2012數(shù)據(jù)集上阱佛，訓(xùn)練了AmoebaNet模型莉擒，并將輸入調(diào)整為480×480。通過增加模型的寬度瘫絮，我們將參數(shù)數(shù)量擴(kuò)大到5.57億，并實(shí)現(xiàn)了84.4%的最高驗(yàn)證精度填硕。在機(jī)器翻譯上麦萤，我們訓(xùn)練了一個(gè)有128層、60億參數(shù)的多語言Transformer模型扁眯，支持103種語言（102種語言到英語）壮莹。我們展示了該模型能夠超越單獨(dú)訓(xùn)練的3.5億參數(shù)的雙語Transformer Big模型在100個(gè)語言對(duì)上的表現(xiàn)。

二姻檀、方法

GPipe開源庫是以Lingvo框架為基礎(chǔ)實(shí)現(xiàn)的命满，不過它的核心設(shè)計(jì)思路是通用的，可以與其他框架結(jié)合绣版。

1. 接口

一個(gè)深度神經(jīng)網(wǎng)絡(luò)可以被定義為層的一個(gè)序列胶台，每一層包括一個(gè)前向傳播計(jì)算函數(shù)和參數(shù)集合。GPipe額外允許用戶為每一層指定一個(gè)可選的計(jì)算花銷評(píng)估函數(shù)杂抽。對(duì)于一個(gè)給定的劃分?jǐn)?shù)量诈唬，這個(gè)神經(jīng)網(wǎng)絡(luò)的層序列可以被劃分為個(gè)復(fù)合層（或者叫做單元）。我們用來代表從到之間的連續(xù)層缩麸。對(duì)應(yīng)于的參數(shù)集合為的并集铸磅，其前向傳播函數(shù)為，相應(yīng)的反向傳播函數(shù)可以通過從使用自動(dòng)符號(hào)微分獲得杭朱，另外花銷估計(jì)函數(shù)為阅仔。

GPipe的接口非常地簡單和直觀，只需要用戶指定：
①模型劃分的數(shù)量弧械；
②微批量的數(shù)量八酒；
③模型的層的序列與定義。

2. 算法

在GPipe中梦谜，用戶定義他們網(wǎng)絡(luò)中的層序列丘跌，包括模型參數(shù)袭景，前向計(jì)算函數(shù)和成本估計(jì)函數(shù)。然后闭树，GPipe將網(wǎng)絡(luò)分割成個(gè)單元耸棒，并將第個(gè)單元放在第個(gè)加速器上。在分割的邊界报辱，GPipe自動(dòng)插入了通信原語与殃，允許鄰近分割之間的數(shù)據(jù)傳輸。分割算法會(huì)試圖最小化所有單元的估計(jì)成本的差異碍现，從而通過同步所有分割的計(jì)算時(shí)間幅疼，以最大化流水線的效率。

在前向傳播過程中昼接，GPipe首先將大小為的每個(gè)小批量（mini-batch）劃分為個(gè)相等的微批量（micro-batch）爽篷，并通過個(gè)加速器進(jìn)行流水線傳輸。在反向傳播過程中慢睡，每個(gè)微批量的梯度都基于用于前向傳播的相同模型參數(shù)來計(jì)算逐工。在每個(gè)小批量結(jié)束時(shí)，所有個(gè)微批量的梯度都被累積起來漂辐，并用于更新所有加速器上的模型參數(shù)泪喊。這個(gè)過程如下圖(c)所示。

GPipe

如果在網(wǎng)絡(luò)中使用了批量歸一化（batch normalization）髓涯，那么在訓(xùn)練期間袒啼，輸入的充分統(tǒng)計(jì)量將在每個(gè)微批量以及在必要的情況下的復(fù)制品（replica）上計(jì)算。我們還會(huì)跟蹤整個(gè)小批量的充分統(tǒng)計(jì)量的移動(dòng)平均值纬纪，以供在評(píng)估期間使用蚓再。

3. 性能優(yōu)化

為了減少激活（activation）的內(nèi)存需求，GPipe支持re-materialization育八。在前向計(jì)算的過程中对途，每個(gè)加速器只保存在單元邊界處的輸出激活。在反向傳播的過程中髓棋，第個(gè)加速器重新計(jì)算復(fù)合前向傳播函數(shù)实檀。在這樣的處理下，峰值激活內(nèi)存需求減少到了按声，這里的是微批量大小膳犹，是每個(gè)單元的層數(shù)。作為對(duì)比签则，沒有re-materialization的情況下內(nèi)存需求為须床，這是因?yàn)樵谟?jì)算梯度需要上層梯度以及緩存的激活。

上圖(c)中所示的分割過程會(huì)引入一些加速器的空閑時(shí)間渐裂，被稱為氣泡開銷（bubble overhead）豺旬。這個(gè)氣泡時(shí)間在每個(gè)微批量上的平攤開銷為钠惩。在實(shí)驗(yàn)中，當(dāng)微批量數(shù)時(shí)族阅，我們發(fā)現(xiàn)氣泡開銷可以忽略不計(jì)篓跛。這部分原因是在反向傳播期間可以提前調(diào)度re-materialization，而無需等待來自前面層的梯度坦刀。

GPipe引入了較低的通信開銷愧沟，因?yàn)槲覀冎恍枰诩铀倨髦g的分割邊界傳遞激活張量。因此鲤遥，即使在沒有高速互聯(lián)的加速器上沐寺，我們也可以實(shí)現(xiàn)高效的擴(kuò)展性性能。

上圖(c)假設(shè)分割是均衡的盖奈。然而混坞，不同層的內(nèi)存需求和計(jì)算量通常是不均衡的。在這種情況下钢坦，不完美的分割算法可能導(dǎo)致負(fù)載不平衡拔第。更好的分割算法有可能改善我們啟發(fā)式方法的性能。

這一部分關(guān)于這一系列復(fù)雜度是如何得到的场钉，可以參看李沐視頻。

三懈涛、實(shí)驗(yàn)

實(shí)驗(yàn)

實(shí)驗(yàn)

實(shí)驗(yàn)

實(shí)驗(yàn)

實(shí)驗(yàn)

實(shí)驗(yàn)