深度學習編譯器模型訓練負載均衡優化方法

王 麗，高 開，趙雅倩，李仁剛，曹 芳，郭振華

浪潮電子信息產業股份有限公司 高效能服務器與存儲技術國家重點實驗室，濟南 250000

對于計算密集型的人工智能（artificial intelligence，AI）訓練任務，隨著AI算法模型參數從千萬級增至千億級，單個AI 硬件設備的算力和內存遠遠不能滿足深度學習模型訓練的需求，模型訓練并行化成為增強其應用時效性的迫切需求。近年來面向AI分布式訓練的硬件加速器層出不窮，如中央處理器（central processing unit，CPU）、圖形處理器（graphic processing unit，GPU）、現場可編程門陣列（field programmable gate array，FPGA）、AI 芯片等，為AI 算法模型并行訓練提供了硬件基礎。同時，國內外研究人員也在致力于推出適用于不同應用場景的深度學習框架，主要 有TensorFlow[1]、PyTorch[2]、MXNet[3]、PaddlePaddle[4]等。隨著深度學習框架和AI 硬件設備的持續多樣化，單純基于手工優化來實現AI 算法模型在深度學習框架和AI 硬件設備之間的高效映射，越來越容易出現性能瓶頸。

為了充分利用各種硬件資源，減輕手動優化每個硬件運行模型的負擔，研究人員試圖通過深度學習編譯器解決框架的靈活性和性能之間的矛盾[5]。深度學習編譯器的通用設計架構如圖1所示，其主要功能是基于編譯優化技術將不同深度學習框架描述的AI 算法模型在底層各種AI 硬件設備上生成有效的代碼實現，并根據模型規范和硬件體系結構完成對生成代碼的高度優化，解決了手工優化帶來的AI模型的性能和效率問題。

圖1 深度學習編譯器通用設計架構Fig.1 General design architecture of deep learning compiler

對于AI 模型并行訓練，研究人員需要將計算任務映射到多個AI 加速設備上并行計算，基于深度學習編譯器完成模型訓練任務到底層硬件設備自動高效映射的同時，如何實現訓練任務在多個加速器上的負載均衡，從而進一步提升訓練效率，已經成為深度學習編譯器面臨的難點問題。當前，業界針對大規模AI 并行訓練提出了各種并行優化方法，以提高模型訓練的效率。其中PipeMare[6]、Wavelet[7]、Chimera[8]等基于流水并行優化方法盡可能減少訓練過程中設備閑置時間，以提高硬件設備資源利用率；Taskflow[9]、Gavel[10]等通過設計不同的調度算法和策略，實現模型訓練任務在多GPU 上的高效分配。但是，上述研究考慮了整體訓練效率的優化，沒有分析模型訓練期間硬件設備之間的負載均衡問題，且主要面向深度學習框架進行，在當前的深度學習編譯器中尚沒有涉及。

本文致力于研究面向模型訓練的深度學習編譯器計算圖負載均衡優化方法。目前業界已經提出幾種主流深度學習編譯器，例如Google XLA[11]（accelerated linear algebra，加速線性代數編譯器）、Apache TVM[12]和Glow[13]等，主要面向AI 應用進行圖層級或算子級編譯優化。其中，Google XLA 旨在為TensorFlow 提供后端的靈活性支持，其集成在TensorFlow內部，基于TensorFlow 前端接口實現對模型訓練任務的支持，XLA 編譯器支持算子融合等圖層級優化，但對于流水調度、負載均衡等優化需要基于TensorFlow 的分布式訓練策略實現；TVM、Glow 主要面向推理應用場景進行算子級編譯優化，對于訓練任務的支持處于開發階段；nGraph[14]目前作為核心映射模塊集成到Intel 高性能推理引擎OpenVINO[15]中，轉向推理應用進行性能優化，不再支持AI 算法模型訓練任務，因此，當前主流深度學習編譯器均無法實現面向AI模型訓練任務的圖層級負載均衡優化。

針對以上問題，本文結合訓練任務的性能瓶頸特點，提出面向模型訓練的深度學習編譯器計算圖負載均衡優化方法，包括數據加載與模型訓練的負載均衡優化設計、計算圖分層調度優化設計以及計算圖分層自動流水優化設計，完成了模型訓練任務中CPU 與后端設備之間以及后端設備內部任務調度的負載均衡優化，從而提高了模型訓練任務從深度學習框架到AI硬件設備映射的計算性能和硬件設備資源利用率，并提升了系統整體能效。本文的貢獻主要有以下三點：

（1）提出數據加載與模型訓練的負載均衡優化設計方法，為數據加載提供最優的并行線程數，在CPU 負載最小的情況下，實現數據加載與模型訓練的負載均衡。同時，通過數據搬運與模型訓練之間的多線程流水并行控制，掩蓋計算設備等待數據搬運時間，提高模型訓練任務的整體性能。

（2）提出計算圖分層調度優化設計方法，通過啟發式分層查找算法自動尋找計算圖的分層和算子調度策略，得到計算圖的最優分層，實現編譯器中計算圖在后端設備上的負載均衡，并通過子圖并行調度，進一步提升了計算性能。

（3）提出計算圖分層自動流水優化設計方法，提出分層后計算子圖間的并行流水設計，解決訓練過程中反向計算和參數更新的依賴，利用二分查找算法自動優化微批次劃分，實現計算圖分層流水并行計算，進一步提高計算過程中AI設備的資源利用率。

1 背景介紹

深度學習編譯器前端優化是基于高層中間表示（intermediate representation，IR）的圖層級優化，雖然不同深度學習編譯器前端計算圖在高層IR 的數據表示和算子定義上有所不同，但獨立于硬件的前端優化均可以分為節點級、塊級和數據流級三個層次，每個層次的優化方法都采用深度學習特有的編譯優化技術，減少計算冗余，提高算法模型在圖層級的性能。

節點級優化包括節點消除和張量消除優化方法，編譯器根據計算圖判斷不參與計算的無用節點和張量，并進行刪除[5,12]。塊級優化包括代數簡化以及算子融合、算子分解等神經網絡相關優化方法，編譯器分析計算圖中各算子之間的基本運算關系，利用結合律、交換律等基本方式調整計算順序，簡化計算，執行算子融合，提升計算性能[16-17]。近期Rammer[18]、DNNFusion[19]、Unity[20]等工作將代數簡化、算子融合、并行優化等技術結合，聯合優化，加速了模型訓練。數據流級優化包括公共表達式消除、數據布局轉換、靜態內存分配等優化方法[21]，編譯器基于公共表達式消除、死代碼消除等數學優化盡可能減少重復和無用計算，并結合就地內存共享和標準內存共享的方式，執行靜態內存規劃優化，以實現內存緩沖區重用，Zhao 等[21]最新研究通過優化數據布局，消除通信過程中數據子集的相互移動，提升了模型訓練過程中通信效率。

針對上述編譯優化方法，各深度學習編譯器具有不同的實現方式。表1 統計了面向AI 模型訓練任務的深度學習編譯器前端優化方法，以及主流深度學習編譯器對各優化方法的支持。統計表明，現有深度學習編譯器采用基本一致的編譯前端優化方法，包括節點級優化、塊級優化和數據流控制優化，而對于AI 算法模型訓練過程中，需要考慮的計算圖負載均衡優化并未涉及。

表1 編譯器前端優化方法Table 1 Front-end optimization methods of compiler

為了實現深度學習編譯器對模型訓練任務的支持，進一步提升模型訓練任務在編譯器中的計算性能，同時降低系統整體能耗，本文提出了面向模型訓練任務的計算圖負載均衡優化方法。該方法從系統全局的角度對計算圖進行分層劃分，并且實現分層劃分后計算子圖間的流水并行調度，使得整個編譯優化過程處于負載均衡狀態，以提高AI 硬件設備的資源利用率，主要包括數據加載與模型訓練負載均衡優化、計算圖分層調度優化以及計算圖分層自動流水優化。

數據加載與模型訓練負載均衡優化主要是解決CPU上數據讀取任務和后端AI硬件設備上模型訓練任務之間的負載不均衡問題。一般通過開啟多線程流水控制，掩蓋數據讀取的時間，以提高模型訓練的性能和效率。AI 算法模型訓練過程中，數據搬運的復雜度顯著提高，容易帶來計算開銷無法覆蓋傳輸開銷的問題，因此數據加載也是影響整體訓練性能的關鍵因素。OneFlow[22]在計算圖中顯式地表達數據搬運，將模型訓練任務過程中的數據讀取、數據傳輸等輸入輸出（input output，IO）操作封裝成與計算圖同一標準的計算子圖，實現數據搬運與模型訓練之間的自動流水并行控制，為深度學習框架在AI 模型訓練中性能優化開啟了新的思路；TensorFlow[1]通過建立文件名隊列和任務隊列，采用多線程流水控制的方式從文件名隊列中讀取圖片數據，進行數據解析，并將解析后的數據加入任務隊列中為訓練任務提供輸入數據，實現了數據加載與訓練任務的并行流水。但是上述實現均沒有考慮數據加載過程中的多線程數量優化問題，僅通過開啟固定單線程或者最大線程數的方式實現數據加載與模型訓練的流水并行，并沒有解決數據加載與模型訓練的負載均衡問題，容易帶來AI 設備資源的利用不均衡，降低AI算法模型的訓練效率。

計算圖分層調度優化以及分層自動流水優化屬于計算子圖在后端設備上調度的負載均衡優化。Xiao、Woo 等[23-24]提出面向分布式訓練的GPU 并行調度系統，旨在提高大規模分布式集群中模型訓練的GPU 設備資源利用率，但以上研究主要針對數據并行訓練任務；Gpipe[25]、DAPPLE[26]、Chimera[8]等工作提出細粒度流水并行優化設計方案，盡可能減少流水并行訓練過程中GPU 設備空閑時間，提高硬件設備資源利用率以及模型訓練性能。但在當前的深度學習編譯器系統中，調度優化主要涉及內存調度以及硬件指令相關的循環調度，在編譯前端對計算圖進行模型并行訓練相關的計算圖分層和流水并行調度優化在當前主流深度學習編譯器中尚未涉及。

2 計算圖負載均衡優化

為了充分利用深度學習編譯器后端各種AI硬件資源，滿足計算密集型模型訓練任務的算力需求，研究人員需要采用大規模AI硬件協同計算提高模型訓練速率，實現模型訓練任務到底層硬件設備的高效映射。目前單一的局部優化算法能夠帶來局部性能提升，但系統整體的計算負載難以達到均衡狀態，導致AI 設備加速器計算效率低下，降低了AI 計算平臺的硬件資源利用率。

本文從編譯器系統全局負載均衡的角度出發，針對目前主流深度學習編譯器有限支持模型訓練任務，以及模型訓練過程中設備之間和設備內部資源利用不均衡等特點，搭建了面向AI 模型訓練的深度學習編譯器框架，并提出了計算圖負載均衡優化策略，主要包括以下兩個方面：第一，通過數據加載與模型訓練的負載均衡優化技術提供數據讀取階段最優并行線程數，實現CPU 上數據加載任務和后端設備上模型計算任務之間的負載均衡；第二，通過計算圖的分層調度優化以及多級分層后的自動流水優化技術實現計算圖在后端設備內部調度的負載均衡。

2.1 面向模型訓練的深度學習編譯器

本文首先搭建自研深度學習編譯器Mofang。Mofang 是基于開源編譯器nGraph[14]開發的一套面向模型訓練任務的深度學習編譯器框架，主要實現模型訓練任務從編譯器前端到編譯器后端AI硬件設備的解耦化映射，其前端支持Tensorflow、PaddlePaddle、ONNX 格式訓練模型的導入，后端支持Intel CPU、NVIDIA V100 GPU以及F10a FPGA加速設備。

Mofang 充分利用nGraph 提供的框架bridge（轉換層）解耦機制，通過bridge 層將不同框架的AI 算法模型輸入轉換至使用其高層中間表示語言描述的統一格式（High-Level IR），并對其進行硬件無關的優化，實現了計算圖重構。由于nGraph 由同時支持訓練和推理的軟件生態，改變為只支持推理任務，且移除了分布式通信接口，本文基于nGraph 添加CPU、GPU、FPGA等AI硬件加速設備的分布式通信算子接口，并集成OpenMPI、NCCL等優化版的通信算子庫，實現了深度學習編譯器對AI 分布式訓練任務的支持，Mofang整體實現架構如圖2所示。

圖2 Mofang深度學習編譯器整體架構Fig.2 Overall architecture of Mofang deep learning compiler

為了提升模型訓練的性能以及后端設備資源利用率，本文在Mofang 框架中實現了數據加載與模型訓練的負載均衡、計算圖分層調度以及分層后自動流水優化等編譯優化技術，其中數據加載與模型訓練的負載均衡優化在Mofang 自定義前端實現，計算圖分層以及分層后自動流水優化基于Mofang高層中間表示實現。

2.2 數據加載與模型訓練的負載均衡

本節首先分析了目前深度學習框架中數據加載和模型訓練任務之間的調度流程。為了解決數據加載的性能瓶頸問題，現有框架通常在CPU 上開啟多線程并行讀取數據，以提高數據加載的效率，使其能夠滿足模型訓練過程的數據讀取需求，數據加載與模型訓練兩個階段通過建立數據隊列緩存的方式實現流水并行執行，主流深度學習框架TensorFlow 中數據加載與模型訓練流水并行的實現流程如圖3 所示。TensorFlow 數據加載過程首先建立文件名隊列和數據隊列，由左至右分為文件名獲取、寫入文件名隊列、按照文件名讀取數據、解析數據、將數據寫入數據隊列五個階段。其中按照文件名讀取數據和解析數據的過程采用啟動固定數量的CPU 線程的方式并行執行，以保證數據讀取和解析的速率能夠滿足訓練任務獲取數據的需求，訓練模塊只需要按照流程從數據隊列中讀取解析后的數據，執行計算過程，減少了訓練階段等待數據加載的時間消耗。

圖3 TensorFlow深度學習框架中的數據加載Fig.3 Data loading in TensorFlow deep learning framework

對于不同的模型訓練任務，模型訓練性能存在較大的差異，單純開啟固定線程數進行數據加載的方式容易導致數據加載與模型訓練任務負載不均衡。例如，當數據加載速度過快，超過模型訓練模塊的數據讀取需求時，會導致數據隊列緩存溢出；而數據加載過慢，不能滿足模型訓練模塊的數據讀取需求時，會出現訓練過程等待數據的問題，降低了訓練性能。

本文測試了在CPU 上開啟不同數量的線程時數據讀取模塊的性能，統計結果如表2所示。線程開啟數量與讀取性能在一定范圍內呈現線性增長的趨勢，但是當開啟最大CPU線程數64時，數據讀取解析性能呈現下降趨勢。這是因為此時開啟最大的線程數，會占用更多CPU 硬件資源，影響系統整體運行，導致數據讀取性能下降。

表2 線程數量對數據讀取性能的影響Table 2 Impact of the number of threads on data read performance

TensorFlow 等深度學習框架通過人工經驗去定義各個階段的線程數量來實現流水，最常用的方式是針對性能瓶頸模塊開啟最大線程數來提高整體性能。根據表2統計結果，開啟最大線程數的方式容易導致占用過多CPU 資源，存在大量的資源浪費，而且會影響其他程序運行。因此在實際模型訓練任務中，需要根據不同的數據批處理大?。╞atch_size）、AI硬件設備類型以及AI硬件設備數量等對性能瓶頸階段設置不同的線程數量，才能夠使整個系統的流水處于負載均衡狀態，提高硬件資源利用率。通常情況下，訓練時batch_size越大，使用的硬件設備數量越多，階段3 模型訓練的性能就會越高，因此在CPU 上需要開啟的處理數據讀取的線程數量越多。

為了更直觀地分析數據加載與模型訓練過程的整體性能瓶頸，本文對整個過程進行多階段劃分，如圖4 所示。數據加載與模型訓練過程主要分為三個階段，其中階段1 為文件名讀取，階段2 為數據的讀取解析，階段3為模型訓練。目前主流深度學習編譯器中，階段1 和階段2 以多線程的方式運行在CPU 處理器上，第三階段模型訓練過程經過前端和后端優化后運行在編譯后端AI 硬件設備上。在AI 模型訓練任務中，數據讀取階段的耗時最多，容易成為整個流水過程的性能瓶頸；而不同的優化方法以及AI 硬件設備對階段3模型訓練的性能影響較大，容易導致三個階段的負載不均衡。

圖4 數據加載與模型訓練階段劃分Fig.4 Data loading and model training phase division

針對上述問題，本文針對深度學習編譯器的自定義前端提出一種自動調整線程數量的方法，進一步優化數據讀取解析階段的線程數量，從而實現三個階段的負載均衡。

本文提出反向查找的方法自動尋找數據讀取階段線程的最優配置，即找出整個系統達到最佳性能時數據讀取階段需要開啟的最少線程數量。假設當數據批處理大小為batch_size時，第三階段經過特定后端優化之后的預估時間為T3，第二階段每個線程處理一個數據的預估時間為T2，第一階段的預估時間為T1。假設CPU 的最大線程數量MAX_thread，線程加速比為?n，如果第二階段在CPU上開啟了n個線程，在流水執行狀態下，那么整個系統的消耗時間就是執行時間最長的階段的時間，如式（1）。因此，本文研究如何開啟最少的線程數，即n盡可能得小，使得整個系統的時間T最小。系統整體時間T的計算公式如下所示：

由于在batch_size固定時，對于同一種AI加速器，T1和T3是固定的，對于式（1），變形后有：

要使得T最小，則有：

在滿足式（4）下選取n的最小正整數即是最優解。

通過上述方法選取的數據讀取階段的線程數，能夠在滿足第三階段模型訓練數據讀取需求的同時，盡可能地節省CPU的資源占用，使得CPU在負載最小的情況下，數據加載與模型訓練兩個階段的任務達到負載均衡的狀態，即提高了CPU 的資源利用率，減少CPU資源浪費，進一步節省了系統的能耗。

2.3 計算圖分層調度優化

深度學習編譯器中模型訓練通常以計算圖的方式在后端AI 硬件設備上進行調度，由于計算圖參數逐漸增大，算子間依賴關系愈加復雜，容易出現算子調度不均衡，模型訓練過程中后端設備資源利用率波動較大的問題[27]。傳統的順序調度算法、貪心調度算法以及算子的變形融合優化（TASO（tensor algebra superoptimizer for deep learning）[28]）技術都存在算子調度不均衡的問題。本文提出了計算圖的分層調度優化算法重新規劃算子間依賴關系，使得計算圖每個階段算子的調度達到負載均衡狀態。

本文以InceptionV3 的部分網絡結構（圖5）為例說明不同的調度算法對算子調度順序以及設備利用率的影響。該網絡結構包括五個卷積操作，每個卷積的通道數和卷積核大小不盡相同。圖6 展示了深度學習編譯器常用的卷積算子的變形融合優化，首先通過補零擴展卷積核的方式把1×1 的卷積擴展為3×3的卷積如圖6（a）所示，再通過簡單堆疊的方式把具有同樣卷積核的Conv(a)、Conv(b)和Conv(d)三個卷積合并為一個卷積Conv(a)(b)(d)。合并后的網絡結構如圖6（b）所示，卷積合并優化之后需要增加一個split 算子，用來處理合并算子的輸出，優化之后的計算圖在執行過程中減少了兩次內核的啟動，能夠獲得相應的性能提升。

圖5 InceptionV3模塊Fig.5 InceptionV3 module

圖6 深度學習編譯器卷積算子合并優化Fig.6 Deep learning compiler convolution operator merging optimization

本節首先基于傳統深度學習框架使用的貪心算法進行算子變形融合優化前后，AI 硬件設備的資源利用率分析?；谒惴ㄔ?，貪心算法每次將最大限度地調度可調度的算子，而不考慮算子的負載均衡問題。在原始的網絡結構中（圖5），Conv(a)、Conv(b)、Conv(d)同時獲取輸入數據，不存在依賴關系，因此在第一階段，貪心算法會同時調度Conv(a)、Conv(b)和Conv(d)三個算子，此時三個算子的調度存在激烈的資源競爭，GPU 的資源利用率會達到頂峰；而在第二階段只調度Conv(c)一個算子，由于計算量小，GPU的利用率大幅降低，導致GPU 硬件資源利用不足，存在資源浪費。

經過卷積算子變形融合技術優化后，貪心算法在第一階段會調度合并后的算子Conv(abd)，由于該算子通道數量合并為三個卷積的通道數量之和，在調度該算子時同樣會占用大量的GPU 內存資源，第二階段調度的算子也只有一個，同樣會由于Conv(c)算子的計算量小，導致資源利用率不高。因此，算子變形融合優化不能解決計算圖前后算子的調度不均衡問題。

基于以上問題，本文提出算子分層調度優化策略，該策略重新定義算子的調度依賴關系，使整個計算圖的調度達到均衡狀態，如圖7所示。首先根據計算圖卷積算子的數據依賴關系，對計算圖執行分層操作，算子Conv(a)和算子Conv(d)為第一層，算子Conv(b)和算子Conv(c)為第二層；然后根據算子卷積核和通道數等參數特點，將第一層中的算子Conv(a)和算子Conv(d)進行算子變形融合優化，形成合并算子Conv(a)(d)。為了提高第二階段AI 設備的資源利用率，將第二層的算子Conv(b)和算子Conv(c)執行并行計算，計算圖第一層和第二層由于存在數據依賴關系，執行串行調度。

圖7 算子分層調度優化策略Fig.7 Operator hierarchical scheduling optimization strategy

為了保證層與層之間的串行調度，本文在層之間增加新的result節點，插入result節點后的整個計算圖如圖8所示，該節點負責上一層計算結果的匯總和下一層輸入數據的分發。

圖8 插入result節點增加層間依賴后的計算圖Fig.8 Calculation graph after inserting result node to add interlayer dependencies

為了解決復雜計算圖中，計算圖高效分層劃分問題，本文提出了一種啟發式分層查找算法，自動計算查找計算圖的最優分層策略。假設計算圖G的所有算子節點集合為V，從V的頂點開始按照算子依賴遍歷整個計算圖，得到一個子集U，子集U將V分成U和V-U兩部分。計算圖的最優調度延遲用cost來表示，那么狀態轉移公式為：

即V的最優延遲是所有可能分層U的延遲和剩余頂點集V-U的最優延遲之和的最小值。

對于每一種可能的分層U，給出一個啟發式評價函數來評估分層的好壞。該啟發式函數如式（6）所示，通過計算分層后的兩層的計算量差別來評估分層結果的質量。

其中，com(U)為分層U的計算量。

具體來說，整個分層包括以下實施步驟：

（1）枚舉計算圖V所有可能的算子分層{U0,U1,…,Um}。

（2）根據啟發評價函數τ計算所有可能分層的評估指標，按照指標由小到大排序，設定閾值φ，若對于Ui,τi＞φ,則刪除對應的分層Ui，得到新的集合{U i,Ui+1,…,Uj}。

（3）通過仿真模擬預估出每層的耗時，劃分后的每層都有兩種調度策略：貪心調度和順序調度。根據預估結果選擇耗時較小的調度策略。預估每層的時間消耗時，將每層的輸入I、層的圖結構G、調度策略S作為輸入參數，該層的耗時T為在某種調度策略下完成整個層計算的時間。

（4）根據分層后的模擬預估數據，選擇出兩層耗時之和最小的分層{Uk}，將該分層結果{Uk}作為第一次分層的最優結果。

（5）對于（4）得到的最優分層[Uk,V-Uk]。將Uk和V-Uk分別賦值給V，重復（1）（2）（3）（4）。

（6）若分層之后的所有可能分層的耗時之和都大于未分層前耗時，即?Ui∈{Ui,Ui+1,…,Uj},cost(V-Ui)+cost(Ui)≥cost(V)，則分層結束。

圖9是啟發式分層的流程示意圖，首先從第一個節點a開始廣度優先遍歷，遍歷過程中需要保證圖的連續性，禁止跳躍節點。枚舉出所有可能的分層U∈{{a},{a,b},{a,c},…,{a,b,c}}。通過啟發式評價函數得到不同分層的指標{τ1,τ2,…,τm}，通過τi＞φ條件剔除評價較差的分層得到子集{τ2,τ3}。對子集中所有可能的分層結果，通過仿真模擬預估出每種情況下的調度耗時cost(U)和cost(V-U)。最后選擇耗時之和最小的分層，即完成第一輪分層。完成第一輪的最優分層之后，分別對U和V-U重復上述過程。直到下一次所有可能的分層預估耗時之和都不小于未分層之前的耗時，此時結束整個分層優化，詳細流程見算法1。

圖9 啟發式分層查找算法Fig.9 Heuristic hierarchical search algorithm

算法1啟發式分層查找算法

啟發式分層查找算法的主要計算量集中在尋找計算圖的所有可能分層上，本文進一步分析評估啟發式查找算法的時間復雜度，以更直接地統計計算圖可能分層的數量。本文給出以下定義：計算圖的寬度w,即在圖中可以找到最多w個算子同時并行執行；計算圖的長度h，即在圖中能夠找到最多h個算子必須串行執行。

通過設置合適的啟發式評價函數的閾值，可以降低啟發式搜索算法的復雜度，假設通過啟發式函數剔除一半的可能分層，則算法的時間復雜度降低為O((N/w+1)w)。

為了更好地評估啟發式分層查找算法自動尋找計算圖分層調度優化策略的執行效率，本節分析了典型神經網絡模型的時間復雜度和經過啟發式算法優化之后的時間復雜度，結果如表3所示。數據分析顯示啟發式分層查找算法可以大大降低搜索的復雜度。

表3 不同模型的時間復雜度Table 3 Time complexity of different models

2.4 計算圖分層自動流水優化

計算圖分層調度策略可以有效地優化算子的調度順序，減緩后端設備利用率的波動。深度學習編譯器中計算圖在經過分層調度優化之后，每層的計算量分布相對均衡，模型訓練過程中后端設備的利用率趨于平穩，但是遠遠達不到100%。

為了進一步提高編譯器后端設備的資源利用率，本文基于分層調度優化提出了多級自動流水優化策略。通過在層與層之間添加多個Result 數據緩存節點，實現層與層間的流水計算，如圖10所示。首先在層與層中間添加兩個Result 節點Result0 和Result1。模型訓練過程中，在T0階段，第一層讀取第一批次數據（Batch_0）進行計算，并把計算好的結果傳遞給Result0 節點；在T1 階段，第二層讀取Result0的數據進行第二層的計算，同時第一層讀取第二批次輸入數據（Batch_1）進行計算，并將計算結果傳遞給Result1，為下一階段的第二層計算提供輸入數據；第二層只需交替從Result0 和Result1 節點讀取輸入數據。因此在計算圖劃分為兩層的情況下，從T1 階段開始，計算圖第一層和第二層即可實現層間流水計算，后端設備需要同時計算第一層和第二層，進一步提高了硬件設備的資源利用率。

當計算圖在經過分層調度優化之后被劃分的層數較多時，如果直接進行流水并行，會帶來設備的超負載問題，如圖11 所示。計算圖被劃分為6 層，且每一層單獨運行時設備的利用率分別為：0.5、0.4、0.3、0.4、0.5 和0.4，如圖11（a）所示，在沒有流水的情況下，各層依次執行計算任務，設備的利用率十分平穩。當6層的計算實現完全流水時，設備資源理論利用率理論上應該達到250%，如圖11（b）所示。此時設備資源已經不能滿足計算需求。

圖11 計算圖分層流水效果Fig.11 Pipeline result after layering of calculation graph

本文設計多級自動流水策略，通過重新劃分流水級解決設備的超負荷問題，如圖12 所示。首先將計算圖的分層劃分為多級流水，每一級流水包含計算圖的多個分層，其中流水級內的分層串行執行，每級流水中設備的利用率為該級流水所包含的分層在單獨計算時設備利用率的最大值。達到完全流水狀態時，設備性能為所有流水級的性能之和，沒有超過100%。

圖12 多層次合并分級流水效果Fig.12 Pipeline result after multi-layer merging and grading

本文基于Gpipe[25]微批次劃分的方法和多級流水策略，在編譯器中實現AI模型訓練流水并行優化，利用多級流水策略確定了流水級之后，設置合理的微批次劃分來構建流水。

本文通過建立流水級長度和微批次數量之間的關系對Gpipe 中的微批次劃分方法進行優化，得到最優的微批次劃分策略。當計算圖的流水級較少時，流水建立過程耗時較短，如果劃分的微批次數量太多，則會導致劃分后的微批次（micro_batch）太小影響性能。當計算圖的流水級較多時，如果微批次數量較小，則無法充滿整個流水線，因此需要劃分足夠的微批次數量來掩蓋流水建立的時間消耗。由于每次訓練時總的批次大小是固定的，微批次數量越多，微批次數據就越少，設備的利用率就會變低。因此在微批次劃分時，在保證流水建立的情況下，要盡可能地增大微批次的大小，即選取滿足流水建立的最小的微批次數量。

假設流水建立時間占比為Ppipe(batch_num)，設備利用率損耗為Dloss(batch_num),微批次數量為batch_num，則batch_num需要滿足式（7）：

滿足式（7）的batch_num能夠使設備間的流水效果和設備的利用率達到一個均衡的狀態，進而整個系統的性能達到最優狀態。假設總的處理批量大小為batch_size，設備最大處理批次為batch_device，流水線長度為pipe_num。流水線建立的時間占比為：

式（7）～（9）聯合計算得到batch_num，劃分后的微批次大小可以通過式（10）計算得到：

最后為了使計算得到的micro_batch為整數，需要對每個micro_batch進行±1微調。

本文通過微批次二分查找法得到batch_num的最優解。batch_num的求解空間為大于1，小于batch_size的正整數，即{1,2,…,batch_size-1,batch_size}，通過算法2 在解空間查找得到最優解batch_num。假設模型訓練每次處理的總批量batch_size為256，GPU 處理的最大的batch_size為64，對于不同的模型計算得到最優的微批次的劃分策略，如表4所示。

表4 不同模型的分層優化和微批次劃分Table 4 Hierarchical optimization and micro-batch partitioning for different models

算法2微批次二分查找法

3 實驗

針對AI 模型訓練應用任務調度負載均衡，本文提出了三種計算圖負載均衡優化方法，在深度學習網絡模型訓練過程實驗中得到了有效的驗證。實驗環境采用AGX-2 GPU服務器，每臺服務器上配置了16塊TeslaV100板卡。其他硬件環境配置如表5所示。

表5 硬件環境配置Table 5 Hardware environment configuration

軟件環境方面，本文采用2.1 節介紹的自研深度學習編譯器Mofang進行計算圖負載均衡優化整體實驗及性能測試。

3.1 線程數對數據加載的影響

為了更直觀地對比數據加載與模型訓練的負載均衡優化對訓練性能的影響，本節首先基于Image-Net2012 數據集測試了Resnet50 和Resnet18 模型訓練時的數據加載、模型計算以及整體訓練的性能，同時測試了開啟不同數量的數據加載線程對訓練性能的影響。

圖13 顯示了數據加載線程數量對Resnet50 模型訓練性能的影響，結果顯示數據加載模塊開啟線程數量的多少不僅影響模型訓練的性能，還會影響CPU 的能耗。對于Resnet50 模型，模型計算操作運行在GPU 上，性能穩定且與數據加載線程數量無關；對于數據加載和整體訓練過程，當線程數開啟到16之前，數據加載和模型整體訓練性能隨著線程開啟數量的增加而線性增加，這是因為此時數據加載依然是整體訓練的性能瓶頸，因此增加線程數提高數據加載速率能夠提升整體訓練的性能；當線程數開啟到16 時，數據加載的性能已經足夠滿足模型計算的需求，繼續增加開啟的線程數，整體訓練性能不再提升，處于平穩狀態，造成了CPU資源浪費。

圖13 Resnet50訓練時線程數量對性能的影響Fig.13 Effect of the number of threads on performance of Resnet50 training

圖14 顯示了線程數量對Resnet18 模型訓練性能的影響，由于模型參數量較小，模型計算性能較高，無論開啟多少線程數量，數據加載的性能都無法滿足模型計算的性能需求，同時可以看到當數據加載開啟最大線程數量時，數據加載的性能會降低，這是因為開啟的線程數量過多，線程間存在資源競爭。此時本文啟發式查找算法自動尋找的最佳線程數量為48，既能保證數據加載的性能最好，又可以降低CPU的利用率。

圖14 Resnet18訓練時線程數量對性能的影響Fig.14 Effect of the number of threads on performance of Resnet18 training

3.2 數據加載與模型訓練負載均衡優化

為了驗證數據加載與模型訓練負載均衡優化技術的可擴展性，本文采用Resnet50網絡模型和ImageNet 2012訓練數據集在多GPU板卡上進行了分布式訓練的驗證。本實驗分別測試了不同Batch_size 和不同GPU 板卡數量下數據加載與模型訓練負載均衡優化技術對模型分布式訓練性能的影響。

圖15 對比了不同GPU 板卡數量下加入負載均衡優化技術的Mofang 框架、TensorFlow 框架以及TensorFlow經nGraph編譯器映射后（TF-ngraph）的分布式訓練性能。結果顯示，隨著GPU 板卡數量的增加，三種框架中模型的分布式訓練的性能均呈現線性增加的趨勢，而本文加入負載均衡優化技術的Mofang 框架性能與TensorFlow 和TF-ngraph 框架相比，提升3%左右。

圖15 GPU數量對負載均衡優化性能的影響Fig.15 Effect of the number of GPU on performance of load balancing optimization

圖16 對比了不同batch_size 對加入負載均衡優化技術的Mofang 框架、TensorFlow 框架以及TFngraph 的分布式訓練性能的影響。實驗結果表明在不同的batch_size下負載均衡優化后的性能比另外兩種框架有2%到8%的提升，在batch_size=2 時性能的提升更加明顯，可以達到8%左右。這是由于batch_size較小，每次模型訓練需要加載的數據量小，在CPU 可以開啟足夠的線程來處理數據加載模塊；而隨著batch_size 增大，數據加載模塊需要開啟的CPU 線程就越多，超出CPU負載時會存在一定的性能損耗。

圖16 Batch大小對負載均衡優化性能的影響Fig.16 Effect of batch size on performance of load balancing optimization

3.3 計算圖分層調度優化

本節基于TensorFlow+Mofang 框架對計算圖的各種調度進行了實驗分析，計算圖經過不同的調度優化后在TeslaV100 GPU 上進行調度執行，不同的調度模式對應算子不同的調度順序。首先在TensorFlow 中構建計算圖，并人為增加算子調度依賴或者算子優化，在確定算子的調度順序之后分別在GPU平臺上執行優化后的計算圖，并統計各個階段的耗時和資源占用。

圖17 展示了利用傳統貪心算法進行計算圖算子調度時GPU 的資源利用率測試結果。數據顯示，第一階段由于同時調度三個算子Conv(a)、Conv(b)、Conv（d），此時GPU 板卡的資源競爭劇烈，GPU 的資源利用率達到95%，第二階段只調度Conv(c)算子，GPU 的利用率只有20%左右，第一階段和第二階段的任務負載存在嚴重失衡的現象。圖18展示了經過算子合并優化后GPU 的資源利用率測試結果，數據顯示，第一階段需要執行合并后的(abd)算子，雖然只有一個算子，但該算子計算量并未減小，由于算子合并后內核啟動次數減少GPU 利用率降為90%左右，第二階段同樣調度Conv(c)算子，資源利用率同樣為20%。綜上，算子合并優化技術沒有解決GPU 資源負載均衡問題。

圖17 貪心算法實現算子調度時的資源占用分析Fig.17 Analysis of resource occupancy when greedy algorithm implements operator scheduling

圖18 算子合并優化后資源占用分析Fig.18 Analysis of resource occupancy after operator merging optimization

圖19 展示了經過計算圖分層調度優化后，各層的資源利用率和執行時間。由于第一層的算子數量由原來的三個算子Conv(a)、Conv(b)、Conv(d)減少為一個合并算子Conv(ad)，GPU 的利用率由90%降為70%，執行時間由0.2 ms 降為0.16 ms。第二層由調度一個算子Conv(c)變成同時并行調度兩個算子Conv(c)和Conv(b)，GPU 利用率由20%升為50%，執行時間由0.16 ms 上升為0.18 ms。但是整個計算圖的執行時間由0.36 ms 減小為0.34 ms。因此經過圖分層優化之后的算子調度更加均衡，GPU 的資源利用率在各階段也更加平穩，同時進一步提高了計算性能。

圖19 計算圖分層調度優化后的資源利用率Fig.19 Resource utilization after optimization of hierarchical scheduling of computational graphs

3.4 計算圖分層自動流水優化

本節基于Gpipe 微批次劃分流水算法與本文的分層自動流水算法進行對比測試。圖20 為Gpipe 流水示意圖，此時微批次劃分數目為4，流水級為5。從圖中可以看出，整個流水計算過程中由于流水級劃分過多導致流水建立時間較長，設備空閑等待的時間占比超過了50%，大部分時間GPU 都處在空閑的狀態，帶來了資源的浪費。

圖20 微批次數目為4，5級流水的Gpipe流水效果Fig.20 Gpipe pipeline results with the number of microbatches of 4 and 5-level flow

圖21 展示了經過分層多級自動流水優化和微批次數目劃分優化之后的流水效果。相比Gpipe，流水級由5 級合并為3 級，微批次劃分的數量由4 增大為通過微批次劃分查找算法自動查找得到的6，從圖中可以看出流水過程中設備空閑的時間占比由原來的高于50%降低為25%左右。

為了更好地評估本文中的三種優化技術對不同模型訓練性能的影響，本節選取兩種不同類型的深度學習網絡模型進行實驗驗證和性能對比：一種是并行分支較少的網絡模型，比如Resnet50；第二種是并行分支較多的模型，比如InceptionV3、NasNet、ResNeXt 等。本文分別在上述兩種不同類型的模型中進行了數據加載與模型訓練負載均衡優化、計算圖分層調度優化和計算圖分層自動流水優化，并與TensorFlow、TensorFlow-XLA、TensorFlow-ngraph 等未進行相關優化的深度學習框架或編譯器進行性能對比。

圖22 展示了歸一化后四種模型經過數據加載與模型訓練負載均衡優化、計算圖分層調度優化和分層自動流水優化后的性能。實驗結果表明兩種不同種類的深度學習網絡模型在經過本文提出的三種優化之后模型訓練的性能均有所提升，由于不同類型的模型中并行分支數量不一樣，性能的加速比也有所不同。對于Resnet50模型，由于該網絡模型具有較少的并行分支，分層優化和自動流水方法取得的性能提升只有2%左右；對于InceptionV3這種具有較多分支并行的模型，單獨的分層優化方法就可以達到5%以上的性能提升，將分層優化和自動流水兩種優化方法結合使用時，相比其他網絡模型訓練框架性能提升達到10%左右。最終，本文將三種優化技術在Mofang 中同時使用，相比TensorFlow 深度學習框架，Mofang可以提升5%～15%的訓練性能。

3.5 資源利用率

為了更好地驗證本文提出的計算圖分層調度優化方法可以實現計算圖算子間的均衡調度，提高設備資源的利用率，本實驗在InceptionV3 網絡模型訓練期間使用自研工具Gsee 實時獲取GPU 板卡的資源利用率數據。

圖23 統計了InceptionV3 網絡模型訓練過程中分層優化前后的GPU 板卡利用率。從實驗結果可以看出未經過分層優化前，GPU 的板卡資源利用率波動較大，最大值可以達到95%，最小值為20%，存在嚴重的資源利用率不均衡問題。經過計算圖分層調度優化后，每個階段的算子的調度相對均衡，GPU板卡的資源利用率穩定在50%～60%，在經過多層分級流水優化之后GPU 板卡的設備利用率均達到90%以上。

圖23 InceptionV3訓練過程中GPU板卡利用率對比Fig.23 Comparison of GPU utilization during InceptionV3 training

3.6 系統能耗分析

本節實驗使用系統采樣的方法統計整個訓練過程中CPU 和GPU 板卡的實時能耗，以進一步分析本文設計的優化方法對系統整體能耗的影響。

對于CPU 處理器，其開啟的線程數和能耗基本成正比關系，線程數越少，CPU 處理的任務就相應較少，處理器能耗越低。數據加載和模型訓練負載均衡優化通過開啟滿足模型訓練需求的最少的線程數，來盡可能地節省CPU 的資源占用，達到了降低CPU處理器能耗的目的。

對于GPU 板卡，本文采用系統采樣的方式統計其設備能耗。在模型的訓練過程中，每隔1 ms 設置一個采樣點，每次采樣時讀取GPU 當前的能耗。由于不同優化方法，模型訓練完成的時間不同，為了實現準確的數據對比，本文統計完成一個epoch 訓練時的GPU 設備能耗。選用InceptionV3 網絡模型進行能耗統計實驗，訓練周期為1 個epoch。經過數據加載和模型訓練自動負載均衡優化，CPU 的線程數由64降低為48，能耗降比達到25%。

圖24 統計了在不同的優化條件下InceptionV3模型訓練1 個epoch 時GPU 板卡的能耗變化以及訓練1 個epoch 的時間消耗。結果顯示，采用計算圖分層和流水并行優化后，訓練1 個epoch 的時間由30 s降低為20 s左右，提高了整體訓練的性能。

圖24 InceptionV3訓練時GPU板卡能耗變化Fig.24 GPU power consumption changes during InceptionV3 training

本文通過能耗曲線和橫坐標圍成的面積來計算GPU 的整體能耗。通過數學計算可得，未優化時GPU 板卡的能耗為4 138 W，經過分層優化之后能耗降低為3 521 W，再經過分層自動流水優化之后能耗降低為3 512 W。實驗結果表明，本文的分層優化算法能有效降低GPU 板卡的整體的能耗，降比達到15%左右。因此本文的計算圖負載均衡優化策略能夠在提高硬件設備資源利用率的同時，優化模型訓練過程中系統整體能耗。

4 結束語

深度學編譯器能夠實現前端深度學習框架與后端AI 硬件設備之間的解耦化映射，目前深度學習編譯器優化技術已經成為深度學習領域十分重要的分支。本文針對深度學習編譯器中前端計算圖的負載均衡問題提出了三種優化技術，分別是數據加載與模型訓練的負載均衡、計算圖的分層調度優化和計算圖分層自動流水優化。實驗結果表明經過上述優化之后，AI模型訓練性能與其他框架相比獲得2%到10%的提升，同時還可以通過負載均衡技術提高硬件設備的資源利用率，降低系統整體能耗，促使性能和能耗達到一個相對平衡的狀態。

目前本文在計算圖的并行調度優化和能耗優化方面做了一些探索性的工作，現階段本文的優化前端支持TensorFlow 深度學習框架以及編譯時自定義前端，后端支持CPU、GPU、FPGA 等硬件設備，后續將繼續優化本文的技術方案，支持更多的前端深度學習框架和后端加速設備。此外，如何自動尋找計算圖在多種后端設備上更優的聯合布局以及如何通過負載均衡技術進一步降低整個系統訓練過程中的能耗，也是未來主要的一個工作方向。