Tesla vill ta kisel för maskininlärning till Dojo

För att släcka törsten efter allt större AI- och maskininlärningsmodeller har Tesla avslöjat en mängd detaljer på Hot Chips 34 om deras helt anpassade superdatorarkitektur som heter Dojo.

Systemet är i huvudsak en massiv komponerbar superdator, även om det till skillnad från vad vi ser på Top 500, den är byggd från en helt anpassad arkitektur som spänner över beräknings-, nätverks- och input/output (I/O) kisel till instruktionsuppsättningsarkitektur (ISA), strömförsörjning, paketering och kylning. Allt gjordes med det uttryckliga syftet att köra skräddarsydda, specifika träningsalgoritmer för maskininlärning i stor skala.

"Databearbetning i den verkliga världen är endast möjlig genom maskininlärningstekniker, vare sig det är naturligt språkbehandling, körning på gator som är gjorda för mänsklig vision till robotik som gränsar till den vardagliga miljön," sa Ganesh Venkataramanan, senior chef för maskinvaruteknik på Tesla, under sitt huvudtal.

Han hävdade dock att traditionella metoder för att skala distribuerade arbetsbelastningar har misslyckats med att accelerera i den takt som är nödvändig för att hålla jämna steg med kraven på maskininlärning. I själva verket minskar inte Moores lag det och inte heller systemen är tillgängliga för AI/ML-träning i stor skala, nämligen någon kombination av CPU/GPU eller i sällsynta fall genom att använda speciella AI-acceleratorer. 

"Traditionellt bygger vi chips, vi sätter dem på paket, paket går på PCB, som går in i system. System går in i rack, säger Venkataramanan. Problemet är att varje gång data flyttas från chipet till paketet och från paketet, medför det en latens- och bandbreddsstraff.

En datacentersmörgås

Så för att komma runt begränsningarna började Venkataramanan och hans team om från början.

"Från min intervju med Elon frågade han mig vad du kan göra som skiljer sig från CPU:er och GPU:er för AI. Jag känner att hela laget fortfarande svarar på den frågan.”

Detta ledde till utvecklingen av Dojo-träningsbrickan, ett fristående datorkluster som upptar en halvkubiksfot som klarar 556 TFLOPS FP32-prestanda i ett 15kW vätskekylt paket.

Varje bricka är utrustad med 11 GB SRAM och är ansluten över ett 9TB/s-tyg med hjälp av ett anpassat transportprotokoll genom hela stacken.

"Denna träningsbricka representerar oöverträffade mängder av integration från dator till minne till strömförsörjning, till kommunikation, utan att behöva några extra switchar," sa Venkataramanan.

I hjärtat av träningsbrickan är Teslas D1, en 50 miljarder transistormatris, baserad på TSMC:s 7nm-process. Tesla säger att varje D1 är kapabel till 22 TFLOPS FP32-prestanda vid en TDP på ​​400W. Tesla noterar dock att chippet kan köra ett brett utbud av flyttalsberäkningar inklusive några få anpassade.

"Om du jämför transistorer för millimeterkvadrat är detta förmodligen den blödande kanten på allt som finns där ute," sa Venkataramanan.

Tesla tog sedan 25 D1:or, lagrade dem för kända bra dies och paketerade dem sedan med TSMC:s system-on-wafer-teknik för att "uppnå en enorm mängd beräkningsintegration med mycket låg latens och mycket hög bandbredd", sa han.

System-on-wafer-designen och den vertikalt staplade arkitekturen introducerade dock utmaningar när det kom till kraftleverans.

Enligt Venkataramanan placerar de flesta acceleratorer idag kraft i direkt anslutning till kislet. Och även om det är bevisat innebär detta tillvägagångssätt att en stor del av acceleratorn måste dedikeras till dessa komponenter, vilket gjorde det opraktiskt för Dojo, förklarade han. Istället designade Tesla sina marker för att leverera kraft direkt genom tärningens botten. 

Få alltid att falla på plats

"Vi skulle kunna bygga ett helt datacenter eller en hel byggnad av den här utbildningsbrickan, men utbildningsbrickan är bara beräkningsdelen. Vi måste också mata den, sa Venkataramanan.

För detta utvecklade Tesla även Dojo Interface Processor (DIP), som fungerar som en brygga mellan värd-CPU och träningsprocessorer. DIP fungerar också som en källa för delat högbandbreddsminne (HBM) och som ett höghastighetsnätverk med 400 Gbit/sek.

Varje DIP har 32 GB HBM och upp till fem av dessa kort kan anslutas till en träningsplatta med 900 GB/s för sammanlagt 4.5 TB/s till värden för totalt 160 GB HBM per bricka.

Teslas V1-konfigurationspar av dessa brickor – eller 150 D1-matriser – i array stödde fyra värd-CPU:er var och en utrustad med fem DIP-kort för att uppnå en påstådd exaflop av BF16- eller CFP8-prestanda.

Sammantaget säger Venkataramanan arkitekturen – detaljerad i djupet här by Nästa plattform – gör det möjligt för Tesla att övervinna begränsningarna förknippade med traditionella acceleratorer från sådana som Nvidia och AMD.

"Hur traditionella acceleratorer fungerar, försöker man vanligtvis att passa in en hel modell i varje accelerator. Replikera det och flöda sedan data genom var och en av dem, sa han. ”Vad händer om vi har större och större modeller? Dessa acceleratorer kan falla platt eftersom de tar slut på minne."

Det här är inget nytt problem, konstaterade han. Nvidias NV-switch gör till exempel att minne kan poolas över stora banker av GPU:er. Men Venkataramanan hävdar att detta inte bara lägger till komplexitet, utan introducerar latens och kompromisser med bandbredd.

– Vi har tänkt på det här redan från början. Våra beräkningsplattor och var och en av formarna gjordes för att passa stora modeller, säger Venkataramanan.

Mjukvara

En sådan specialiserad beräkningsarkitektur kräver en specialiserad mjukvarustapel. Men Venkataramanan och hans team insåg att programmerbarhet antingen skulle göra eller bryta Dojo.

"Enkel programmerbarhet för programvarumotsvarigheter är avgörande när vi designar dessa system," sa han. "Forskare väntar inte på att dina programvarufolk ska skriva en handskriven kärna för anpassning till en ny algoritm som vi vill köra."

För att göra detta övergav Tesla idén om att använda kärnor och designade Dojos arkitektur kring kompilatorer.

"Vad vi gjorde var att vi använde PiTorch. Vi skapade ett mellanlager, som hjälper oss att parallellisera för att skala ut hårdvaran under det. Under allt finns kompilerad kod”, sa han. "Detta är det enda sättet att skapa mjukvarustackar som är anpassningsbara till alla dessa framtida arbetsbelastningar."

Trots betoningen på mjukvaruflexibilitet, noterar Venkataramanan att plattformen, som för närvarande körs i deras labb, är begränsad till Tesla-användning för tillfället.

"Vi fokuserar på våra interna kunder först," sa han. "Elon har gjort det offentligt att vi med tiden kommer att göra detta tillgängligt för forskare, men vi har ingen tidsram för det. ®