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

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

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

1 引言

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

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

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

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

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

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

本文的具體工作如下:

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

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

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

2 背景

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

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

Table 1 Explanations of important symbols in GNN表1 GNN中重要符號(hào)說(shuō)明

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

(1)

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

(2)

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

(3)

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

(4)

2.2 GNN框架的計(jì)算模式

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

Figure 1 Computation workflow of GNN圖1 GNN計(jì)算流程

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

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

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

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

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

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

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

3.1 設(shè)計(jì)概要

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

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

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

3.2 任務(wù)管理機(jī)制

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

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

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

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

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

3.3 顯存消耗估計(jì)

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

(5)

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

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

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

(6)

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

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

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

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

3.4 調(diào)度策略

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

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

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

算法2 BQT調(diào)度策略輸入:升序隊(duì)列Deque,分配顯存大小MA,推理任務(wù)的PMC總和SP;輸出:任務(wù)組列表taskGroup。1:queSize←Deque.size;//原始隊(duì)列大小2:gTH=ceil(SP/ceil(SP/MA));//分組閾值//初始化組PMC和組計(jì)數(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([]);//前進(jìn)到下一組12: gPMC,gCounter←0,gCounter+1;13: end if14: taskGroup[gCounter].append(task);15: gPMC←gPMC+task.PMC;//累加PMC16:end for

3.5 實(shí)現(xiàn)細(xì)節(jié)

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

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

4 實(shí)驗(yàn)與結(jié)果分析

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

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

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

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

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

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

4.3 延遲評(píng)估

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

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

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

4.4 性能評(píng)估

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

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

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

4.5 顯存估計(jì)精度

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

Figure 9 Relative errors of PMC estimation圖9 PMC估計(jì)的相對(duì)誤差

4.6 開(kāi)銷分析

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

Figure 10 Time breakdown of processing overhead圖10 處理開(kāi)銷的時(shí)間分解

5 相關(guān)工作

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

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

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

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

6 結(jié)束語(yǔ)

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