GPU-uri ocupate: metoda de eșantionare și pipelining accelerează învățarea profundă pe grafice mari

Graficele, o rețea potențial extinsă de noduri conectate prin margini, pot fi folosite pentru a exprima și interoga relațiile dintre date, cum ar fi conexiunile sociale, tranzacțiile financiare, traficul, rețelele energetice și interacțiunile moleculare. Pe măsură ce cercetătorii colectează mai multe date și construiesc aceste imagini grafice, cercetătorii vor avea nevoie de metode mai rapide și mai eficiente, precum și de mai multă putere de calcul, pentru a le desfășura învățarea profundă, în modul rețelelor neuronale grafice (GNN).  

Acum, o nouă metodă, numită SALIENT (Eșantionare, tăiere și mișcare a datelor), dezvoltată de cercetătorii de la MIT și IBM Research, îmbunătățește performanța de instruire și inferență abordând trei blocaje cheie în calcul. Acest lucru reduce drastic timpul de rulare a GNN-urilor pe seturi de date mari, care, de exemplu, conțin la scara de 100 de milioane de noduri și 1 miliard de margini. În plus, echipa a descoperit că tehnica se scalează bine atunci când se adaugă putere de calcul de la una la 16 unități de procesare grafică (GPU). Lucrarea a fost prezentată la a cincea conferință privind învățarea automată și sisteme.

„Am început să ne uităm la provocările cu care se confruntă sistemele actuale la scalarea tehnicilor de învățare automată de ultimă generație pentru grafice la seturi de date foarte mari. S-a dovedit că era mult de lucru de făcut, deoarece multe dintre sistemele existente obțin performanțe bune în primul rând pe seturi de date mai mici, care se potrivesc în memoria GPU”, spune Tim Kaler, autorul principal și postdoctorat în MIT Computer Science. și Laboratorul de Inteligență Artificială (CSAIL).

Prin seturi vaste de date, experții înțeleg scale precum întreaga rețea Bitcoin, în care anumite modele și relații de date ar putea explica tendințele sau jocul greșit. „Există aproape un miliard de tranzacții Bitcoin pe blockchain și, dacă dorim să identificăm activități ilicite în cadrul unei astfel de rețele comune, atunci ne confruntăm cu un grafic de o asemenea scară”, spune coautorul Jie Chen, cercetător și manager senior. al IBM Research și al MIT-IBM Watson AI Lab. „Vrem să construim un sistem care să fie capabil să gestioneze acest tip de grafic și să permită procesarea să fie cât mai eficientă posibil, pentru că în fiecare zi dorim să ținem pasul cu ritmul noilor date care sunt generate.”

Co-autorii lui Kaler și Chen includ Nickolas Stathas MEng '21 de la Jump Trading, care a dezvoltat SALIENT ca parte a lucrării sale de absolvent; fost stagiar MIT-IBM Watson AI Lab și studentă absolventă MIT Anne Ouyang; MIT CSAIL postdoc Alexandros-Stavros Iliopoulos; Cercetător MIT CSAIL Tao B. Schardl; și Charles E. Leiserson, Edwin Sibley Webster Profesor de Inginerie Electrică la MIT și cercetător la MIT-IBM Watson AI Lab.     

Pentru această problemă, echipa a adoptat o abordare orientată spre sisteme în dezvoltarea metodei lor: SALIENT, spune Kaler. Pentru a face acest lucru, cercetătorii au implementat ceea ce au considerat ca fiind importante, optimizări de bază ale componentelor care se potrivesc în cadrele de învățare automată existente, cum ar fi PyTorch Geometric și biblioteca de grafice profunde (DGL), care sunt interfețe pentru construirea unui model de învățare automată. Stathas spune că procesul este ca și cum ați schimba motoare pentru a construi o mașină mai rapidă. Metoda lor a fost concepută pentru a se potrivi în arhitecturile GNN existente, astfel încât experții în domeniu să poată aplica cu ușurință această muncă în domeniile lor specificate pentru a accelera instruirea modelului și a scoate mai repede informații în timpul inferenței. Trucul, a stabilit echipa, a fost să mențină tot hardware-ul (CPU-uri, legături de date și GPU-uri) ocupat în orice moment: în timp ce CPU-ul prelevănește graficul și pregătește mini-loturi de date care vor fi apoi transferate prin legătura de date. , cu atât GPU-ul mai critic lucrează pentru a antrena modelul de învățare automată sau pentru a efectua inferențe. 

Cercetătorii au început prin a analiza performanța unei biblioteci de învățare automată utilizată în mod obișnuit pentru GNN-uri (PyTorch Geometric), care a arătat o utilizare uimitor de scăzută a resurselor GPU disponibile. Aplicând optimizări simple, cercetătorii au îmbunătățit utilizarea GPU-ului de la 10 la 30%, rezultând o îmbunătățire a performanței de 1.4 până la de două ori față de codurile de referință publice. Acest cod de bază rapid ar putea executa o trecere completă peste un set mare de date de antrenament prin algoritm (o epocă) în 50.4 secunde.                          

Căutând îmbunătățiri suplimentare de performanță, cercetătorii și-au propus să examineze blocajele care apar la începutul conductei de date: algoritmii pentru eșantionarea graficelor și pregătirea mini-loturi. Spre deosebire de alte rețele neuronale, GNN-urile efectuează o operațiune de agregare de vecinătate, care calculează informații despre un nod folosind informațiile prezente în alte noduri din apropiere din grafic - de exemplu, într-un grafic al unei rețele sociale, informații de la prietenii prietenilor unui utilizator. Pe măsură ce numărul de straturi din GNN crește, numărul de noduri la care trebuie să ajungă rețeaua pentru informații poate exploda, depășind limitele unui computer. Algoritmii de eșantionare de vecinătate ajută prin selectarea unui subset aleatoriu mai mic de noduri de adunat; cu toate acestea, cercetătorii au descoperit că implementările actuale ale acestui lucru au fost prea lente pentru a ține pasul cu viteza de procesare a GPU-urilor moderne. Ca răspuns, ei au identificat un amestec de structuri de date, optimizări algoritmice și așa mai departe, care a îmbunătățit viteza de eșantionare, îmbunătățind în cele din urmă doar operațiunea de eșantionare de aproximativ trei ori, ducând timpul de rulare pe epocă de la 50.4 la 34.6 secunde. Ei au descoperit, de asemenea, că eșantionarea, la o rată adecvată, poate fi făcută în timpul inferenței, îmbunătățind eficiența energetică și performanța generală, un punct care a fost trecut cu vederea în literatură, notează echipa.      

În sistemele anterioare, acest pas de eșantionare a fost o abordare multi-proces, creând date suplimentare și mișcare inutilă a datelor între procese. Cercetătorii și-au făcut metoda SALIENT mai agilă prin crearea unui singur proces cu fire ușoare care păstrau datele de pe CPU în memoria partajată. În plus, SALIENT profită de un cache de procesoare moderne, spune Stathas, paralelizând funcția de tăiere, care extrage informații relevante din nodurile de interes și din vecinii și marginile lor din jur, în memoria partajată a memoriei cache-ului central al procesorului. Acest lucru a redus din nou durata generală de rulare pe epocă de la 34.6 la 27.8 secunde.

Ultimul blocaj abordat de cercetători a fost să canalizeze transferurile de date în mini-loturi între CPU și GPU folosind un pas de preluare prealabilă, care ar pregăti datele chiar înainte de a fi necesare. Echipa a calculat că aceasta ar maximiza utilizarea lățimii de bandă în legătura de date și ar aduce metoda la o utilizare perfectă; cu toate acestea, au văzut doar aproximativ 90 la sută. Ei au identificat și remediat o eroare de performanță într-o bibliotecă populară PyTorch care a cauzat comunicații inutile dus-întors între CPU și GPU. Cu această eroare remediată, echipa a atins un timp de rulare de 16.5 secunde pe epocă cu SALIENT.

„Cred că munca noastră a arătat că diavolul este în detalii”, spune Kaler. „Când acordați o atenție deosebită detaliilor care influențează performanța atunci când antrenați o rețea neuronală grafică, puteți rezolva un număr mare de probleme de performanță. Cu soluțiile noastre, am ajuns să fim complet blocați de calculul GPU, care este obiectivul ideal al unui astfel de sistem.”

Viteza lui SALIENT a fost evaluată pe trei seturi de date standard ogbn-arxiv, ogbn-products și ogbn-papers100M, precum și în setări multi-mașină, cu diferite niveluri de fanout (cantitatea de date pe care CPU-ul ar pregăti-o pentru GPU) și în mai multe arhitecturi, inclusiv cea mai recentă, GraphSAGE-RI. În fiecare setare, SALIENT a depășit performanța PyTorch Geometric, în special pe setul mare de date ogbn-papers100M, care conține 100 de milioane de noduri și peste un miliard de margini Aici, a fost de trei ori mai rapid, rulând pe un singur GPU, decât linia de bază optimizată care a fost creată inițial pentru acest lucru; cu 16 GPU-uri, SALIENT a fost de opt ori mai rapid. 

În timp ce alte sisteme aveau hardware și configurații experimentale ușor diferite, deci nu a fost întotdeauna o comparație directă, SALIENT le-a depășit totuși. Printre sistemele care au obținut o precizie similară, cifrele reprezentative de performanță includ 99 de secunde folosind un GPU și 32 de procesoare și 13 secunde folosind 1,536 de procesoare. În schimb, timpul de rulare al lui SALIENT folosind un GPU și 20 de procesoare a fost de 16.5 secunde și a fost de doar două secunde cu 16 GPU și 320 de procesoare. „Dacă te uiți la cifrele de bază pe care le raportează lucrările anterioare, timpul nostru de rulare a 16 GPU (două secunde) este cu un ordin de mărime mai rapid decât alte numere care au fost raportate anterior pe acest set de date”, spune Kaler. Cercetătorii și-au atribuit îmbunătățirile de performanță, în parte, abordării lor de a-și optimiza codul pentru o singură mașină înainte de a trece la setarea distribuită. Stathas spune că lecția aici este că, pentru banii tăi, „este mai logic să folosești hardware-ul pe care îl ai în mod eficient și până la extrem, înainte de a începe extinderea la mai multe computere”, ceea ce poate oferi economii semnificative la costuri și la emisiile de carbon. care poate veni cu antrenamentul modelului.

Această nouă capacitate va permite acum cercetătorilor să abordeze și să sape mai adânc în grafice din ce în ce mai mari. De exemplu, rețeaua Bitcoin care a fost menționată mai devreme conținea 100,000 de noduri; sistemul SALIENT poate gestiona cu capacitate un grafic de 1,000 de ori (sau de trei ordine de mărime) mai mare.

„În viitor, ne-am uita nu doar să rulăm acest sistem de antrenament al rețelei neuronale grafice pe algoritmii existenți pe care i-am implementat pentru clasificarea sau prezicerea proprietăților fiecărui nod, dar dorim și să facem sarcini mai aprofundate, cum ar fi identificarea tiparelor comune într-un grafic (modele subgraf), [care] poate fi de fapt interesantă pentru a indica infracțiuni financiare”, spune Chen. „Vrem, de asemenea, să identificăm noduri într-un grafic care sunt similare în sensul că ar putea corespunde aceluiași actor rău într-o crimă financiară. Aceste sarcini ar necesita dezvoltarea unor algoritmi suplimentari și, posibil, de asemenea, arhitecturi de rețele neuronale.”

Această cercetare a fost susținută de Laboratorul MIT-IBM Watson AI și, în parte, de Laboratorul de Cercetare al Forțelor Aeriene din SUA și Acceleratorul de Inteligență Artificială a Forțelor Aeriene din SUA.