Travle GPU'er: Sampling og pipelining-metode fremskynder dyb læring på store grafer

Grafer, et potentielt omfattende net af knudepunkter forbundet af kanter, kan bruges til at udtrykke og udspørge relationer mellem data, såsom sociale forbindelser, finansielle transaktioner, trafik, energinet og molekylære interaktioner. Efterhånden som forskere indsamler flere data og bygger disse grafiske billeder ud, vil forskerne have brug for hurtigere og mere effektive metoder samt mere beregningskraft til at udføre dyb læring på dem i form af grafiske neurale netværk (GNN).  

Nu forbedrer en ny metode, kaldet SALIENT (Sampling, sLIcing, and data movemeNT), udviklet af forskere ved MIT og IBM Research, trænings- og inferensydelsen ved at adressere tre vigtige flaskehalse i beregningen. Dette skærer dramatisk ned på køretiden for GNN'er på store datasæt, som for eksempel indeholder på skalaen 100 millioner noder og 1 milliard kanter. Yderligere fandt teamet ud af, at teknikken kan skaleres godt, når beregningskraft tilføjes fra én til 16 grafiske behandlingsenheder (GPU'er). Arbejdet blev præsenteret på den femte konference om maskinlæring og systemer.

"Vi begyndte at se på de udfordringer, de nuværende systemer oplevede, når de skalering af state-of-the-art maskinlæringsteknikker for grafer til virkelig store datasæt. Det viste sig, at der var meget arbejde at gøre, fordi mange af de eksisterende systemer opnåede god ydeevne primært på mindre datasæt, der passer ind i GPU-hukommelsen,” siger Tim Kaler, hovedforfatter og postdoc i MIT Computer Science og Artificial Intelligence Laboratory (CSAIL).

Med store datasæt mener eksperter skalaer som hele Bitcoin-netværket, hvor visse mønstre og datarelationer kunne udskrive tendenser eller urimeligt spil. "Der er næsten en milliard Bitcoin-transaktioner på blockchain, og hvis vi ønsker at identificere ulovlige aktiviteter inde i et sådant fælles netværk, så står vi over for en graf af en sådan skala," siger medforfatter Jie Chen, seniorforsker og leder af IBM Research og MIT-IBM Watson AI Lab. "Vi vil bygge et system, der er i stand til at håndtere den slags grafer og tillader behandlingen at være så effektiv som muligt, fordi vi hver dag vil følge med tempoet i de nye data, der genereres."

Kaler og Chens medforfattere omfatter Nickolas Stathas MEng '21 fra Jump Trading, som udviklede SALIENT som en del af sit kandidatarbejde; tidligere MIT-IBM Watson AI Lab praktikant og MIT kandidatstuderende Anne Ouyang; MIT CSAIL postdoc Alexandros-Stavros Iliopoulos; MIT CSAIL-forsker Tao B. Schardl; og Charles E. Leiserson, Edwin Sibley Webster professor i elektroteknik ved MIT og en forsker ved MIT-IBM Watson AI Lab.     

Til dette problem tog teamet en systemorienteret tilgang til at udvikle deres metode: SALIENT, siger Kaler. For at gøre dette implementerede forskerne, hvad de så som vigtige, grundlæggende optimeringer af komponenter, der passer ind i eksisterende maskinlæringsrammer, såsom PyTorch Geometric og deep graph library (DGL), som er grænseflader til at bygge en maskinlæringsmodel. Stathas siger, at processen er som at skifte motorer ud for at bygge en hurtigere bil. Deres metode blev designet til at passe ind i eksisterende GNN-arkitekturer, så domæneeksperter nemt kunne anvende dette arbejde på deres specificerede felter for at fremskynde modeltræning og pirre indsigt under inferens hurtigere. Tricket, bestemte teamet, var at holde al hardware (CPU'er, datalinks og GPU'er) beskæftiget på alle tidspunkter: mens CPU'en prøver grafen og forbereder mini-batches af data, som derefter vil blive overført via datalinket , arbejder den mere kritiske GPU på at træne maskinlæringsmodellen eller udføre konklusioner. 

Forskerne begyndte med at analysere ydeevnen af ​​et almindeligt brugt maskinlæringsbibliotek for GNN'er (PyTorch Geometric), som viste en forbløffende lav udnyttelse af tilgængelige GPU-ressourcer. Ved at anvende simple optimeringer forbedrede forskerne GPU-udnyttelsen fra 10 til 30 procent, hvilket resulterede i en 1.4 til to gange ydeevneforbedring i forhold til offentlige benchmark-koder. Denne hurtige basislinjekode kunne udføre en komplet passage over et stort træningsdatasæt gennem algoritmen (en epoke) på 50.4 sekunder.                          

For at søge yderligere præstationsforbedringer satte forskerne sig for at undersøge de flaskehalse, der opstår i begyndelsen af ​​datapipelinen: algoritmerne til grafprøvetagning og mini-batch-forberedelse. I modsætning til andre neurale netværk udfører GNN'er en nabolagsaggregeringsoperation, som beregner information om en knude ved hjælp af information, der findes i andre nærliggende knudepunkter i grafen - for eksempel i en social netværksgraf, information fra venner af en brugers venner. Efterhånden som antallet af lag i GNN stiger, kan antallet af noder, netværket skal nå ud til for at få information, eksplodere og overskride grænserne for en computer. Neighborhood sampling algoritmer hjælper ved at vælge en mindre tilfældig delmængde af noder til at indsamle; dog fandt forskerne ud af, at de nuværende implementeringer af dette var for langsomme til at holde trit med behandlingshastigheden af ​​moderne GPU'er. Som svar identificerede de en blanding af datastrukturer, algoritmiske optimeringer og så videre, der forbedrede samplingshastigheden, hvilket i sidste ende forbedrede samplingsoperationen alene med omkring tre gange, hvilket tog den per-epoke runtime fra 50.4 til 34.6 sekunder. De fandt også ud af, at prøvetagning, med en passende hastighed, kan udføres under inferens, hvilket forbedrer den overordnede energieffektivitet og ydeevne, et punkt, der var blevet overset i litteraturen, bemærker holdet.      

I tidligere systemer var dette prøveudtagningstrin en multi-proces tilgang, der skabte ekstra data og unødvendig databevægelse mellem processerne. Forskerne gjorde deres SALIENT-metode mere smidig ved at skabe en enkelt proces med letvægtstråde, der holdt dataene på CPU'en i delt hukommelse. Ydermere udnytter SALIENT en cache af moderne processorer, siger Stathas, der paralleliserer feature slicing, som udtrækker relevant information fra interesseknuder og deres omkringliggende naboer og kanter i CPU-kernecachens delte hukommelse. Dette reducerede igen den samlede per-epoke køretid fra 34.6 til 27.8 sekunder.

Den sidste flaskehals, som forskerne adresserede, var at pipeline mini-batch-dataoverførsler mellem CPU'en og GPU'en ved hjælp af et forhåndshentningstrin, som ville forberede data lige før det er nødvendigt. Holdet beregnede, at dette ville maksimere båndbreddeforbruget i datalinket og bringe metoden op til perfekt udnyttelse; dog så de kun omkring 90 pct. De identificerede og rettede en ydeevnefejl i et populært PyTorch-bibliotek, der forårsagede unødvendig rundturskommunikation mellem CPU'en og GPU'en. Med denne fejl rettet opnåede teamet en 16.5 sekunders runtime pr. epoke med SALIENT.

"Vores arbejde viste, tror jeg, at djævelen er i detaljerne," siger Kaler. "Når du er meget opmærksom på de detaljer, der påvirker ydeevnen, når du træner et grafisk neuralt netværk, kan du løse et stort antal præstationsproblemer. Med vores løsninger endte vi med at blive fuldstændig flaskehalsede af GPU-beregning, som er det ideelle mål for et sådant system."

SALIENTs hastighed blev evalueret på tre standarddatasæt ogbn-arxiv, ogbn-products og ogbn-papers100M, såvel som i multi-maskine indstillinger, med forskellige niveauer af fanout (mængde af data, som CPU'en ville forberede til GPU'en), og på tværs af adskillige arkitekturer, herunder den seneste state-of-the-art, GraphSAGE-RI. I hver indstilling klarede SALIENT sig bedre end PyTorch Geometric, især på det store ogbn-papers100M-datasæt, der indeholder 100 millioner noder og over en milliard kanter. dette arbejde; med 16 GPU'er var SALIENT yderligere otte gange hurtigere. 

Mens andre systemer havde lidt anderledes hardware og eksperimentelle opsætninger, så det ikke altid var en direkte sammenligning, overgik SALIENT dem stadig. Blandt systemer, der opnåede lignende nøjagtighed, inkluderer repræsentative ydeevnetal 99 sekunder ved brug af én GPU og 32 CPU'er og 13 sekunder ved brug af 1,536 CPU'er. I modsætning hertil var SALIENTs køretid ved brug af en GPU og 20 CPU'er 16.5 sekunder og var kun to sekunder med 16 GPU'er og 320 CPU'er. "Hvis du ser på de bundlinjetal, som tidligere arbejde rapporterer, er vores 16 GPU-runtid (to sekunder) en størrelsesorden hurtigere end andre tal, der tidligere er blevet rapporteret på dette datasæt," siger Kaler. Forskerne tilskrev deres præstationsforbedringer til dels deres tilgang til at optimere deres kode til en enkelt maskine, før de flyttede til den distribuerede indstilling. Stathas siger, at lektionen her er, at for dine penge "giver det mere mening at bruge den hardware, du har, effektivt og til dets ekstreme, før du begynder at skalere op til flere computere," hvilket kan give betydelige besparelser på omkostninger og kulstofemissioner der kan komme med modeltræning.

Denne nye kapacitet vil nu give forskere mulighed for at tackle og grave dybere ned i større og større grafer. For eksempel indeholdt Bitcoin-netværket, der blev nævnt tidligere, 100,000 noder; SALIENT-systemet kan håndtere en graf 1,000 gange (eller tre størrelsesordener) større.

"I fremtiden vil vi se på ikke kun at køre dette grafiske neurale netværkstræningssystem på de eksisterende algoritmer, som vi implementerede til at klassificere eller forudsige egenskaberne for hver knude, men vi vil også lave mere dybtgående opgaver, som f.eks. at identificere almindelige mønstre i en graf (undergrafmønstre), [som] faktisk kan være interessant til at indikere økonomisk kriminalitet,” siger Chen. ”Vi ønsker også at identificere noder i en graf, der ligner hinanden i en forstand, at de muligvis ville svare til den samme dårlige aktør i en økonomisk kriminalitet. Disse opgaver ville kræve udvikling af yderligere algoritmer og muligvis også neurale netværksarkitekturer."

Denne forskning blev støttet af MIT-IBM Watson AI Lab og til dels af US Air Force Research Laboratory og US Air Force Artificial Intelligence Accelerator.