Upptagna GPU:er: Samplings- och pipeliningmetoder påskyndar djupinlärning av stora grafer

Grafer, ett potentiellt omfattande nät av noder sammankopplade av kanter, kan användas för att uttrycka och förhöra relationer mellan data, som sociala kopplingar, finansiella transaktioner, trafik, energinät och molekylära interaktioner. När forskare samlar in mer data och bygger ut dessa grafiska bilder, kommer forskarna att behöva snabbare och effektivare metoder, såväl som mer beräkningskraft, för att genomföra djupinlärning på dem, i form av grafiska neurala nätverk (GNN).  

Nu förbättrar en ny metod, kallad SALIENT (Sampling, sLIcing, and data movemeNT), utvecklad av forskare vid MIT och IBM Research, tränings- och slutledningsprestanda genom att ta itu med tre viktiga flaskhalsar i beräkningar. Detta minskar dramatiskt körtiden för GNN:er på stora datamängder, som till exempel innehåller en skala av 100 miljoner noder och 1 miljard kanter. Vidare fann teamet att tekniken skalas bra när beräkningskraft läggs till från en till 16 grafiska bearbetningsenheter (GPU). Arbetet presenterades vid den femte konferensen om maskininlärning och system.

"Vi började titta på de utmaningar som nuvarande system upplevde när vi skalade toppmoderna maskininlärningstekniker för grafer till riktigt stora datamängder. Det visade sig att det fanns mycket arbete att göra, eftersom många av de befintliga systemen uppnådde bra prestanda främst på mindre datauppsättningar som passar in i GPU-minnet, säger Tim Kaler, huvudförfattare och postdoc vid MIT Computer Science och Artificiell Intelligens Laboratory (CSAIL).

Med stora datamängder menar experter skalor som hela Bitcoin-nätverket, där vissa mönster och datarelationer kan förklara trender eller fult spel. "Det finns nästan en miljard Bitcoin-transaktioner på blockkedjan, och om vi vill identifiera olagliga aktiviteter i ett sådant gemensamt nätverk, då står vi inför en graf i en sådan skala", säger medförfattare Jie Chen, senior forskare och chef av IBM Research och MIT-IBM Watson AI Lab. "Vi vill bygga ett system som kan hantera den typen av grafer och gör att bearbetningen blir så effektiv som möjligt, eftersom vi varje dag vill hänga med i takten i den nya data som genereras."

Kaler och Chens medförfattare inkluderar Nickolas Stathas MEng '21 från Jump Trading, som utvecklade SALIENT som en del av sitt examensarbete; tidigare MIT-IBM Watson AI Lab praktikant och MIT doktorand Anne Ouyang; MIT CSAIL postdoc Alexandros-Stavros Iliopoulos; MIT CSAIL forskare Tao B. Schardl; och Charles E. Leiserson, Edwin Sibley Webster professor i elektroteknik vid MIT och en forskare vid MIT-IBM Watson AI Lab.     

För detta problem tog teamet ett systemorienterat tillvägagångssätt när de utvecklade sin metod: SALIENT, säger Kaler. För att göra detta implementerade forskarna vad de såg som viktiga, grundläggande optimeringar av komponenter som passar in i befintliga ramverk för maskininlärning, som PyTorch Geometric och deep graph library (DGL), som är gränssnitt för att bygga en maskininlärningsmodell. Stathas säger att processen är som att byta ut motorer för att bygga en snabbare bil. Deras metod var designad för att passa in i befintliga GNN-arkitekturer, så att domänexperter enkelt kunde tillämpa detta arbete på sina specificerade områden för att påskynda modellutbildning och reta ut insikter under slutledning snabbare. Tricket, bestämde teamet, var att hålla all hårdvara (CPU:er, datalänkar och GPU:er) sysselsatta hela tiden: medan CPU:n samplar grafen och förbereder minibatcher av data som sedan kommer att överföras via datalänken , den mer kritiska GPU:n arbetar för att träna maskininlärningsmodellen eller göra slutledningar. 

Forskarna började med att analysera prestandan hos ett vanligt förekommande maskininlärningsbibliotek för GNN:er (PyTorch Geometric), som visade ett häpnadsväckande lågt utnyttjande av tillgängliga GPU-resurser. Genom att tillämpa enkla optimeringar förbättrade forskarna GPU-användningen från 10 till 30 procent, vilket resulterade i en prestandaförbättring på 1.4 till två gånger jämfört med offentliga benchmark-koder. Denna snabba baslinjekod skulle kunna köra ett helt pass över en stor träningsdatauppsättning genom algoritmen (en epok) på 50.4 sekunder.                          

För att söka ytterligare prestandaförbättringar, försökte forskarna undersöka flaskhalsarna som uppstår i början av datapipelinen: algoritmerna för grafsampling och förberedelse av minibatch. Till skillnad från andra neurala nätverk utför GNN:er en grannskapsaggregationsoperation, som beräknar information om en nod med hjälp av information som finns i andra närliggande noder i grafen - till exempel i ett socialt nätverksdiagram, information från vänner till vänner till en användare. När antalet lager i GNN ökar, kan antalet noder som nätverket måste nå ut till för information explodera och överskrida gränserna för en dator. Algoritmer för granskning av grannskap hjälper genom att välja en mindre slumpmässig delmängd av noder att samla in; forskarna fann dock att nuvarande implementeringar av detta var för långsamma för att hålla jämna steg med bearbetningshastigheten hos moderna GPU:er. Som svar identifierade de en blandning av datastrukturer, algoritmiska optimeringar och så vidare som förbättrade samplingshastigheten, vilket i slutändan förbättrade bara samplingsoperationen med ungefär tre gånger, vilket tog körtiden per epok från 50.4 till 34.6 sekunder. De fann också att provtagning, i en lämplig takt, kan göras under slutledning, vilket förbättrar den totala energieffektiviteten och prestanda, en punkt som hade förbisetts i litteraturen, noterar teamet.      

I tidigare system var detta provtagningssteg ett tillvägagångssätt med flera processer, vilket skapade extra data och onödig datarörelse mellan processerna. Forskarna gjorde sin SALIENT-metod smidigare genom att skapa en enda process med lätta trådar som höll data på CPU:n i delat minne. Vidare drar SALIENT fördel av en cache med moderna processorer, säger Stathas, parallelliserad funktionsdelning, som extraherar relevant information från noder av intresse och deras omgivande grannar och kanter, inom det delade minnet av CPU-kärncachen. Detta minskade återigen den totala körtiden per epok från 34.6 till 27.8 sekunder.

Den sista flaskhalsen som forskarna tog upp var att pipeline mini-batch-dataöverföringar mellan CPU och GPU med hjälp av ett förhämtningssteg, som skulle förbereda data precis innan det behövs. Teamet beräknade att detta skulle maximera bandbreddsanvändningen i datalänken och få metoden att utnyttjas perfekt; dock såg de bara runt 90 procent. De identifierade och fixade en prestandabugg i ett populärt PyTorch-bibliotek som orsakade onödig rundturskommunikation mellan CPU och GPU. Med denna bugg åtgärdad uppnådde teamet en körtid på 16.5 sekunder per epok med SALIENT.

"Vårt arbete visade, tror jag, att djävulen finns i detaljerna", säger Kaler. "När du är mycket uppmärksam på detaljerna som påverkar prestanda när du tränar ett grafiskt neuralt nätverk, kan du lösa ett stort antal prestationsproblem. Med våra lösningar blev vi helt flaskhalsade av GPU-beräkning, vilket är det idealiska målet för ett sådant system.”

SALIENTs hastighet utvärderades på tre standarddatauppsättningar ogbn-arxiv, ogbn-products och ogbn-papers100M, såväl som i inställningar för flera maskiner, med olika nivåer av fanout (mängd data som processorn skulle förbereda för GPU), och över flera arkitekturer, inklusive den senaste toppmoderna, GraphSAGE-RI. I varje inställning överträffade SALIENT PyTorch Geometric, framför allt på den stora ogbn-papers100M-datauppsättningen, som innehåller 100 miljoner noder och över en miljard kanter. Här var den tre gånger snabbare, körd på en GPU, än den optimerade baslinjen som ursprungligen skapades för detta jobb; med 16 GPU:er var SALIENT ytterligare åtta gånger snabbare. 

Medan andra system hade något annorlunda hårdvara och experimentella inställningar, så det var inte alltid en direkt jämförelse, överträffade SALIENT dem. Bland system som uppnådde liknande noggrannhet inkluderar representativa prestandasiffror 99 sekunder med en GPU och 32 processorer och 13 sekunder med 1,536 20 processorer. Däremot var SALIENTs körtid med en GPU och 16.5 processorer 16 sekunder och var bara två sekunder med 320 GPU:er och 16 processorer. "Om du tittar på siffrorna på den nedersta raden som tidigare arbete rapporterar, är vår XNUMX GPU-körtid (två sekunder) en storleksordning snabbare än andra siffror som tidigare har rapporterats på denna datamängd", säger Kaler. Forskarna tillskrev sina prestandaförbättringar, delvis, till deras tillvägagångssätt att optimera sin kod för en enda maskin innan de flyttade till den distribuerade inställningen. Stathas säger att lärdomen här är att för dina pengar är det "mer meningsfullt att använda hårdvaran du har effektivt, och till dess extrema, innan du börjar skala upp till flera datorer", vilket kan ge betydande besparingar på kostnader och koldioxidutsläpp som kan komma med modellutbildning.

Denna nya kapacitet kommer nu att tillåta forskare att ta itu med och gräva djupare in i allt större grafer. Till exempel innehöll Bitcoin-nätverket som nämndes tidigare 100,000 1,000 noder; SALIENT-systemet kan hantera en graf som är XNUMX XNUMX gånger (eller tre storleksordningar) större.

"I framtiden skulle vi titta på att inte bara köra detta grafiska neurala nätverksutbildningssystem på de befintliga algoritmerna som vi implementerade för att klassificera eller förutsäga egenskaperna för varje nod, utan vi vill också göra mer djupgående uppgifter, som t.ex. identifiera vanliga mönster i en graf (undergrafmönster), [vilket] faktiskt kan vara intressant för att indikera ekonomiska brott”, säger Chen. "Vi vill också identifiera noder i en graf som är lika i en mening att de möjligen skulle motsvara samma dåliga aktör i ett ekonomiskt brott. Dessa uppgifter skulle kräva utveckling av ytterligare algoritmer, och möjligen även neurala nätverksarkitekturer."

Denna forskning stöddes av MIT-IBM Watson AI Lab och delvis av US Air Force Research Laboratory och US Air Force Artificial Intelligence Accelerator.