[[425806]]

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

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

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

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

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

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

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

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

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

RL 策略網(wǎng)絡架構

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

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

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

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

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

結果

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

加速：

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

我們的實證結果表明，GO-one始終優(yōu)于專家布局、TensorFlow METIS 布局和分層設備布局(HDP)。由于 GO 的設計方式可以擴展到由超過 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%?？傮w而言，GO-one 平均減少了 20.5% 和 18.2% 的運行時間分別與 HP 和 HDP 相比，包含 14 個不同的圖形集。重要的是，憑借高效的端到端單次布局，GO-one的布局網(wǎng)絡在 HDP 上的收斂時間加快了15 倍。

泛化：

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

我們還直接在預訓練模型上運行推理，而不對目標保持圖進行任何微調(diào)，并將其命名為GO-generalization-zeroshot。這種未調(diào)整模型的性能僅比GO-generalization+finetune差一點，而略好于專家布局和 HDP。這表明圖嵌入和學習到的策略都有效地轉移，允許模型泛化到看不見的數(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)化每個任務相比，多任務 GO ( pl+sch+fu ) 比單任務 GO ( p | sch | fu )——一次優(yōu)化所有任務——平均高 7.8%。此外，對于所有工作負載，與優(yōu)化任意兩個任務并為第三個任務使用默認策略相比，共同優(yōu)化所有三個任務可提供更快的運行時間。

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

結論

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