GNNSched:面向GPU的圖神經(jīng)網(wǎng)絡(luò)推理任務(wù)調(diào)度框架*

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

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

1 引言

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

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

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

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

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

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

本文的具體工作如下:

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

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

(3)實現(xiàn)了并發(fā)GNN推理任務(wù)調(diào)度框架GNNSched,其可以有效地調(diào)度和管理在GPU上的并發(fā)GNN推理任務(wù)。實驗結(jié)果表明,GNNSched能夠滿足服務(wù)質(zhì)量要求并降低推理任務(wù)響應(yīng)時延。

2 背景

2.1 圖神經(jīng)網(wǎng)絡(luò)簡介

近年來,一些研究致力于將深度學(xué)習(xí)應(yīng)用于圖等非結(jié)構(gòu)化數(shù)據(jù)[2]。不同于傳統(tǒng)深度學(xué)習(xí)模型處理的圖像和文本等密集數(shù)據(jù),圖表示稀疏且連接不規(guī)則。圖中每個節(jié)點與一個特征向量相關(guān)聯(lián),節(jié)點之間的邊表示圖拓撲結(jié)構(gòu),并以邊的權(quán)重進行量化。GNN以圖結(jié)構(gòu)數(shù)據(jù)作為輸入,綜合圖結(jié)構(gòu)和節(jié)點特征來學(xué)習(xí)數(shù)據(jù)關(guān)系。表1列出了重要的GNN符號。

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

圖卷積神經(jīng)網(wǎng)絡(luò)GCN(Graph Convolutional Network)[17]是面向圖學(xué)習(xí)的最成功的網(wǎng)絡(luò)之一。GCN緩解了圖的局部鄰域結(jié)構(gòu)過擬合的問題,其圖操作如式(1)所示:

(1)

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

(2)

其中SN(v)是節(jié)點v的隨機采樣鄰居,fk(·,·)是聚合函數(shù)。圖同構(gòu)網(wǎng)絡(luò)GIN(Graph Isomorphism Network)[19]通過可學(xué)習(xí)參數(shù)εk調(diào)整中心節(jié)點的權(quán)重,其圖操作如式(3)所示:

(3)

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

(4)

2.2 GNN框架的計算模式

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

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

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

2.3 推理任務(wù)的GPU共享

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

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

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

上述機制針對的是具有固定大小輸入的傳統(tǒng)神經(jīng)網(wǎng)絡(luò),沒有量化不規(guī)則圖對顯存消耗的影響。此外,對于同一批次到達的推理任務(wù),可以結(jié)合時間共享和空間共享在GPU上實現(xiàn)更靈活高效的調(diào)度機制。

3 方法設(shè)計與實現(xiàn)

3.1 設(shè)計概要

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

Figure 2 Design overview of GNNSched圖2 GNNSched設(shè)計概要

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

3.2 任務(wù)管理機制

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

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

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

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

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

3.3 顯存消耗估計

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

(5)

其中,節(jié)點opi表示數(shù)學(xué)調(diào)用的算子,邊edj指定執(zhí)行依賴。TO=〈ed1,ed2,…,edm〉是DAG規(guī)定的拓撲順序,其通過查閱DL后端[23]內(nèi)的拓撲順序預(yù)先生成。GNNSched利用TO遍歷計算圖并根據(jù)張量的分配和釋放更新PMC。

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

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

(6)

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

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

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

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

3.4 調(diào)度策略

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

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

算法1 SQTF調(diào)度策略輸入:升序隊列Deque,分配顯存大小MA,推理任務(wù)的PMC總和SP;輸出:任務(wù)組列表taskGroup。1:queSize←Deque.size;//原始隊列大小2:gTH=ceil(SP/ceil(SP/MA));//分組閾值//初始化組PMC和組計數(shù)器3:gPMC,gCounter←MA,-1;4:for i in range [0,queSize)do5: task←Deque.popleft();//從右側(cè)彈出任務(wù)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策略降低了短期任務(wù)的排隊延遲,但將具有高QT的推理任務(wù)置于同一組可能會加劇資源沖突。BQT策略用于平衡任務(wù)時長和排隊時間。BQT策略的原則是將計算強度高的任務(wù)和強度低的任務(wù)歸為一組,從而降低并發(fā)執(zhí)行的性能干擾。算法2為BQT策略的實現(xiàn)細節(jié)。該策略根據(jù)QT的升序?qū)⑼评砣蝿?wù)推入雙端隊列(Deque)。每次迭代從Deque的右端或左端彈出任務(wù)(第5～9行),再將任務(wù)分到當前組并累加PMC(第14,15行)。當組PMC大于分組閾值時,前進到下一組并初始化組計數(shù)器(第10～13行)。重復(fù)上述步驟直到Deque為空,當前批次的任務(wù)組生成結(jié)束。

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

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 實現(xiàn)細節(jié)

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

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

4 實驗與結(jié)果分析

4.1 實驗設(shè)置

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

Table 3 Hardware configuration表3 硬件配置

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

Table 4 Graph datasets表4 圖數(shù)據(jù)集

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

4.2 服務(wù)質(zhì)量評估

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

Figure 5 Comparison of QoS violation rate圖5 服務(wù)質(zhì)量違反率對比

4.3 延遲評估

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

Figure 6 Comparison of request latency圖6 響應(yīng)延遲對比

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

4.4 性能評估

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

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

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

4.5 顯存估計精度

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

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

4.6 開銷分析

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

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

5 相關(guān)工作

5.1 圖神經(jīng)網(wǎng)絡(luò)的算子加速

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

5.2 深度學(xué)習(xí)推理調(diào)度系統(tǒng)

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

6 結(jié)束語

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