Tesla ønsker å ta maskinlæringssilisium til Dojoen

For å slukke tørsten etter stadig større AI- og maskinlæringsmodeller, har Tesla avslørt et vell av detaljer på Hot Chips 34 om deres helt tilpassede superdatabehandlingsarkitektur kalt Dojo.

Systemet er i hovedsak en massiv komponerbar superdatamaskin, men i motsetning til det vi ser på Top 500, den er bygget fra en helt tilpasset arkitektur som spenner over databehandling, nettverk og input/output (I/O) silisium til instruksjonssettarkitektur (ISA), strømforsyning, pakking og kjøling. Alt ble gjort med det uttrykkelige formålet å kjøre skreddersydde, spesifikke maskinlæringsalgoritmer i skala.

"Databehandling i den virkelige verden er bare mulig gjennom maskinlæringsteknikker, enten det er naturlig språkbehandling, kjøring i gater som er laget for menneskesyn til robotikk som kommuniserer med det daglige miljøet," sa Ganesh Venkataramanan, seniordirektør for maskinvareteknikk i Tesla. under hans hovedtale.

Han hevdet imidlertid at tradisjonelle metoder for å skalere distribuerte arbeidsbelastninger ikke har klart å akselerere med den hastigheten som er nødvendig for å holde tritt med kravene til maskinlæring. Faktisk er ikke Moores lov å kutte det, og heller ikke systemene er tilgjengelige for AI/ML-trening i stor skala, nemlig en kombinasjon av CPU/GPU eller under sjeldnere omstendigheter ved å bruke spesial-AI-akseleratorer. 

"Tradisjonelt bygger vi brikker, vi legger dem på pakker, pakker går på PCB, som går inn i systemer. Systemer går inn i stativer, sa Venkataramanan. Problemet er at hver gang data flyttes fra brikken til pakken og ut av pakken, pådrar det seg en ventetid og båndbreddestraff.

En datasentersandwich

Så for å omgå begrensningene startet Venkataramanan og teamet hans på nytt fra bunnen av.

"Helt fra intervjuet mitt med Elon spurte han meg hva du kan gjøre som er forskjellig fra CPUer og GPUer for AI. Jeg føler at hele teamet fortsatt svarer på det spørsmålet.»

Dette førte til utviklingen av Dojo-treningsflisen, en selvstendig dataklynge som opptar en halvkubikkfot som er i stand til 556 TFLOPS FP32-ytelse i en 15kW væskekjølt pakke.

Hver flis er utstyrt med 11 GB SRAM og er koblet over et 9TB/s-stoff ved hjelp av en tilpasset transportprotokoll gjennom hele stabelen.

"Denne treningsflisen representerer enestående mengder integrasjon fra datamaskin til minne til strømforsyning, til kommunikasjon, uten å kreve noen ekstra brytere," sa Venkataramanan.

I hjertet av treningsflisen er Teslas D1, en 50 milliarder transistordyse, basert på TSMCs 7nm-prosess. Tesla sier at hver D1 er i stand til 22 TFLOPS FP32-ytelse ved en TDP på ​​400W. Tesla bemerker imidlertid at brikken er i stand til å kjøre et bredt spekter av flyttallsberegninger, inkludert noen få tilpassede.

"Hvis du sammenligner transistorer for millimeter kvadrat, er dette sannsynligvis den blødende kanten til alt som er der ute," sa Venkataramanan.

Tesla tok deretter 25 D1-er, kasserte dem for kjente gode dies, og pakket dem deretter ved hjelp av TSMCs system-on-wafer-teknologi for å "oppnå en enorm mengde dataintegrasjon med svært lav ventetid og veldig høy båndbredde," sa han.

System-på-wafer-designet og den vertikalt stablede arkitekturen introduserte imidlertid utfordringer når det kom til strømforsyning.

I følge Venkataramanan plasserer de fleste akseleratorer i dag kraft rett ved siden av silisiumet. Og selv om denne tilnærmingen er bevist, betyr det at et stort område av akseleratoren må dedikeres til disse komponentene, noe som gjorde det upraktisk for Dojo, forklarte han. I stedet designet Tesla sjetongene deres for å levere strøm direkte gjennom bunnen av terningen. 

Sette alt sammen

"Vi kan bygge et helt datasenter eller en hel bygning av denne treningsflisen, men treningsflisen er bare beregningsdelen. Vi må også mate den, sa Venkataramanan.

For dette utviklet Tesla også Dojo Interface Processor (DIP), som fungerer som en bro mellom verts-CPU og treningsprosessorer. DIP-en fungerer også som en kilde for delt høybåndbredde-minne (HBM) og som en høyhastighets 400 Gbit/sek. NIC.

Hver DIP har 32 GB HBM og opptil fem av disse kortene kan kobles til en treningsbrikke med 900 GB/s for totalt 4.5 TB/s til verten for totalt 160 GB HBM per brikke.

Teslas V1-konfigurasjonspar av disse flisene – eller 150 D1-matriser – i array støttet fire verts-CPUer hver utstyrt med fem DIP-kort for å oppnå en påstått eksaflop av BF16- eller CFP8-ytelse.

Til sammen sier Venkataramanan arkitekturen – detaljert i dybde her by Neste plattform – gjør det mulig for Tesla å overvinne begrensningene knyttet til tradisjonelle akseleratorer fra slike som Nvidia og AMD.

"Hvordan tradisjonelle akseleratorer fungerer, prøver du vanligvis å passe en hel modell inn i hver akselerator. Repliker det, og flyt deretter dataene gjennom hver av dem," sa han. «Hva skjer hvis vi har større og større modeller? Disse akseleratorene kan falle flatt fordi de går tom for minne.»

Dette er ikke et nytt problem, bemerket han. Nvidias NV-switch gjør det for eksempel mulig å samle minne på tvers av store GPU-er. Venkataramanan hevder imidlertid at dette ikke bare gir kompleksitet, men introduserer latens og kompromisser med båndbredde.

– Vi har tenkt på dette helt fra starten. Dataflisene våre og hver av diesene ble laget for å passe store modeller, sa Venkataramanan.

Software

En slik spesialisert dataarkitektur krever en spesialisert programvarestabel. Venkataramanan og teamet hans erkjente imidlertid at programmerbarhet enten ville gjøre eller ødelegge Dojo.

"Enkel programmerbarhet for programvaremotparter er avgjørende når vi designer disse systemene," sa han. "Forskere vil ikke vente på at programvarefolkene dine skal skrive en håndskrevet kjerne for å tilpasse seg en ny algoritme som vi ønsker å kjøre."

For å gjøre dette, droppet Tesla ideen om å bruke kjerner, og designet Dojos arkitektur rundt kompilatorer.

"Det vi gjorde var at vi brukte PiTorch. Vi opprettet et mellomlag, som hjelper oss å parallellisere for å skalere ut maskinvare under det. Under alt er kompilert kode," sa han. "Dette er den eneste måten å lage programvarestabler som kan tilpasses alle fremtidige arbeidsbelastninger."

Til tross for vektleggingen av programvarefleksibilitet, bemerker Venkataramanan at plattformen, som for tiden kjører i laboratoriene deres, er begrenset til bruk av Tesla foreløpig.

"Vi er først og fremst fokusert på våre interne kunder," sa han. «Elon har gjort det offentlig at vi over tid vil gjøre dette tilgjengelig for forskere, men vi har ingen tidsramme for det. ®