Tesla ønsker at tage maskinlæringssilicium til Dojo'en

For at slukke tørsten efter stadig større AI- og maskinlæringsmodeller har Tesla afsløret et væld af detaljer på Hot Chips 34 om deres fuldt tilpassede supercomputing-arkitektur kaldet Dojo.

Systemet er i det væsentlige en massiv komponerbar supercomputer, men i modsætning til det, vi ser på Top 500, den er bygget fra en helt tilpasset arkitektur, der spænder over databehandling, netværk og input/output (I/O) silicium til instruktionssætarkitektur (ISA), strømforsyning, pakning og køling. Det hele blev gjort med det udtrykkelige formål at køre skræddersyede, specifikke maskinlærings-træningsalgoritmer i skala.

"Den virkelige verden databehandling er kun mulig gennem maskinlæringsteknikker, hvad enten det er naturligt sprogbehandling, kørsel i gader, der er skabt til menneskesyn til robotteknologi, der interfacer med det daglige miljø," sagde Ganesh Venkataramanan, seniordirektør for hardware engineering hos Tesla. under hans hovedtale.

Han argumenterede dog for, at traditionelle metoder til at skalere distribuerede arbejdsbelastninger ikke har formået at accelerere med den hastighed, der er nødvendig for at holde trit med maskinlæringskravene. Faktisk skærer Moores lov ikke det, og systemerne er heller ikke tilgængelige til AI/ML-træning i skala, nemlig en kombination af CPU/GPU eller under sjældnere omstændigheder ved at bruge specielle AI-acceleratorer. 

"Traditionelt bygger vi chips, vi sætter dem på pakker, pakker går på PCB'er, som går ind i systemer. Systemer går ind i stativer,” sagde Venkataramanan. Problemet er, at hver gang data flyttes fra chippen til pakken og ud af pakken, medfører det en latens- og båndbreddestraf.

En datacentersandwich

Så for at komme uden om begrænsningerne startede Venkataramanan og hans team forfra fra bunden.

"Lige fra mit interview med Elon spurgte han mig, hvad du kan gøre, der er anderledes end CPU'er og GPU'er til AI. Jeg føler, at hele holdet stadig svarer på det spørgsmål."

Dette førte til udviklingen af ​​Dojo-træningsflisen, en selvstændig computerklynge, der optager en halv kubikfod, der er i stand til 556 TFLOPS FP32-ydelse i en 15 kW væskekølet pakke.

Hver flise er udstyret med 11 GB SRAM og er forbundet over et 9TB/s stof ved hjælp af en brugerdefineret transportprotokol gennem hele stakken.

"Denne træningsflise repræsenterer enestående mængder af integration fra computer til hukommelse til strømforsyning, til kommunikation uden at kræve yderligere kontakter," sagde Venkataramanan.

I hjertet af træningsflisen er Teslas D1, en 50 milliarder transistor-matrice, baseret på TSMC's 7nm-proces. Tesla siger, at hver D1 er i stand til 22 TFLOPS FP32-ydelse ved en TDP på ​​400W. Tesla bemærker dog, at chippen er i stand til at køre en bred vifte af flydende kommaberegninger, herunder nogle få brugerdefinerede.

"Hvis du sammenligner transistorer for millimeter kvadrat, er dette sandsynligvis den blødende kant af alt, hvad der er derude," sagde Venkataramanan.

Tesla tog derefter 25 D1'er, kasserede dem for kendte gode dies og pakkede dem derefter ved hjælp af TSMC's system-on-wafer-teknologi for at "opnå en enorm mængde computerintegration med meget lav latenstid og meget høj båndbredde," sagde han.

System-on-wafer-designet og den vertikalt stablede arkitektur introducerede imidlertid udfordringer, når det kom til strømforsyning.

Ifølge Venkataramanan placerer de fleste acceleratorer i dag kraft direkte ved siden af ​​silicium. Og selvom det er bevist, betyder denne tilgang, at et stort område af acceleratoren skal dedikeres til disse komponenter, hvilket gjorde det upraktisk for Dojo, forklarede han. I stedet designede Tesla deres chips til at levere strøm direkte gennem bunden af ​​formen. 

Samler det hele

"Vi kunne bygge et helt datacenter eller en hel bygning ud af denne træningsflise, men træningsflisen er kun beregningsdelen. Vi er også nødt til at fodre den,” sagde Venkataramanan.

Til dette udviklede Tesla også Dojo Interface Processor (DIP), der fungerer som en bro mellem værts-CPU'en og træningsprocessorer. DIP'en fungerer også som en kilde til delt hukommelse med høj båndbredde (HBM) og som et NIC med høj hastighed på 400 Gbit/sek.

Hver DIP har 32 GB HBM, og op til fem af disse kort kan forbindes til en træningsflise ved 900 GB/s for i alt 4.5 TB/s til værten for i alt 160 GB HBM pr.

Teslas V1-konfigurationspar af disse fliser – eller 150 D1-matricer – i array understøttede fire værts-CPU'er hver udstyret med fem DIP-kort for at opnå en påstået exaflop af BF16- eller CFP8-ydeevne.

Sammenlagt siger Venkataramanan arkitekturen - detaljeret i dybden her by Den næste platform – gør det muligt for Tesla at overvinde de begrænsninger, der er forbundet med traditionelle acceleratorer fra f.eks. Nvidia og AMD.

”Hvordan traditionelle acceleratorer fungerer, prøver man typisk at passe en hel model ind i hver accelerator. Repliker det, og flyt derefter dataene gennem hver af dem,” sagde han. ”Hvad sker der, hvis vi har større og større modeller? Disse acceleratorer kan falde fladt, fordi de løber tør for hukommelse."

Dette er ikke et nyt problem, bemærkede han. Nvidias NV-switch gør det for eksempel muligt at samle hukommelse på tværs af store banker af GPU'er. Venkataramanan hævder dog, at dette ikke kun tilføjer kompleksitet, men introducerer latens og går på kompromis med båndbredden.

"Vi tænkte på det her lige fra starten. Vores computerfliser og hver af matricerne blev lavet til at passe store modeller,” sagde Venkataramanan.

Software

En sådan specialiseret computerarkitektur kræver en specialiseret softwarestak. Venkataramanan og hans team erkendte dog, at programmerbarhed enten ville gøre eller ødelægge Dojo.

"Nem programmerbarhed for software-modparter er altafgørende, når vi designer disse systemer," sagde han. "Forskere vil ikke vente på, at dine softwarefolk skriver en håndskrevet kerne til tilpasning til en ny algoritme, som vi ønsker at køre."

For at gøre dette droppede Tesla ideen om at bruge kerner og designede Dojos arkitektur omkring compilere.

"Det, vi gjorde, var, at vi brugte PiTorch. Vi oprettede et mellemlag, som hjælper os med at parallelisere for at udskalere hardware under det. Under alt er kompileret kode," sagde han. "Dette er den eneste måde at skabe softwarestakke, der kan tilpasses til alle disse fremtidige arbejdsbelastninger."

På trods af vægten på softwarefleksibilitet bemærker Venkataramanan, at platformen, som i øjeblikket kører i deres laboratorier, er begrænset til Tesla-brug i øjeblikket.

"Vi fokuserer først på vores interne kunder," sagde han. ”Elon har gjort det offentligt, at vi over tid vil gøre dette tilgængeligt for forskere, men vi har ikke en tidsramme for det. ®