Travle GPUer: Sampling og pipelining-metoden fremskynder dyp læring på store grafer

Grafer, et potensielt omfattende nett av noder forbundet med kanter, kan brukes til å uttrykke og undersøke forhold mellom data, som sosiale forbindelser, økonomiske transaksjoner, trafikk, energinett og molekylære interaksjoner. Etter hvert som forskere samler inn mer data og bygger ut disse grafiske bildene, vil forskerne trenge raskere og mer effektive metoder, samt mer beregningskraft, for å gjennomføre dyp læring på dem, i form av grafiske nevrale nettverk (GNN).  

Nå forbedrer en ny metode, kalt SALIENT (Sampling, sLIcing, and data movemeNT), utviklet av forskere ved MIT og IBM Research, trenings- og slutningsytelsen ved å adressere tre viktige flaskehalser i beregningen. Dette kutter dramatisk ned på kjøretiden til GNN-er på store datasett, som for eksempel inneholder på skalaen 100 millioner noder og 1 milliard kanter. Videre fant teamet at teknikken kan skaleres godt når beregningskraft legges til fra én til 16 grafiske prosesseringsenheter (GPUer). Arbeidet ble presentert på den femte konferansen om maskinlæring og systemer.

"Vi begynte å se på utfordringene dagens systemer opplevde når de skalere toppmoderne maskinlæringsteknikker for grafer til virkelig store datasett. Det viste seg at det var mye arbeid å gjøre, fordi mange av de eksisterende systemene oppnådde god ytelse, primært på mindre datasett som passet inn i GPU-minnet, sier Tim Kaler, hovedforfatter og postdoc i MIT Computer Science og Artificial Intelligence Laboratory (CSAIL).

Med enorme datasett mener eksperter skalaer som hele Bitcoin-nettverket, der visse mønstre og dataforhold kan stave ut trender eller stygt spill. "Det er nesten en milliard Bitcoin-transaksjoner på blokkjeden, og hvis vi ønsker å identifisere ulovlige aktiviteter i et slikt felles nettverk, så står vi overfor en graf av en slik skala," sier medforfatter Jie Chen, seniorforsker og leder fra IBM Research og MIT-IBM Watson AI Lab. "Vi ønsker å bygge et system som er i stand til å håndtere den typen grafer og lar behandlingen være så effektiv som mulig, fordi vi hver dag ønsker å følge med på tempoet til de nye dataene som genereres."

Kaler og Chens medforfattere inkluderer Nickolas Stathas MEng '21 fra Jump Trading, som utviklet SALIENT som en del av sitt doktorgradsarbeid; tidligere MIT-IBM Watson AI Lab praktikant og MIT graduate student Anne Ouyang; MIT CSAIL postdoc Alexandros-Stavros Iliopoulos; MIT CSAIL forsker Tao B. Schardl; og Charles E. Leiserson, Edwin Sibley Webster professor i elektroteknikk ved MIT og en forsker ved MIT-IBM Watson AI Lab.     

For dette problemet tok teamet en systemorientert tilnærming i utviklingen av metoden deres: SALIENT, sier Kaler. For å gjøre dette implementerte forskerne det de så som viktige, grunnleggende optimaliseringer av komponenter som passer inn i eksisterende maskinlæringsrammeverk, som PyTorch Geometric og deep graph library (DGL), som er grensesnitt for å bygge en maskinlæringsmodell. Stathas sier at prosessen er som å bytte ut motorer for å bygge en raskere bil. Metoden deres ble designet for å passe inn i eksisterende GNN-arkitekturer, slik at domeneeksperter enkelt kunne bruke dette arbeidet til sine spesifiserte felt for å fremskynde modelltrening og pirre ut innsikt under slutning raskere. Trikset, bestemte teamet, var å holde all maskinvaren (CPU-er, datakoblinger og GPU-er) opptatt til enhver tid: mens CPU-en prøver grafen og forbereder mini-batcher med data som deretter vil bli overført gjennom datakoblingen , jobber den mer kritiske GPU-en for å trene maskinlæringsmodellen eller utføre slutninger. 

Forskerne begynte med å analysere ytelsen til et vanlig brukt maskinlæringsbibliotek for GNN-er (PyTorch Geometric), som viste en oppsiktsvekkende lav utnyttelse av tilgjengelige GPU-ressurser. Ved å bruke enkle optimaliseringer forbedret forskerne GPU-utnyttelsen fra 10 til 30 prosent, noe som resulterte i en ytelsesforbedring på 1.4 til to ganger i forhold til offentlige benchmark-koder. Denne raske grunnlinjekoden kunne utføre en fullstendig pass over et stort treningsdatasett gjennom algoritmen (en epoke) på 50.4 sekunder.                          

For å søke ytterligere ytelsesforbedringer, forsøkte forskerne å undersøke flaskehalsene som oppstår i begynnelsen av datapipelinen: algoritmene for grafprøvetaking og forberedelse av minibatch. I motsetning til andre nevrale nettverk, utfører GNN-er en nabolagsaggregering, som beregner informasjon om en node ved å bruke informasjon som er tilstede i andre nærliggende noder i grafen - for eksempel i en sosial nettverksgraf, informasjon fra venner av venner til en bruker. Ettersom antall lag i GNN øker, kan antallet noder nettverket må nå ut til for informasjon eksplodere, og overskride grensene for en datamaskin. Algoritmer for sampling av nabolag hjelper ved å velge et mindre tilfeldig delsett av noder som skal samles; forskerne fant imidlertid at nåværende implementeringer av dette var for trege til å holde tritt med prosesseringshastigheten til moderne GPUer. Som svar identifiserte de en blanding av datastrukturer, algoritmiske optimaliseringer og så videre som forbedret samplingshastigheten, og til slutt forbedret samplingsoperasjonen alene med omtrent tre ganger, og tok kjøretiden per epoke fra 50.4 til 34.6 sekunder. De fant også at prøvetaking, med en passende hastighet, kan gjøres under inferens, og forbedrer den generelle energieffektiviteten og ytelsen, et punkt som hadde blitt oversett i litteraturen, bemerker teamet.      

I tidligere systemer var dette prøvetakingstrinnet en multi-prosess tilnærming, og skapte ekstra data og unødvendig databevegelse mellom prosessene. Forskerne gjorde SALIENT-metoden deres mer kvikk ved å lage en enkelt prosess med lette tråder som holdt dataene på CPU-en i delt minne. Videre drar SALIENT nytte av en hurtigbuffer med moderne prosessorer, sier Stathas, og parallelliserer funksjonsslicing, som trekker ut relevant informasjon fra noder av interesse og deres omkringliggende naboer og kanter, innenfor det delte minnet til CPU-kjernebufferen. Dette reduserte igjen den totale kjøretiden per epoke fra 34.6 til 27.8 sekunder.

Den siste flaskehalsen forskerne tok for seg var å pipeline mini-batch-dataoverføringer mellom CPU og GPU ved å bruke et forhåndshentingstrinn, som ville forberede data rett før det er nødvendig. Teamet beregnet at dette ville maksimere båndbreddebruken i datalinken og bringe metoden opp til perfekt utnyttelse; imidlertid så de bare rundt 90 prosent. De identifiserte og fikset en ytelsesfeil i et populært PyTorch-bibliotek som forårsaket unødvendig rundturskommunikasjon mellom CPU og GPU. Med denne feilen fikset, oppnådde teamet en 16.5 sekunders kjøretid per epoke med SALIENT.

"Vårt arbeid viste, tror jeg, at djevelen er i detaljene," sier Kaler. "Når du følger nøye med på detaljene som påvirker ytelsen når du trener et grafisk nevralt nettverk, kan du løse et stort antall ytelsesproblemer. Med løsningene våre endte vi opp med å bli fullstendig flaskehalset av GPU-beregning, som er det ideelle målet for et slikt system.»

SALIENTs hastighet ble evaluert på tre standard datasett ogbn-arxiv, ogbn-products og ogbn-papers100M, så vel som i innstillinger for flere maskiner, med forskjellige nivåer av fanout (mengde data som CPUen ville forberede for GPU), og på tvers av flere arkitekturer, inkludert den nyeste toppmoderne, GraphSAGE-RI. I hver innstilling overgikk SALIENT PyTorch Geometric, spesielt på det store ogbn-papers100M-datasettet, som inneholder 100 millioner noder og over en milliard kanter. Her var det tre ganger raskere, kjører på én GPU, enn den optimaliserte grunnlinjen som opprinnelig ble opprettet for denne jobben; med 16 GPUer var SALIENT ytterligere åtte ganger raskere. 

Mens andre systemer hadde litt annerledes maskinvare og eksperimentelle oppsett, så det var ikke alltid en direkte sammenligning, utkonkurrerte SALIENT dem. Blant systemer som oppnådde lignende nøyaktighet, inkluderer representative ytelsestall 99 sekunder ved bruk av én GPU og 32 CPUer, og 13 sekunder ved bruk av 1,536 CPUer. Derimot var SALIENTs kjøretid ved bruk av én GPU og 20 CPUer 16.5 sekunder og var bare to sekunder med 16 GPUer og 320 CPUer. "Hvis du ser på tallene på bunnlinjen som tidligere arbeid rapporterer, er vår 16 GPU-kjøringstid (to sekunder) en størrelsesorden raskere enn andre tall som har blitt rapportert tidligere på dette datasettet," sier Kaler. Forskerne tilskrev ytelsesforbedringene deres, delvis til deres tilnærming til å optimalisere koden for en enkelt maskin før de flyttet til den distribuerte innstillingen. Stathas sier at lærdommen her er at for pengene dine, "gir det mer mening å bruke maskinvaren du har effektivt, og til det ytterste, før du begynner å skalere opp til flere datamaskiner," som kan gi betydelige besparelser på kostnader og karbonutslipp som kan følge med modelltrening.

Denne nye kapasiteten vil nå tillate forskere å takle og grave dypere inn i større og større grafer. For eksempel inneholdt Bitcoin-nettverket som ble nevnt tidligere 100,000 1,000 noder; SALIENT-systemet kan håndtere en graf som er XNUMX ganger (eller tre størrelsesordener) større.

"I fremtiden vil vi se på ikke bare å kjøre dette grafiske nevrale nettverksopplæringssystemet på de eksisterende algoritmene som vi implementerte for å klassifisere eller forutsi egenskapene til hver node, men vi ønsker også å gjøre mer dyptgående oppgaver, som f.eks. identifisere vanlige mønstre i en graf (undergrafmønstre), [som] faktisk kan være interessant for å indikere økonomisk kriminalitet, sier Chen. "Vi ønsker også å identifisere noder i en graf som er like i en forstand at de muligens ville tilsvare den samme dårlige aktøren i en økonomisk kriminalitet. Disse oppgavene vil kreve utvikling av ytterligere algoritmer, og muligens også nevrale nettverksarkitekturer.

Denne forskningen ble støttet av MIT-IBM Watson AI Lab og delvis av US Air Force Research Laboratory og US Air Force Artificial Intelligence Accelerator.