GNNSched:面向GPU的圖神經網絡推理任務調度框架*

孫慶驍,劉 軼,楊海龍,王一晴,賈 婕,欒鐘治,錢德沛

(北京航空航天大學計算機學院,北京 100191)

1 引言

深度學習DL(Deep Learning)已在大量應用領域中得到廣泛使用,從目標檢測和圖像分類到自然語言處理和機器翻譯。隨著更多新興深度神經網絡DNN(Deep Neural Network)模型被提出,DNN模型對算力的需求呈現快速增長的趨勢,研究人員開始利用TPU(Tensor Processing Unit)和GPU等硬件加速器來提高DL任務運行性能。特別是GPU,由于其擅長處理DNN模型中的大量高度并行化矩陣計算,且得到了主流DNN框架的普遍支持,已經成為主流服務器中提供DNN模型算力的主體[1]。

與此同時,由于強大的節點表示能力,圖神經網絡GNN(Graph Neural Network)在基于圖的預測任務上取得了不錯的效果[2]。GNN結合圖操作和神經計算來表征數據關系。由于圖數據的不規則性,在GPU上實現高性能GNN極具挑戰性。學術界雖然提出了基于節點分區和緩存合并的優化策略來解決GNN執行中的負載不均衡和線程分歧[3]等問題,但是圖相關算子的實現仍使得GPU利用率低。例如,PyG(PyTorch Geometric)[4]通過消息傳遞單獨更新節點特征,但其頻繁的數據移動會導致計算停頓。DGL(Deep Graph Library)[5]使用類SpMM(Sparse Matrix-matrix Multiplication)的內核來實現同時更新,但稀疏數據讀取會降低訪存效率。

為了簡化集群管理,最常見的方法是將GPU資源分配的最小粒度設置為整個GPU[6]。而由于GPU算力的不斷提高,單個DL任務很難充分利用GPU資源[7],特別是對于GNN這類訪存密集型任務,其性能會隨著分配的計算資源的增加而達到飽和狀態[8]。通過對DNN的研究發現,多個DNN任務可以在GPU上共置以提升資源利用率。工業界實現了多進程服務器MPS(Multi-Process Server)和多實例GPU MIG(Multi-Instance GPU),以使多個CUDA(Compute Unified Device Architecture)進程通過資源分區共享GPU。在學術界,時間共享[9,10]通過重疊預處理和計算來降低流水線延遲,而空間共享[11,12]允許并發執行DNN,以提供更高的吞吐量。上述機制只適用于具有固定大小輸入的DNN,無法直接適配到顯存消耗和計算強度與模型輸入動態相關的GNN[13]。

相比訓練框架,基于GPU的推理框架需要應對的問題更加復雜。除改進整體吞吐量以外,還必須在限定時間內提供推理結果,以滿足服務質量目標QT(Quality-of-service Target)[14]。然而,單個GPU上運行多個推理任務可能會因為顯存過載導致執行失敗或延遲顯著增加。因此,推理系統需要提前估計推理任務的顯存占用情況,避免其需求超出顯存容量進而觸發開銷較大的統一虛擬內存UVM(Unified Virtual Memory)數據交換[15]。另一方面,云服務商通常以多租戶方式共享GPU集群資源[16]。在這種情況下,需要根據推理任務的計算模式和顯存占用特點設計靈活的調度機制,以滿足服務質量要求并降低推理響應時延。

為了應對上述挑戰,本文提出并發GNN推理任務調度框架GNNSched(GNN Scheduler),其在GPU上高效地調度和管理并發GNN推理任務。GNNSched首先將推理任務打包到隊列中,并提取有關任務輸入和網絡結構信息。之后,GNNSched分析每個推理任務的計算圖,并量化算子對顯存占用的影響。最后,GNNSched利用多種調度策略對任務進行分組并迭代地分配顯存以供執行。本文最后開展了大量的實驗來評估GNNSched,以驗證其在滿足服務質量和降低延遲等方面的有效性。

GNNSched在任務組內和任務組間分別使用了空間共享和時間共享技術。具體來說,組內的任務通過空間共享提高整體GPU吞吐量,而組間通過重疊數據預處理與計算降低流水線延遲。本文是首次針對并發GNN推理任務的調度優化和顯存管理進行研究。GNNSched已開源于:https://github.com/sunqingxiao/GNNSched。

本文的具體工作如下:

(1)提出了并發GNN推理任務管理機制,通過細粒度的顯存管理和工作器(Worker)分配以自動執行GNN推理任務。此外,提出了多種調度策略對任務進行分組。

(2)提出了GNN推理任務顯存占用估計策略,針對GNN算子設計了顯存成本函數,并通過遍歷前向傳播的計算圖來估計GPU顯存占用情況。

(3)實現了并發GNN推理任務調度框架GNNSched,其可以有效地調度和管理在GPU上的并發GNN推理任務。實驗結果表明,GNNSched能夠滿足服務質量要求并降低推理任務響應時延。

2 背景

2.1 圖神經網絡簡介

近年來,一些研究致力于將深度學習應用于圖等非結構化數據[2]。不同于傳統深度學習模型處理的圖像和文本等密集數據,圖表示稀疏且連接不規則。圖中每個節點與一個特征向量相關聯,節點之間的邊表示圖拓撲結構,并以邊的權重進行量化。GNN以圖結構數據作為輸入,綜合圖結構和節點特征來學習數據關系。表1列出了重要的GNN符號。

Table 1 Explanations of important symbols in GNN表1 GNN中重要符號說明

圖卷積神經網絡GCN(Graph Convolutional Network)[17]是面向圖學習的最成功的網絡之一。GCN緩解了圖的局部鄰域結構過擬合的問題,其圖操作如式(1)所示:

(1)

SAGE(SAmple and aggreGatE)[18]進一步應用采樣的方式來為每個節點獲取固定數目的鄰居。SAGE的圖操作如式(2)所示:

(2)

其中SN(v)是節點v的隨機采樣鄰居,fk(·,·)是聚合函數。圖同構網絡GIN(Graph Isomorphism Network)[19]通過可學習參數εk調整中心節點的權重,其圖操作如式(3)所示:

(3)

基于上述分析,GNN的核心計算可以抽象為式(4):

(4)

2.2 GNN框架的計算模式

圖1展示了GNN前向傳播的計算流程。典型的GNN層包含聚合階段和更新階段,其結合了圖操作和神經計算。聚合階段從節點的每個鄰居檢索一個特征向量,并將這些向量聚合成一個新的特征向量。更新階段執行多層感知機MLP(MultiLayer Perceptron)等神經操作以轉換每個節點的特征向量。GNN框架根據圖結構進行圖操作,其中邊表示數據傳輸。DGL通過中心鄰居模式引入了節點級并行,它從特征矩陣中獲取數據,再執行類SpMM的歸約操作以同時更新節點特征。PyG通過MessagePassing抽象引入了邊級并行,它通過消息傳遞在所有邊上直接生成消息,再分別執行歸約操作。

Figure 1 Computation workflow of GNN圖1 GNN計算流程

DGL和PyG都受到了GPU計算資源不足的限制。DGL應用類SpMM內核實現節點特征的更新,但是由于圖結構的不規則性,顯存訪問成為了其性能瓶頸。PyG通過聚合內核實現單獨的節點特征更新,以提高訪存效率,但是耗時的數據移動會導致計算停頓。為解決上述問題,可以利用GNN推理任務的共置機制最大化GPU吞吐量。然而,這需要預知推理任務的顯存消耗情況以免顯存過載。不同于擁有固定大小輸入的神經網絡,GNN層的輸出維度與圖維度和特征長度緊密相關。此外,為了得到更精確的估計值,圖傳播的顯存消耗也需要被納入考慮。

2.3 推理任務的GPU共享

在工業界,主流做法是將GPU分配的最小粒度設置為整個GPU[6]。雖然這樣的設置簡化了集群資源管理,但導致GPU資源的利用率較低。因此,GPU共享逐漸成為在GPU上共置推理任務的一項基本技術。例如,即使是單個服務也可能包含多個異構推理任務[21],如何將其映射到GPU是重要挑戰。對具有自身計算要求的不同推理任務必須適時地加載到GPU,以滿足其服務質量目標,同時提高整體吞吐量。

現有工作提出了基于時間或空間共享的機制[9,12],以實現深度學習任務在GPU上的共置。時間共享高度靈活,將GPU顯存和核心專用于特定持續時間的單次執行。PipeSwitch[9]利用流水線模型傳輸和主備Worker來最小化切換開銷,從而滿足推理任務的嚴格服務質量目標。REEF[10]改造了GPU驅動以支持軟件隊列和計算單元的重置,并主動搶占批量內核從而在微秒級啟動實時推理任務。盡管時間共享通過重疊數據預處理和計算來隱藏延遲,但仍難以充分發掘GPU的計算潛力。例如,對于遞歸神經網絡構成的語言模型,計算單元往往會長時間閑置[16]。

相比之下,空間共享允許在不違反服務質量的情況下提供更高的GPU吞吐量。應用空間共享的一個限制是并發任務的工作集大小。如果工作集大小超過GPU顯存,系統必須將數據交換到主機,這會掩蓋空間共享的優勢。GSLICE[14]采用自調整算法并根據性能反饋調整每個推理任務的線程占用率。Abacus[11]通過確定性算子重疊實現了并發推理任務的延遲可預測性,并設計基于配額的控制器以確定算子執行順序。Choi等人[12]創建GPU資源抽象層為推理任務分配有效資源,再通過性能預測模型評估空間共享的潛在干擾開銷。然而,這些方法均未考慮并發推理任務可能帶來的顯存過載,這可能會導致任務運行失敗或耗時的設備-主機數據交換[15]。

上述機制針對的是具有固定大小輸入的傳統神經網絡,沒有量化不規則圖對顯存消耗的影響。此外,對于同一批次到達的推理任務,可以結合時間共享和空間共享在GPU上實現更靈活高效的調度機制。

3 方法設計與實現

3.1 設計概要

本節提出GNN推理調度框架GNNSched來維護GNN推理隊列,每次從隊列頭部取同一批次的推理任務進行高效調度和管理。如圖2所示,灰色模塊是由GNNSched設計或擴展的。GNNSched由4個重要組件組成,包括顯存管理器、峰值顯存消耗PMC(Peak Memory Consumption)分析器、任務調度器和Worker分配器。顯存管理器維護統一的顯存池,按需分配顯存;PMC分析器提取運行時信息以估計推理任務的顯存消耗;任務調度器根據調度策略確定分組和任務執行順序;Worker分配器將訓練任務綁定到具體的Worker上,執行后返回結果。注意Worker是指負責任務執行的進程,其常駐在推理系統中,跨不同組串行執行推理任務。

Figure 2 Design overview of GNNSched圖2 GNNSched設計概要

圖2為GNNSched的設計概要。GNNSched將CUDA Allocator集成到DL后端以實現顯式的GPU顯存管理。GNNSched將使用GNN框架實現的推理任務打包到任務隊列中(T0,T1和T2為推理請求到達時間),其中PMC分析器提取模型輸入和網絡結構的詳細信息。PMC分析器將GNN模型的計算圖表示為有向無環圖DAG(Directed Acyclic Graph),并使用公式量化每個算子對顯存消耗的影響。PMC分析器將PMC信息發送給任務調度器,任務調度器使用特定策略對推理任務進行分組和重排,再迭代地從隊列中彈出任務組。在每次迭代中,顯存管理器和Worker分配器接收信號以分配共享GPU顯存和執行推理任務。

3.2 任務管理機制

并發任務空間共享的關鍵在于GPU顯存的細粒度管理。本文擴展了DL后端(PyTorch)的統一顯存池,針對DL任務的特點,將顯存池劃分為預留顯存和分配顯存。預留顯存存儲框架內部數據,例如CUDA上下文和模型工作區,通常在任務執行前預先分配,分配顯存存儲任務運行時產生的張量,例如層輸出。GNNSched從分配顯存的一端連續插入任務Buffer,以確保分配顯存被完全占用(見圖2)。GNNSched通過顯存消耗估計來指定每個任務Buffer的分配大小,再將其插入到特定的顯存位置。對于分配顯存,GNNSched通過滿足對齊要求的額外顯存填充來處理內部張量碎片。

圖3給出了GNNSched中任務管理的整體工作流程。對于同一批次的推理任務,任務調度器執行分組操作,Worker分配器以任務組為粒度迭代地提交到GPU。以這種方式,GNNSched避免了顯存碎片問題。顯存碎片可能浪費并發推理機會,同時使內存維護復雜化。在每次迭代中,Worker分配器將推理任務映射到具體Worker(wk0,wk1和wk2)。每個Worker串行處理Buffer插入、任務執行和結果返回等操作;Worker間并行以重疊相關操作,其中不同推理任務以空間共享的方式執行。并行Worker均返回結果后,GNNSched清除所有Buffer并前進到下一組。注意,GNNSched在GPU上執行當前任務組的同時,在主機端預先處理下一組的輸入數據,因而有效地降低了流水線延遲。

Figure 3 Overall workflow of task management圖3 任務管理的整體工作流程

Figure 4 Computation graph of two-layer SAGE model圖4 2層SAGE模型的計算圖

從以上分析可知,GNNSched靈活的管理機制可以支持推理任務的任意分組和調度順序。這對于提高輸入敏感GNN的訪存效率是必要的,其PMC隨圖維度和特征長度而顯著變化。

3.3 顯存消耗估計

受文獻[22]的啟發,本文將GNN推理的計算圖CG表示為DAG,如式(5)所示:

(5)

其中,節點opi表示數學調用的算子,邊edj指定執行依賴。TO=〈ed1,ed2,…,edm〉是DAG規定的拓撲順序,其通過查閱DL后端[23]內的拓撲順序預先生成。GNNSched利用TO遍歷計算圖并根據張量的分配和釋放更新PMC。

根據拓撲遍歷,估計GPU顯存消耗可以形式化為計算圖上每個算子所需顯存的累加。本文為每個算子定義了DL后端無關的顯存成本函數MCF(Memory Cost Function)。顯存成本函數返回一組已分配的具有類別和形狀的張量,這些張量通過輸入維度和形狀推斷來推導得到。算子op的MCF可以表示為式(6):

MCF(op)=W(op)∪O(op)∪E(op)

(6)

其中,W、O和E分別是權重張量、輸出張量和臨時張量的集合。

本文還需要考慮GNN的定制化細節。GCN補充自循環以使中心節點的聚合表示包含其自身特征。即便如此,GNN框架中仍不可避免地會創建部分臨時張量來維護任務執行。因此,GNNSched將估計PMC乘以固定閾值threshhold,以保證插入的Buffer能夠滿足計算過程的顯存需求。

Table 2 Allocated tensors and their sizes表2 分配張量的類型和大小

上述抽象和形式化對于各種GNN框架的PMC估計是通用的。GNNSched還可以通過修改計算圖或顯存成本函數來適配其他GNN。

3.4 調度策略

GNNSched的靈活管理機制為調度策略提供了很大的設計空間。與文獻[11]一致,本文將每個推理任務的QT設置為其單獨運行時間的2倍。因此,高QT任務具有更長的執行時間,而執行時間通常與計算強度正相關。本文實現了2種非搶占式的調度策略,包括最短QT優先SQTF(Shortest-QT-First)和平衡QTBQT(Balanced-QT)。所有策略均具有“安全”條件,以確保并發任務的顯存使用不超過GPU顯存容量。幸運的是,GNNSched可以估計推理任務的PMC并進行累加,從而使每個任務組都處于“安全”位置。接下來,說明本文2種調度策略的具體細節:

(1)SQTF策略:最簡單的先進先出FIFO(First-In-First-Out)算法可以在滿足“安全”條件下盡可能地將更多的任務打包到同一組中。然而,FIFO算法可能會導致短期任務因等待大型正在進行的任務完成而遭受長時間的排隊延遲。因此,提出了最短作業優先算法和最短剩余時間優先算法[24],以降低隊列中任務的排隊延遲。然而,任務持續時間或剩余時間需要離線分析,這會影響可用性和運營成本[25]。為了解決以上問題,本文提出SQTF策略(見算法1),其將隊列中的任務索引按照QT的升序進行排序,然后遍歷排序生成任務組。此外,本文設置了分組閾值以改善時間維度的負載均衡。具體來說,分組閾值由當前批次的推理任務的PMC總和計算得到,在分組過程中保證不同任務組的顯存分配大小相近。

算法1 SQTF調度策略輸入:升序隊列Deque,分配顯存大小MA,推理任務的PMC總和SP;輸出:任務組列表taskGroup。1:queSize←Deque.size;//原始隊列大小2:gTH=ceil(SP/ceil(SP/MA));//分組閾值//初始化組PMC和組計數器3:gPMC,gCounter←MA,-1;4:for i in range [0,queSize)do5: task←Deque.popleft();//從右側彈出任務6: if gPMC>gTHthen7: taskGroup.append([]);//前進到下一組8: gPMC,gCounter←0,gCounter+1;9: end if10: taskGroup[gCounter].append(task);11: gPMC←gPMC+task.PMC;//累加PMC12:end for

(2)BQT策略:該策略繼承了SQTF策略的QT排序和分組閾值機制。盡管SQTF策略降低了短期任務的排隊延遲,但將具有高QT的推理任務置于同一組可能會加劇資源沖突。BQT策略用于平衡任務時長和排隊時間。BQT策略的原則是將計算強度高的任務和強度低的任務歸為一組,從而降低并發執行的性能干擾。算法2為BQT策略的實現細節。該策略根據QT的升序將推理任務推入雙端隊列(Deque)。每次迭代從Deque的右端或左端彈出任務(第5～9行),再將任務分到當前組并累加PMC(第14,15行)。當組PMC大于分組閾值時,前進到下一組并初始化組計數器(第10～13行)。重復上述步驟直到Deque為空,當前批次的任務組生成結束。

算法2 BQT調度策略輸入:升序隊列Deque,分配顯存大小MA,推理任務的PMC總和SP;輸出:任務組列表taskGroup。1:queSize←Deque.size;//原始隊列大小2:gTH=ceil(SP/ceil(SP/MA));//分組閾值//初始化組PMC和組計數器3:gPMC,gCounter←MA,-1;4:for i in range [0,queSize)do5: if i%2==1then6: task←Deque.pop();//從右側彈出任務7: else8: task←Deque.popleft();//從左側彈出

9: end if10: if gPMC>gTHthen11: taskGroup.append([]);//前進到下一組12: gPMC,gCounter←0,gCounter+1;13: end if14: taskGroup[gCounter].append(task);15: gPMC←gPMC+task.PMC;//累加PMC16:end for

3.5 實現細節

GNNSched的系統原型由C++和Python代碼實現,并構建在PyG和PyTorch[26]之上。然而,GNNSched背后的思想可通用于其他GNN框架或DL后端。本文使用CUDA IPC擴展PyTorch的Allocator模塊從而實現GPU顯存池的顯式管理。本文通過添加函數來支持在特定CUDA流中插入Buffer以及從顯存池中清除Buffer。Buffer插入函數可以被多次調用,以實現GNN推理任務的空間共享。任務組執行完成后,Buffer清除函數被調用來清除顯存分配。注意PyTorch為張量分配的實際顯存大小需滿足一定的對齊要求。為了解決這一問題,本文通過填充來使得顯存分配大小滿足512字節的倍數。

GNNSched由分別負責隊列管理和任務執行的調度進程和Worker進程組成。調度進程監聽客戶端通過TCP端口發送的任務請求,將任務打包到隊列中并加載模型結構。接下來,調度進程通過計算圖遍歷來分析PMC并生成任務組。在每個組迭代中,調度進程將任務的哈希索引發送到Worker進程。Worker進程根據哈希索引識別任務并將其分配給Worker線程。每個Worker線程加載相應的GNN模型并將其附加到CUDA流。結果返回后,Worker線程通過PyTorch Pipe API向調度進程發送“完成”信號。當接受到的信號數量等于組大小時,調度進程迭代到下一個任務組。

4 實驗與結果分析

4.1 實驗設置

(1)硬件和軟件配置。硬件規格如表3所示。操作系統為Ubuntu 20.04,編譯器版本為GCC v9.3和NVCC v11.1。GNNSched基于PyG v1.7和PyTorch v1.8構建。本文修改了PyTorch以支持GNNSched的顯式任務共置和顯存管理。

Table 3 Hardware configuration表3 硬件配置

(2)圖數據集和任務隊列。用于實驗的圖數據集如表4所示。本文為每個圖數據集隨機生成25個子圖,用作網絡輸入。本文選取3個典型GNN(包括GCN、SAGE和GIN),其中層數和層寬分別設置為8和256。本文將SAGE的采樣率設置為0.5(與PyG默認設置一致)。本文為每個GNN生成由100個推理任務組成的隊列,不同任務具有不同的子圖輸入。此外,本文從以上隊列中均勻采樣100個任務以獲得具有不同GNN的任務隊列(命名為MIX)。隊列中任務到達的時間遵循泊松分布。同時到達的任務數量均值被設置為1和2,分別代表低負載(Low Load)和高負載(High Load)。多樣的任務組織用于對GNNSched的有效性進行全面評估。

Table 4 Graph datasets表4 圖數據集

(3)對比方法和指標。本文將具有2種調度策略(SQTF和BQT)的GNNSched與Default和MPS進行對比。Default采取FIFO的方式串行執行推理任務。為了公平起見,本文基于MPS實現了2個變體MPS-Base和MPS-Aggr。MPS-Base的最大并發數量為4,MPS-Aggr并發執行同一批次的所有任務。本文為MPS啟動UVM以處理可能的顯存過載。為了評估推理任務的服務質量和執行效率,本文選取服務質量違反率(QoS violation)、響應延遲(medium/90%-ile/99%-ile latency)、作業完成時間JCT(job completion time)和排隊時間QUET(QUEuing Time)作為評估指標。本文將每種方法運行10次并給出平均結果以隔離隨機性的影響。

4.2 服務質量評估

圖5給出了不同方法的服務質量違反率對比。相比Default和MPS,GNNSched在所有任務隊列下均取得更低或等同的服務質量違反率。在低負載下,GNNSched使得所有任務均滿足其服務質量要求。即便在高負載下,GNNSched的服務質量違反率最高僅為8%,而Default和MPS的最高違反率分別達到了93%和91%。另外,還注意到,MPS-Base和MPS-Aggr呈現明顯差異性的實驗結果。MPS-Base的服務質量違反率總是低于20%,而MPS-Aggr在半數隊列下甚至比Default的違反率更高。這是因為MPS的貪心機制使得顯存過載,UVM的分頁開銷顯著降低了并發任務的性能。盡管MPS-Base的靜態分組有效規避了UVM的使用,然而其欠缺任務調度的靈活性。GNNSched通過PMC預估計和高效調度來保證在顯存安全下盡可能地充分利用計算資源。

Figure 5 Comparison of QoS violation rate圖5 服務質量違反率對比

4.3 延遲評估

圖6給出了歸一化為QT的響應延遲對比。Default的串行執行忽略了任務并發的機會,在GIN隊列下的medium延遲、90%-ile延遲和99%-ile延遲分別達到了QT的6.6倍、12.2倍和17.3倍。MPS-Aggr由于UVM開銷表現不穩定,在GCN隊列下取得7.5倍QT的99%-ile延遲。相比之下,GNNSched和MPS-Base在所有任務隊列下均取得小于2倍QT的99%-ile延遲。以上結果表明,GNNSched和MPS-Base具有更低的尾延遲,這對于提升用戶粘性尤為重要。

Figure 6 Comparison of request latency圖6 響應延遲對比

從圖6可以觀察到,GNNSched在大多數案例中取得了最低延遲。GNNSched相比MPS-Base分別平均降低了26.2%,27.4%和31.9%的medium延遲、90-ile延遲和99-ile延遲。這說明GNNSched能夠提供穩定的服務質量,有效避免了處理批量任務時的長尾延遲。GNNSched的優越性主要來自于其充分發掘了推理任務的并發機會。另外,GNNSched設置的分組閾值不但改善了時間維度的負載均衡,還在很大程度上緩解了任務空間共享時的性能干擾。

4.4 性能評估

圖7給出了歸一化為QT的平均JCT對比。GNNSched在所有任務隊列下均取得了最短的JCT。與Default、MPS-Base和MPS-Aggr對比,GNNSched的JCT分別平均降低了60.6%,26.8%和62.7%。除了靈活的任務并發機制,GNNSched的顯存池管理同樣有助于性能提升,其避免了張量的頻繁分配和釋放。值得注意的是,GNNSched-BQT在大多數任務隊列下略微好于GNNSched-SQTF,JCT平均降低了0.6%。這是因為BQT策略將計算強度高的和計算強度低的任務歸為一組,降低了對計算和緩存資源的競爭。

Figure 7 Comparison of average JCT圖7 平均作業完成時間對比

圖8給出了歸一化為QT的平均排隊時間對比。對于低負載,MPS-Base和GNNSched在所有任務隊列下的排隊時間均為0。這意味著上一批次的推理任務總能在當前批次到達前完成,側面證明了空間共享的有效性。對于高負載,MPS-Base在所有隊列下均存在等待延遲,而GNNSched在SAGE隊列和GIN隊列下的排隊時間仍為0。GNNSched的靈活共享機制重疊了數據預處理和計算,改進了任務組間的流水線執行效率。另外,GNNSched-SQTF明顯比GNNSched-BQT的排隊時間更短,平均降低了0.8%。原因是SQTF策略優先調度短時執行任務,緩解了等待長時任務完成時的隊頭堵塞。

4.5 顯存估計精度

PMC的精確估計對于保證并發執行的顯存安全是必要的。本文使用相對誤差RE(relative error)來度量估計精度。圖9給出了GNN推理任務的PMC估計的相對誤差,可以觀察到所有任務的誤差均低于8%。原因是顯存成本函數準確地獲取了GNN算子的顯存消耗。另一方面,隨著子圖節點數和邊數的變化,GNNSched取得了穩定的估計精度。這表明基于網絡結構生成的DAG很好地表示了前向傳播的計算流程。根據以上結果,本文將乘法閾值(threshhold)設置為1.1從而確保空間共享下的顯存安全。

Figure 9 Relative errors of PMC estimation圖9 PMC估計的相對誤差

4.6 開銷分析

本節將GNNSched的處理開銷歸一化為任務推理時間。處理開銷可以分為PMC估計、任務調度和Buffer插入3個部分。PMC估計通過顯存成本函數遍歷計算圖以獲取PMC信息。任務調度將并發推理任務組織為隊列,再根據調度策略對隊列進行重排并生成任務組。Buffer插入基于PMC信息將Buffer插入到分配顯存的特定位置。圖10顯示了GNNSched處理開銷的時間分解。可以看到,處理開銷相對于任務執行可以忽略不計,在低負載和高負載下的平均時間僅為推理時間的2.4和3.0。這表明使用GNNSched來管理相異的推理任務可以有效提高GPU吞吐量。

Figure 10 Time breakdown of processing overhead圖10 處理開銷的時間分解

5 相關工作

5.1 圖神經網絡的算子加速

最近研究工作深入挖掘了GNN的計算特征,并針對GPU架構進行細粒度優化。FeatGraph[27]結合圖劃分與特征維度優化聚合階段的緩存利用。Huang等人[3]通過局部敏感哈希對節點進行聚類,再對鄰居進行分組以緩解負載不均衡。GNNAdvisor[13]通過引入warp對齊的線程映射和維度劃分來減少線程分歧。QGTC(Quantized graph neural networks via GPU Tensor Core)[28]基于低位數據表示和位分解的量化技術實現了張量核心定制的計算內核。本文工作與上述工作在性能優化上互補,GNNSched的重點在于GNN推理任務的并發調度。

5.2 深度學習推理調度系統

許多研究工作提出了各種調度系統來改進部署在集群上的深度學習推理任務的服務質量。Clipper[29]設計了模型抽象層以在框架之上實現緩存和自適應批處理策略。MArk[30]動態處理批量推理請求并適時地使用GPU以改進性能。Nexus[21]通過打包機制將推理任務組調度到集群并指定每個任務的GPU使用。Clockwork[31]提供性能可預測系統,其通過中央控制器調度請求并預先放置模型。這些系統均未考慮推理任務在單GPU上的空間共享,從而導致計算資源的利用率較低。

6 結束語

本文提出了并發GNN推理任務調度框架GNNSched,其可以高效管理GPU上共置的推理任務。GNNSched將推理任務組織為隊列,并提取圖輸入數據和模型網絡結構信息;之后,GNNSched對每個任務的計算圖進行分析,并利用算子成本函數估計任務顯存占用;最后,GNNSched實現了多種調度策略用于生成任務組,并將其迭代地提交到GPU上執行。實驗結果表明,GNNSched可以滿足推理任務服務質量要求并降低響應時延。