Tesla vuole portare il silicio per l'apprendimento automatico nel Dojo

Per soddisfare la sete di modelli di intelligenza artificiale e apprendimento automatico sempre più grandi, Tesla ha rivelato una vasta gamma di dettagli a Hot Chips 34 sulla loro architettura di supercalcolo completamente personalizzata chiamata Dojo.

Il sistema è essenzialmente un enorme supercomputer componibile, anche se a differenza di quello che vediamo sul Top 500, è costruito a partire da un'architettura interamente personalizzata che spazia dal silicio di elaborazione, rete e input/output (I/O) all'architettura del set di istruzioni (ISA), all'erogazione dell'alimentazione, al confezionamento e al raffreddamento. Tutto ciò è stato fatto con lo scopo esplicito di eseguire algoritmi di addestramento di machine learning specifici e personalizzati su larga scala.

"L'elaborazione dei dati nel mondo reale è fattibile solo attraverso tecniche di apprendimento automatico, sia che si tratti di elaborazione del linguaggio naturale, dalla guida in strade fatte per la visione umana alla robotica che si interfaccia con l'ambiente quotidiano", ha affermato Ganesh Venkataramanan, direttore senior dell'ingegneria hardware di Tesla durante il suo discorso di apertura.

Tuttavia, ha affermato che i metodi tradizionali per la scalabilità dei carichi di lavoro distribuiti non sono riusciti ad accelerare alla velocità necessaria per tenere il passo con le richieste dell'apprendimento automatico. In effetti, la legge di Moore non lo sta tagliando e nemmeno i sistemi sono disponibili per l'addestramento AI/ML su larga scala, vale a dire una combinazione di CPU/GPU o in circostanze più rare utilizzando acceleratori AI speciali. 

“Tradizionalmente costruiamo chip, li mettiamo sui pacchetti, i pacchetti vanno sui PCB, che vanno nei sistemi. I sistemi vanno nei rack", ha affermato Venkataramanan. Il problema è che ogni volta che i dati si spostano dal chip al pacchetto e fuori dal pacchetto, incorre in una penalità di latenza e larghezza di banda.

Un sandwich per datacenter

Quindi, per aggirare i limiti, Venkataramanan e il suo team hanno ricominciato da capo.

"Proprio dalla mia intervista con Elon, mi ha chiesto cosa puoi fare di diverso dalle CPU e dalle GPU per l'IA. Sento che l'intera squadra sta ancora rispondendo a questa domanda".

Ciò ha portato allo sviluppo del modulo di addestramento Dojo, un cluster di calcolo autonomo che occupa un mezzo piede cubo in grado di fornire 556 TFLOPS di prestazioni FP32 in un pacchetto raffreddato a liquido da 15 kW.

Ogni tile è dotato di 11 GB di SRAM ed è connesso su un fabric da 9 TB/s utilizzando un protocollo di trasporto personalizzato nell'intero stack.

"Questa tessera di formazione rappresenta una quantità senza precedenti di integrazione dal computer alla memoria, all'erogazione di energia, alla comunicazione, senza richiedere alcun interruttore aggiuntivo", ha affermato Venkataramanan.

Il cuore della tessera di addestramento è il D1 di Tesla, un die di transistor da 50 miliardi, basato sul processo a 7 nm di TSMC. Tesla afferma che ogni D1 è in grado di fornire 22 TFLOPS di prestazioni FP32 a un TDP di 400 W. Tuttavia, Tesla osserva che il chip è in grado di eseguire un'ampia gamma di calcoli in virgola mobile, inclusi alcuni personalizzati.

"Se confronti i transistor con un millimetro quadrato, questo è probabilmente il vantaggio di tutto ciò che è là fuori", ha detto Venkataramanan.

Tesla ha quindi preso 25 D1, li ha cestinati per dadi noti e quindi li ha impacchettati utilizzando la tecnologia system-on-wafer di TSMC per "raggiungere un'enorme quantità di integrazione di calcolo a latenza molto bassa e larghezza di banda molto elevata", ha affermato.

Tuttavia, il design del sistema su wafer e l'architettura impilata verticalmente hanno introdotto sfide per quanto riguarda l'erogazione di potenza.

Secondo Venkataramanan, la maggior parte degli acceleratori oggi colloca l'energia direttamente adiacente al silicio. E sebbene dimostrato, questo approccio significa che una vasta area dell'acceleratore deve essere dedicata a quei componenti, il che lo ha reso poco pratico per Dojo, ha spiegato. Invece, Tesla ha progettato i propri chip per fornire energia direttamente attraverso la parte inferiore del die. 

Mettere tutto insieme

“Potremmo costruire un intero data center o un intero edificio da questo riquadro di addestramento, ma il riquadro di addestramento è solo la parte di calcolo. Dobbiamo anche dargli da mangiare", ha detto Venkataramanan.

Per questo, Tesla ha anche sviluppato il Dojo Interface Processor (DIP), che funge da ponte tra la CPU host e i processori di training. Il DIP funge anche da fonte di memoria HBM (High Bandwidth Memory) condivisa e da NIC ad alta velocità da 400 Gbit/sec.

Ciascun DIP dispone di 32 GB di HBM e fino a cinque di queste schede possono essere collegate a un riquadro di formazione a 900 GB/s per un totale di 4.5 TB/s all'host per un totale di 160 GB di HBM per riquadro.

Le coppie di configurazione V1 di Tesla di queste tessere - o 150 dies D1 - in array supportavano quattro CPU host ciascuna dotata di cinque schede DIP per ottenere un exaflop dichiarato di prestazioni BF16 o CFP8.

Messo insieme, Venkataramanan dice l'architettura – dettagliata in profondità qui by La prossima piattaforma – consente a Tesla di superare i limiti associati agli acceleratori tradizionali di Nvidia e AMD.

“Come funzionano gli acceleratori tradizionali, in genere si tenta di inserire un intero modello in ciascun acceleratore. Replicalo e poi fai scorrere i dati attraverso ciascuno di essi", ha affermato. “Cosa succede se abbiamo modelli sempre più grandi? Questi acceleratori possono fallire perché esauriscono la memoria".

Questo non è un problema nuovo, ha osservato. Lo switch Nvidia di Nvidia, ad esempio, consente il raggruppamento della memoria su grandi banchi di GPU. Tuttavia, Venkataramanan sostiene che ciò non solo aggiunge complessità, ma introduce latenza e compromessi sulla larghezza di banda.

“Ci abbiamo pensato fin dall'inizio. Le nostre tessere di calcolo e ciascuno degli stampi sono stati realizzati per adattarsi a modelli di grandi dimensioni", ha affermato Venkataramanan.

Software

Un'architettura di calcolo così specializzata richiede uno stack software specializzato. Tuttavia, Venkataramanan e il suo team hanno riconosciuto che la programmabilità avrebbe creato o distrutto Dojo.

"La facilità di programmabilità per le controparti software è fondamentale quando progettiamo questi sistemi", ha affermato. "I ricercatori non aspetteranno che i tuoi esperti di software scrivano un kernel scritto a mano per adattarsi a un nuovo algoritmo che vogliamo eseguire".

Per fare ciò, Tesla ha abbandonato l'idea di utilizzare i kernel e ha progettato l'architettura di Dojo attorno ai compilatori.

“Quello che abbiamo fatto è stato usare PiTorch. Abbiamo creato uno strato intermedio, che ci aiuta a parallelizzare per ridimensionare l'hardware sottostante. Sotto ogni cosa c'è un codice compilato", ha detto. "Questo è l'unico modo per creare stack software adattabili a tutti quei carichi di lavoro futuri".

Nonostante l'enfasi sulla flessibilità del software, Venkataramanan osserva che la piattaforma, attualmente in esecuzione nei loro laboratori, è limitata all'uso di Tesla per il momento.

"Ci concentriamo innanzitutto sui nostri clienti interni", ha affermato. “Elon ha reso pubblico che nel tempo lo renderemo disponibile ai ricercatori, ma non abbiamo un lasso di tempo per questo. ®