Tesla wil machine learning-silicium naar de Dojo . brengen

Om de dorst naar steeds grotere AI- en machine learning-modellen te lessen, heeft Tesla op Hot Chips 34 een schat aan details onthuld over hun volledig aangepaste supercomputing-architectuur genaamd Dojo.

Het systeem is in wezen een enorme, samen te stellen supercomputer, hoewel in tegenstelling tot wat we op de Top 500, het is opgebouwd uit een volledig op maat gemaakte architectuur die het compute-, netwerk- en input/output (I/O)-silicium omvat tot instructiesetarchitectuur (ISA), stroomvoorziening, verpakking en koeling. Dit alles werd gedaan met het uitdrukkelijke doel om op maat gemaakte, specifieke trainingsalgoritmen voor machine learning op grote schaal uit te voeren.

"Gegevensverwerking in de echte wereld is alleen mogelijk door middel van machine learning-technieken, of het nu gaat om natuurlijke taalverwerking, rijden in straten die zijn gemaakt voor menselijk zicht tot robotica die interageert met de dagelijkse omgeving", zegt Ganesh Venkataramanan, senior director hardware engineering bij Tesla. tijdens zijn keynote speech.

Hij voerde echter aan dat traditionele methoden voor het schalen van gedistribueerde workloads niet zijn versneld in het tempo dat nodig is om aan de eisen van machine learning te voldoen. In feite snijdt de wet van Moore het niet en zijn de systemen ook niet beschikbaar voor AI/ML-training op grote schaal, namelijk een combinatie van CPU/GPU of in zeldzamere omstandigheden door speciale AI-versnellers te gebruiken. 

“Traditioneel bouwen we chips, we zetten ze op verpakkingen, verpakkingen gaan op PCB's, die in systemen gaan. Systemen gaan in racks”, zegt Venkataramanan. Het probleem is dat elke keer dat gegevens van de chip naar het pakket en van het pakket worden verplaatst, dit een vertraging en bandbreedteverlies oplevert.

Een datacenter sandwich

Dus om de beperkingen te omzeilen, begonnen Venkataramanan en zijn team helemaal opnieuw.

"Vanaf mijn interview met Elon vroeg hij me wat je kunt doen dat anders is dan CPU's en GPU's voor AI. Ik heb het gevoel dat het hele team die vraag nog steeds beantwoordt."

Dit leidde tot de ontwikkeling van de Dojo-trainingstegel, een op zichzelf staand rekencluster dat een halve kubieke voet beslaat en in staat is tot 556 TFLOPS aan FP32-prestaties in een 15 kW vloeistofgekoeld pakket.

Elke tegel is uitgerust met 11 GB SRAM en is verbonden via een stof van 9 TB/s met behulp van een aangepast transportprotocol door de hele stapel.

"Deze trainingstegel vertegenwoordigt een ongeëvenaarde mate van integratie van computer naar geheugen tot stroomvoorziening en communicatie, zonder dat er extra schakelaars nodig zijn", zei Venkataramanan.

Het hart van de trainingstegel is Tesla's D1, een transistorchip van 50 miljard, gebaseerd op het 7nm-proces van TSMC. Tesla zegt dat elke D1 in staat is tot 22 TFLOPS aan FP32-prestaties bij een TDP van 400W. Tesla merkt echter op dat de chip in staat is een breed scala aan drijvende-kommaberekeningen uit te voeren, waaronder enkele aangepaste.

"Als je transistors vergelijkt voor vierkante millimeters, is dit waarschijnlijk de bloeding van alles wat er is," zei Venkataramanan.

Tesla nam vervolgens 25 D1's, stopte ze voor bekende goede dies en verpakte ze vervolgens met behulp van TSMC's systeem-op-wafer-technologie om "een enorme hoeveelheid computerintegratie te bereiken met een zeer lage latentie en een zeer hoge bandbreedte", zei hij.

Het systeem-op-waferontwerp en de verticaal gestapelde architectuur brachten echter uitdagingen met zich mee als het ging om vermogensafgifte.

Volgens Venkataramanan plaatsen de meeste versnellers tegenwoordig stroom direct naast het silicium. En hoewel bewezen, betekent deze aanpak dat een groot deel van het gaspedaal moet worden gewijd aan die componenten, wat het onpraktisch maakte voor Dojo, legde hij uit. In plaats daarvan ontwierp Tesla hun chips om stroom rechtstreeks via de onderkant van de dobbelsteen te leveren. 

Putting het allemaal samen

"We zouden een heel datacenter of een heel gebouw uit deze trainingstegel kunnen bouwen, maar de trainingstegel is slechts het rekengedeelte. We moeten het ook voeden', zei Venkataramanan.

Hiervoor ontwikkelde Tesla ook de Dojo Interface Processor (DIP), die fungeert als een brug tussen de host-CPU en trainingsprocessors. De DIP dient ook als een bron van gedeeld geheugen met hoge bandbreedte (HBM) en als een snelle 400Gbit/sec NIC.

Elke DIP beschikt over 32 GB HBM en maximaal vijf van deze kaarten kunnen worden aangesloten op een trainingstegel met 900 GB/s voor een totaal van 4.5 TB/s naar de host, voor een totaal van 160 GB HBM per tegel.

Tesla's V1-configuratieparen van deze tegels - of 150 D1-dies - in array ondersteunden vier host-CPU's die elk waren uitgerust met vijf DIP-kaarten om een ​​geclaimde exaflop van BF16- of CFP8-prestaties te bereiken.

Samengesteld, zegt Venkataramanan de architectuur - gedetailleerd in de diepte hier by Het volgende platform – stelt Tesla in staat om de beperkingen te overwinnen die gepaard gaan met traditionele versnellers van Nvidia en AMD.

“Hoe traditionele versnellers werken, probeer je meestal een heel model in elke versneller te passen. Repliceer het en laat de gegevens vervolgens door elk van hen stromen, "zei hij. “Wat gebeurt er als we steeds grotere modellen hebben? Deze versnellers kunnen plat vallen omdat ze geen geheugen meer hebben.”

Dit is geen nieuw probleem, merkte hij op. Nvidia's NV-switch maakt het bijvoorbeeld mogelijk om geheugen te poolen over grote banken met GPU's. Venkataramanan stelt echter dat dit niet alleen complexiteit toevoegt, maar ook latentie en compromissen op bandbreedte introduceert.

“Daar hebben we vanaf het begin over nagedacht. Onze computertegels en elk van de matrijzen zijn gemaakt voor het passen van grote modellen, "zei Venkataramanan.

Software

Zo'n gespecialiseerde rekenarchitectuur vraagt ​​om een ​​gespecialiseerde softwarestack. Venkataramanan en zijn team erkenden echter dat programmeerbaarheid Dojo zou maken of breken.

"Eenvoud van programmeerbaarheid voor software-tegenhangers is van het grootste belang wanneer we deze systemen ontwerpen", zei hij. "Onderzoekers wachten niet tot jullie softwaremensen een handgeschreven kernel schrijven voor aanpassing aan een nieuw algoritme dat we willen gebruiken."

Om dit te doen, liet Tesla het idee van het gebruik van kernels varen en ontwierp Dojo's architectuur rond compilers.

“Wat we deden, was dat we PiTorch gebruikten. We hebben een tussenlaag gemaakt, die ons helpt om de hardware eronder uit te schalen. Onder alles zit gecompileerde code', zei hij. "Dit is de enige manier om softwarestacks te maken die aanpasbaar zijn aan al die toekomstige workloads."

Ondanks de nadruk op softwareflexibiliteit, merkt Venkataramanan op dat het platform, dat momenteel in hun laboratoria draait, voorlopig beperkt is tot Tesla-gebruik.

"We richten ons in de eerste plaats op onze interne klanten", zei hij. “Elon heeft openbaar gemaakt dat we dit op termijn beschikbaar zullen stellen aan onderzoekers, maar daar hebben we geen tijdschema voor. ®