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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

結(jié)果

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

加速：

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

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

泛化：

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

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

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

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

pl+sch：多任務(wù) GO 使用默認融合協(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ù)雜和多樣化，使得開發(fā)健壯且適應(yīng)性強的 ML 框架變得既繁瑣又耗時，通常需要數(shù)百名工程師的多年努力。在本文中，我們證明了此類框架中的許多優(yōu)化問題都可以使用精心設(shè)計的學(xué)習(xí)方法有效且最優(yōu)地解決。

更新說明：優(yōu)先更新微信公眾號“雨夜的博客”，后更新博客，之后才會陸續(xù)分發(fā)到各個平臺，如果先提前了解更多，請關(guān)注微信公眾號“雨夜的博客”。