[[425806]]

越來(lái)越多的應(yīng)用程序是由在不同加速器集上訓(xùn)練的大型復(fù)雜神經(jīng)網(wǎng)絡(luò)驅(qū)動(dòng)的。這個(gè)過(guò)程由 ML 編譯器促進(jìn)，將高級(jí)計(jì)算圖映射到低級(jí)、特定于設(shè)備的可執(zhí)行文件。為此，ML 編譯器需要解決許多優(yōu)化問(wèn)題，包括圖形重寫、設(shè)備上的操作分配、操作融合、張量的布局和平鋪以及調(diào)度。例如，在設(shè)備放置問(wèn)題中，編譯器需要確定計(jì)算圖中的操作到目標(biāo)物理設(shè)備之間的映射，以便可以最小化目標(biāo)函數(shù)，例如訓(xùn)練步驟時(shí)間。放置性能由復(fù)雜因素的混合決定，包括設(shè)備間網(wǎng)絡(luò)帶寬、峰值設(shè)備內(nèi)存、協(xié)同定位約束等，這使得啟發(fā)式算法或基于搜索的算法具有挑戰(zhàn)性，這些算法通常滿足于快速但低于-最佳，解決方案。此外，啟發(fā)式很難開(kāi)發(fā)和維護(hù)，尤其是在出現(xiàn)新的模型架構(gòu)時(shí)。

最近使用基于學(xué)習(xí)的方法的嘗試已經(jīng)證明了有希望的結(jié)果，但它們有許多限制，使其無(wú)法在實(shí)踐中部署。首先，這些方法不容易推廣到看不見(jiàn)的圖，尤其是那些由較新的模型架構(gòu)產(chǎn)生的圖，其次，它們的樣本效率很差，導(dǎo)致訓(xùn)練過(guò)程中的資源消耗很高。最后，它們只能解決單個(gè)優(yōu)化任務(wù)，因此無(wú)法捕獲編譯堆棧中緊耦合優(yōu)化問(wèn)題之間的依賴關(guān)系。

在最近在NeurIPS 2020 上作為口頭論文發(fā)表的“Transferable Graph Optimizers for ML Compilers ”中，我們提出了一種用于計(jì)算圖優(yōu)化 (GO) 的端到端、可轉(zhuǎn)移的深度強(qiáng)化學(xué)習(xí)方法，該方法克服了上述所有限制。與TensorFlow默認(rèn)優(yōu)化相比，我們?cè)谌齻€(gè)圖優(yōu)化任務(wù)上展示了 33%-60% 的加速。在包括Inception-v3、Transformer-XL和WaveNet在內(nèi)的由多達(dá) 80,000 個(gè)節(jié)點(diǎn)組成的不同代表性圖集上，GO 比專家優(yōu)化平均提高了 21%，比現(xiàn)有技術(shù)提高了 18%，提高了 15 倍更快的收斂。

ML 編譯器中的圖優(yōu)化問(wèn)題 ML 編譯器

中經(jīng)常出現(xiàn)三個(gè)耦合優(yōu)化任務(wù)，我們將它們表述為可以使用學(xué)習(xí)策略解決的決策問(wèn)題。每個(gè)任務(wù)的決策問(wèn)題可以重新定義為為計(jì)算圖中的每個(gè)節(jié)點(diǎn)做出決策。

第一個(gè)優(yōu)化任務(wù)是設(shè)備放置，其目標(biāo)是確定如何最好地將圖的節(jié)點(diǎn)分配給它運(yùn)行的物理設(shè)備，從而最大限度地減少端到端的運(yùn)行時(shí)間。

第二個(gè)優(yōu)化任務(wù)是操作調(diào)度。計(jì)算圖中的操作已準(zhǔn)備就緒當(dāng)它的傳入張量存在于設(shè)備內(nèi)存中時(shí)運(yùn)行。一個(gè)常用的調(diào)度策略是為每個(gè)設(shè)備維護(hù)一個(gè)就緒的操作隊(duì)列，并按先進(jìn)先出的順序調(diào)度操作。但是，這種調(diào)度策略沒(méi)有考慮到其他設(shè)備上可能被某個(gè)操作阻塞的下游操作，并且經(jīng)常導(dǎo)致調(diào)度未充分利用的設(shè)備。為了找到可以跟蹤此類跨設(shè)備依賴關(guān)系的調(diào)度，我們的方法使用基于優(yōu)先級(jí)的調(diào)度算法，該算法根據(jù)每個(gè)操作的優(yōu)先級(jí)調(diào)度就緒隊(duì)列中的操作。類似于設(shè)備放置，操作調(diào)度可以被表述為學(xué)習(xí)策略的問(wèn)題，該策略為圖中的每個(gè)節(jié)點(diǎn)分配優(yōu)先級(jí)，以根據(jù)運(yùn)行時(shí)間最大化獎(jiǎng)勵(lì)。

第三個(gè)優(yōu)化任務(wù)是操作融合。為簡(jiǎn)潔起見(jiàn)，我們?cè)谶@里省略了對(duì)這個(gè)問(wèn)題的詳細(xì)討論，而只是注意類似于基于優(yōu)先級(jí)的調(diào)度，操作融合也可以使用基于優(yōu)先級(jí)的算法來(lái)決定融合哪些節(jié)點(diǎn)。在這種情況下，策略網(wǎng)絡(luò)的目標(biāo)再次是為圖中的每個(gè)節(jié)點(diǎn)分配優(yōu)先級(jí)。

最后，重要的是要認(rèn)識(shí)到在三個(gè)優(yōu)化問(wèn)題中的每一個(gè)中做出的決定都會(huì)影響其他問(wèn)題的最佳決策。例如，將兩個(gè)節(jié)點(diǎn)放在兩個(gè)不同的設(shè)備上會(huì)有效地禁用融合并引入可能影響調(diào)度的通信延遲。

RL 策略網(wǎng)絡(luò)架構(gòu)

我們的研究提出了 GO，這是一種深度 RL 框架，可適用于單獨(dú)或聯(lián)合解決上述每個(gè)優(yōu)化問(wèn)題。所提議的架構(gòu)有三個(gè)關(guān)鍵方面：

首先，我們使用圖神經(jīng)網(wǎng)絡(luò)（特別是GraphSAGE）來(lái)捕獲計(jì)算圖中編碼的拓?fù)湫畔ⅰraphSAGE 的歸納網(wǎng)絡(luò)利用節(jié)點(diǎn)屬性信息泛化到以前看不見(jiàn)的圖，這使得對(duì)看不見(jiàn)的數(shù)據(jù)做出決策而不會(huì)產(chǎn)生大量的訓(xùn)練成本。

其次，許多模型的計(jì)算圖通常包含超過(guò) 1 萬(wàn)個(gè)節(jié)點(diǎn)。在如此大規(guī)模上有效地解決優(yōu)化問(wèn)題需要網(wǎng)絡(luò)能夠捕獲節(jié)點(diǎn)之間的遠(yuǎn)程依賴關(guān)系。GO 的架構(gòu)包括一個(gè)可擴(kuò)展的注意力網(wǎng)絡(luò)，該網(wǎng)絡(luò)使用段級(jí)遞歸來(lái)捕獲這種遠(yuǎn)程節(jié)點(diǎn)依賴關(guān)系。

第三，ML 編譯器需要解決來(lái)自不同應(yīng)用領(lǐng)域的各種圖形的優(yōu)化問(wèn)題。用異構(gòu)圖訓(xùn)練共享策略網(wǎng)絡(luò)的幼稚策略不太可能捕捉到特定類別圖的特質(zhì)。為了克服這個(gè)問(wèn)題，GO 使用了一種特征調(diào)制機(jī)制，該機(jī)制允許網(wǎng)絡(luò)專門針對(duì)特定的圖類型而無(wú)需增加參數(shù)數(shù)量。

為了共同解決多個(gè)相關(guān)優(yōu)化任務(wù)，GO 能夠?yàn)槊總€(gè)任務(wù)添加額外的循環(huán)注意層，并在不同任務(wù)之間共享參數(shù)。具有剩余動(dòng)作連接的循環(huán)注意層能夠跟蹤任務(wù)間的依賴關(guān)系。

結(jié)果

接下來(lái)，我們展示了基于真實(shí)硬件測(cè)量的設(shè)備放置任務(wù)的單任務(wù)加速評(píng)估結(jié)果，對(duì)具有不同 GO 變體的不可見(jiàn)圖的泛化，以及聯(lián)合優(yōu)化操作融合、設(shè)備放置和調(diào)度的多任務(wù)性能。

加速：

為了評(píng)估該架構(gòu)的性能，我們將 GO 應(yīng)用于基于真實(shí)硬件評(píng)估的設(shè)備放置問(wèn)題，我們首先在每個(gè)工作負(fù)載上分別訓(xùn)練模型。這種稱為GO-one 的方法始終優(yōu)于專家手動(dòng)放置 (HP)、TensorFlow METIS放置和分層設(shè)備放置(HDP)——當(dāng)前最先進(jìn)的基于強(qiáng)化學(xué)習(xí)的設(shè)備放置。重要的是，憑借高效的端到端單次布局，GO-one 的布局網(wǎng)絡(luò)在 HDP 上的收斂時(shí)間加快了 15 倍。

我們的實(shí)證結(jié)果表明，GO-one始終優(yōu)于專家布局、TensorFlow METIS 布局和分層設(shè)備布局(HDP)。由于 GO 的設(shè)計(jì)方式可以擴(kuò)展到由超過(guò) 80,000 個(gè)節(jié)點(diǎn)組成的超大圖，例如 8 層Google 神經(jīng)機(jī)器翻譯(GNMT) 模型，因此它優(yōu)于以前的方法，包括 HDP、REGAL和Placeto。GO 為 GNMT 等大型圖實(shí)現(xiàn)了優(yōu)化的圖運(yùn)行時(shí)間，分別比 HP 和 HDP 快 21.7% 和 36.5%。總體而言，GO-one 平均減少了 20.5% 和 18.2% 的運(yùn)行時(shí)間分別與 HP 和 HDP 相比，包含 14 個(gè)不同的圖形集。重要的是，憑借高效的端到端單次布局，GO-one的布局網(wǎng)絡(luò)在 HDP 上的收斂時(shí)間加快了15 倍。

泛化：

GO 使用離線預(yù)訓(xùn)練對(duì)未見(jiàn)過(guò)的圖進(jìn)行泛化，然后對(duì)未見(jiàn)過(guò)的圖進(jìn)行微調(diào)。在預(yù)訓(xùn)練期間，我們?cè)谟?xùn)練集中的圖的異構(gòu)子集上訓(xùn)練 GO。在切換到下一個(gè)之前，我們?cè)诿總€(gè)這樣的批次圖上訓(xùn)練 GO 1000 步。然后在保留圖上對(duì)這個(gè)預(yù)訓(xùn)練模型進(jìn)行微調(diào)（GO-generalization+finetune），步驟少于 50 步，通常不到一分鐘。GO-generalization+finetune for hold-out graphs 在所有數(shù)據(jù)集上始終優(yōu)于專家放置和 HDP，并且平均匹配GO-one。

我們還直接在預(yù)訓(xùn)練模型上運(yùn)行推理，而不對(duì)目標(biāo)保持圖進(jìn)行任何微調(diào)，并將其命名為GO-generalization-zeroshot。這種未調(diào)整模型的性能僅比GO-generalization+finetune差一點(diǎn)，而略好于專家布局和 HDP。這表明圖嵌入和學(xué)習(xí)到的策略都有效地轉(zhuǎn)移，允許模型泛化到看不見(jiàn)的數(shù)據(jù)。

協(xié)同優(yōu)化布局、調(diào)度和融合（pl+sch+fu）：

同時(shí)優(yōu)化布局、調(diào)度和融合與未優(yōu)化的單 GPU 情況相比提供 30%-73% 的加速，與 TensorFlow 相比提供 33%-60% 的加速默認(rèn)放置、調(diào)度和融合。與單獨(dú)優(yōu)化每個(gè)任務(wù)相比，多任務(wù) GO ( pl+sch+fu ) 比單任務(wù) GO ( p | sch | fu )——一次優(yōu)化所有任務(wù)——平均高 7.8%。此外，對(duì)于所有工作負(fù)載，與優(yōu)化任意兩個(gè)任務(wù)并為第三個(gè)任務(wù)使用默認(rèn)策略相比，共同優(yōu)化所有三個(gè)任務(wù)可提供更快的運(yùn)行時(shí)間。

pl+sch：多任務(wù) GO 使用默認(rèn)融合協(xié)同優(yōu)化布局和調(diào)度。sch+fu：多任務(wù) GO 協(xié)同優(yōu)化調(diào)度和融合與人類安置。pl | sch | fu：GO 分別優(yōu)化布局、調(diào)度和融合。pl+sch+fu：多任務(wù) GO 協(xié)同優(yōu)化布局、調(diào)度和融合。

結(jié)論

硬件加速器日益復(fù)雜和多樣化，使得開(kāi)發(fā)健壯且適應(yīng)性強(qiáng)的 ML 框架變得既繁瑣又耗時(shí)，通常需要數(shù)百名工程師的多年努力。在本文中，我們證明了此類框架中的許多優(yōu)化問(wèn)題都可以使用精心設(shè)計(jì)的學(xué)習(xí)方法有效且最優(yōu)地解決。