Tesla vrea să ducă siliciul de învățare automată în Dojo

Pentru a potoli setea de modele din ce în ce mai mari de inteligență artificială și de învățare automată, Tesla a dezvăluit o mulțime de detalii la Hot Chips 34 cu privire la arhitectura lor de supercalculare complet personalizată numită Dojo.

Sistemul este în esență un supercomputer composabil masiv, deși spre deosebire de ceea ce vedem pe Top 500, este construit dintr-o arhitectură complet personalizată care cuprinde siliciul de calcul, rețea și intrare/ieșire (I/O) până la arhitectura set de instrucțiuni (ISA), livrarea energiei, ambalarea și răcirea. Toate acestea au fost realizate cu scopul expres de a rula algoritmi de antrenament de învățare automată personalizați și specifici la scară.

„Prelucrarea datelor din lumea reală este fezabilă doar prin tehnici de învățare automată, fie că este vorba de procesare în limbaj natural, de conducere pe străzi care sunt create pentru viziunea umană până la robotică care interfață cu mediul de zi cu zi”, a declarat Ganesh Venkataramanan, director senior de inginerie hardware la Tesla, în timpul discursului său principal.

Cu toate acestea, el a susținut că metodele tradiționale de scalare a sarcinilor de lucru distribuite nu au reușit să accelereze la ritmul necesar pentru a ține pasul cu cerințele învățării automate. De fapt, Legea lui Moore nu o reduce și nici sistemele nu sunt disponibile pentru antrenamentul AI/ML la scară, și anume o combinație de CPU/GPU sau, în circumstanțe mai rare, prin utilizarea acceleratoarelor AI de specialitate. 

„În mod tradițional, construim cipuri, le punem pe pachete, pachetele merg pe PCB-uri, care intră în sisteme. Sistemele intră în rafturi”, a spus Venkataramanan. Problema este că de fiecare dată când datele se mută de la cip la pachet și în afara pachetului, acestea implică o penalizare de latență și lățime de bandă.

Un sandviș de centru de date

Așa că pentru a ocoli limitările, Venkataramanan și echipa sa au început de la zero.

„De la interviul meu cu Elon, el m-a întrebat ce poți face, diferit de procesoarele și GPU-urile pentru AI. Simt că întreaga echipă încă răspunde la această întrebare.”

Acest lucru a condus la dezvoltarea plăcii de antrenament Dojo, un cluster de calcul autonom care ocupă o jumătate de picior cub capabil de 556 TFLOPS de performanță FP32 într-un pachet de 15 kW răcit cu lichid.

Fiecare placă este echipată cu 11 GB de SRAM și este conectată printr-o țesătură de 9 TB/s folosind un protocol de transport personalizat în întreaga stivă.

„Această piesă de antrenament reprezintă cantități inegalabile de integrare de la computer la memorie la livrarea energiei, la comunicare, fără a necesita comutatoare suplimentare”, a spus Venkataramanan.

În centrul plăcii de antrenament se află Tesla D1, o matriță de tranzistor de 50 de miliarde, bazată pe procesul de 7 nm al TSMC. Tesla spune că fiecare D1 este capabil de 22 TFLOPS de performanță FP32 la un TDP de 400 W. Cu toate acestea, Tesla observă că cipul este capabil să ruleze o gamă largă de calcule cu virgulă mobilă, inclusiv câteva personalizate.

„Dacă comparați tranzistorii pentru milimetri pătrați, aceasta este probabil marginea de sângerare a oricărui lucru care este acolo”, a spus Venkataramanan.

Apoi, Tesla a luat 25 de D1, le-a împărțit pentru matrițe bune cunoscute și apoi le-a împachetat folosind tehnologia system-on-wafer a TSMC pentru a „realiza o cantitate uriașă de integrare a calculului la latență foarte scăzută și lățime de bandă foarte mare”, a spus el.

Cu toate acestea, designul system-on-wafer și arhitectura stivuită vertical au introdus provocări în ceea ce privește livrarea energiei.

Potrivit lui Venkataramanan, majoritatea acceleratoarelor plasează astăzi puterea direct în apropierea siliciului. Și, deși este dovedită, această abordare înseamnă că o zonă mare a acceleratorului trebuie să fie dedicată acelor componente, ceea ce a făcut-o nepractică pentru Dojo, a explicat el. În schimb, Tesla și-a proiectat cipurile pentru a furniza putere direct prin partea de jos a matriței. 

Rezumând

„Am putea construi un întreg centru de date sau o clădire întreagă din această țiglă de antrenament, dar țiglă de antrenament este doar porțiunea de calcul. De asemenea, trebuie să-l hrănim”, a spus Venkataramanan.

Pentru aceasta, Tesla a dezvoltat și procesorul de interfață Dojo (DIP), care funcționează ca o punte între procesorul gazdă și procesoarele de antrenament. DIP servește, de asemenea, ca o sursă de memorie partajată cu lățime de bandă mare (HBM) și ca o NIC de mare viteză de 400 Gbit/sec.

Fiecare DIP dispune de 32 GB de HBM și până la cinci dintre aceste carduri pot fi conectate la o piesă de antrenament la 900 GB/s pentru un total de 4.5 TB/s la gazdă pentru un total de 160 GB de HBM pe piesă.

Perechile de configurație V1 ale Tesla ale acestor plăci – sau 150 de matrițe D1 – în matrice au suportat patru procesoare gazdă echipate fiecare cu cinci carduri DIP pentru a obține un exaflop revendicat de performanță BF16 sau CFP8.

Împreună, Venkataramanan spune că arhitectura – detaliată în profunzime aici by Următoarea platformă – permite Tesla să depășească limitările asociate cu acceleratoarele tradiționale precum Nvidia și AMD.

„Cum funcționează acceleratoarele tradiționale, de obicei încerci să potriviți un model întreg în fiecare accelerator. Replicați-l și apoi transmiteți datele prin fiecare dintre ele”, a spus el. „Ce se întâmplă dacă avem modele din ce în ce mai mari? Aceste acceleratoare se pot prăbuși pentru că rămân fără memorie.”

Aceasta nu este o problemă nouă, a remarcat el. Comutatorul NV de la Nvidia, de exemplu, permite ca memoria să fie grupată pe bănci mari de GPU-uri. Cu toate acestea, Venkataramanan susține că acest lucru nu numai că adaugă complexitate, dar introduce latență și compromisuri asupra lățimii de bandă.

„Ne-am gândit la asta chiar de la început. Plăcile noastre de calcul și fiecare matriță au fost făcute pentru a se potrivi modele mari”, a spus Venkataramanan.

Software

O astfel de arhitectură de calcul specializată necesită o stivă de software specializată. Cu toate acestea, Venkataramanan și echipa sa au recunoscut că programabilitatea fie va face, fie va distruge Dojo.

„Ușurința de programare pentru omologii software este primordială atunci când proiectăm aceste sisteme”, a spus el. „Cercetătorii nu vor aștepta ca oamenii tăi să scrie un nucleu scris de mână pentru adaptarea la un nou algoritm pe care vrem să-l rulăm.”

Pentru a face acest lucru, Tesla a renunțat la ideea de a folosi nuclee și a proiectat arhitectura Dojo în jurul compilatoarelor.

„Ceea ce am făcut a fost că am folosit PiTorch. Am creat un strat intermediar, care ne ajută să paralelizăm pentru a extinde hardware-ul de sub el. Sub totul este cod compilat”, a spus el. „Acesta este singurul mod de a crea stive de software care sunt adaptabile la toate aceste sarcini de lucru viitoare.”

În ciuda accentului pus pe flexibilitatea software-ului, Venkataramanan observă că platforma, care rulează în prezent în laboratoarele lor, este limitată la utilizarea Tesla pentru moment.

„Ne concentrăm mai întâi pe clienții noștri interni”, a spus el. „Elon a făcut public că, în timp, vom pune acest lucru la dispoziția cercetătorilor, dar nu avem un interval de timp pentru asta. ®