Tesla will Machine-Learning-Silizium ins Dojo bringen

Um den Durst nach immer größeren KI- und maschinellen Lernmodellen zu stillen, hat Tesla auf der Hot Chips 34 eine Fülle von Details zu seiner vollständig benutzerdefinierten Supercomputing-Architektur namens Dojo enthüllt.

Das System ist im Wesentlichen ein massiver zusammensetzbarer Supercomputer, obwohl es anders ist als das, was wir auf dem sehen Top 500, basiert es auf einer vollständig benutzerdefinierten Architektur, die die Rechen-, Netzwerk- und Eingabe-/Ausgabe-(E/A-)Chips bis hin zur Befehlssatzarchitektur (ISA), Stromversorgung, Verpackung und Kühlung umfasst. All dies wurde mit dem ausdrücklichen Ziel durchgeführt, maßgeschneiderte, spezifische maschinelle Lernalgorithmen in großem Umfang auszuführen.

„Die Datenverarbeitung in der realen Welt ist nur durch maschinelle Lerntechniken möglich, sei es die Verarbeitung natürlicher Sprache, das Fahren auf Straßen, die für das menschliche Sehen geschaffen sind, oder die Robotik, die mit der alltäglichen Umgebung verbunden ist“, sagte Ganesh Venkataramanan, Senior Director of Hardware Engineering bei Tesla während seiner Grundsatzrede.

Er argumentierte jedoch, dass herkömmliche Methoden zur Skalierung verteilter Workloads nicht mit der erforderlichen Geschwindigkeit beschleunigt werden konnten, um mit den Anforderungen des maschinellen Lernens Schritt zu halten. Tatsächlich schneidet das Mooresche Gesetz nicht ab, und die Systeme sind auch nicht für KI/ML-Training in großem Umfang verfügbar, nämlich eine Kombination aus CPU/GPU oder in selteneren Fällen durch die Verwendung spezieller KI-Beschleuniger. 

„Traditionell bauen wir Chips, wir setzen sie auf Gehäuse, Gehäuse kommen auf Leiterplatten, die in Systeme eingebaut werden. Systeme kommen in Racks“, sagte Venkataramanan. Das Problem besteht darin, dass jedes Mal, wenn sich Daten vom Chip zum Paket und aus dem Paket bewegen, eine Latenz- und Bandbreitenstrafe entsteht.

Ein Rechenzentrums-Sandwich

Um die Einschränkungen zu umgehen, fingen Venkataramanan und sein Team also von vorne an.

„Gleich bei meinem Interview mit Elon fragte er mich, was man anders als CPUs und GPUs für KI machen könne. Ich habe das Gefühl, dass das gesamte Team diese Frage immer noch beantwortet.“

Dies führte zur Entwicklung der Dojo-Schulungskachel, eines eigenständigen Rechenclusters, der einen halben Kubikfuß einnimmt und 556 TFLOPS FP32-Leistung in einem flüssigkeitsgekühlten 15-kW-Paket erbringen kann.

Jede Kachel ist mit 11 GB SRAM ausgestattet und über ein 9-TB/s-Fabric mit einem benutzerdefinierten Transportprotokoll im gesamten Stack verbunden.

„Diese Trainingskachel stellt ein beispielloses Maß an Integration vom Computer über den Speicher bis hin zur Stromversorgung und Kommunikation dar, ohne dass zusätzliche Schalter erforderlich sind“, sagte Venkataramanan.

Das Herzstück der Trainingskachel ist Teslas D1, ein 50-Milliarden-Transistorchip, der auf dem 7-nm-Prozess von TSMC basiert. Tesla sagt, dass jeder D1 22 TFLOPS FP32-Leistung bei einer TDP von 400 W leisten kann. Tesla merkt jedoch an, dass der Chip in der Lage ist, eine breite Palette von Gleitkommaberechnungen auszuführen, einschließlich einiger benutzerdefinierter Berechnungen.

„Wenn man Transistoren auf Quadratmillimeter vergleicht, ist dies wahrscheinlich das Neueste von allem, was es da draußen gibt“, sagte Venkataramanan.

Tesla nahm dann 25 D1s, ordnete sie für bekannte gute Chips ein und verpackte sie dann mit der System-on-Wafer-Technologie von TSMC, um „eine enorme Menge an Rechenintegration bei sehr geringer Latenz und sehr hoher Bandbreite zu erreichen“, sagte er.

Das System-on-Wafer-Design und die vertikal gestapelte Architektur brachten jedoch Herausforderungen mit sich, wenn es um die Stromversorgung ging.

Laut Venkataramanan platzieren die meisten Beschleuniger heute die Leistung direkt neben dem Silizium. Und obwohl sich dieser Ansatz bewährt hat, bedeutet er, dass ein großer Bereich des Beschleunigers diesen Komponenten gewidmet werden muss, was ihn für Dojo unpraktisch machte, erklärte er. Stattdessen hat Tesla seine Chips so konzipiert, dass sie Strom direkt durch die Unterseite des Chips liefern. 

Dass sie alle zusammen

„Wir könnten ein ganzes Rechenzentrum oder ein ganzes Gebäude aus dieser Trainingskachel bauen, aber die Trainingskachel ist nur der Rechenteil. Wir müssen es auch füttern“, sagte Venkataramanan.

Dafür hat Tesla auch den Dojo Interface Processor (DIP) entwickelt, der als Brücke zwischen der Host-CPU und den Trainingsprozessoren fungiert. Das DIP dient auch als Quelle für gemeinsam genutzten Speicher mit hoher Bandbreite (HBM) und als Hochgeschwindigkeits-NIC mit 400 Gbit/s.

Jeder DIP verfügt über 32 GB HBM, und bis zu fünf dieser Karten können mit einer Trainingskachel mit 900 GB/s für insgesamt 4.5 TB/s zum Host verbunden werden, was insgesamt 160 GB HBM pro Kachel ergibt.

Teslas V1-Konfigurationspaare dieser Kacheln – oder 150 D1-Dies – in einem Array unterstützten vier Host-CPUs, die jeweils mit fünf DIP-Karten ausgestattet waren, um eine behauptete Exaflop-Leistung von BF16 oder CFP8 zu erreichen.

Zusammengefasst sagt Venkataramanan die Architektur – detailliert hier ausführlich by Die nächste Plattform – ermöglicht es Tesla, die Einschränkungen zu überwinden, die mit traditionellen Beschleunigern von Nvidia und AMD verbunden sind.

„Wie herkömmliche Beschleuniger funktionieren, versucht man normalerweise, ein ganzes Modell in jeden Beschleuniger zu integrieren. Replizieren Sie es und lassen Sie die Daten dann durch jeden von ihnen fließen“, sagte er. „Was passiert, wenn wir immer größere Modelle haben? Diese Beschleuniger können versagen, weil ihnen der Arbeitsspeicher ausgeht.“

Dies ist kein neues Problem, bemerkte er. Der NV-Switch von Nvidia beispielsweise ermöglicht es, Speicher über große GPU-Bänke hinweg zu bündeln. Venkataramanan argumentiert jedoch, dass dies nicht nur die Komplexität erhöht, sondern auch Latenz und Kompromisse bei der Bandbreite einführt.

„Daran haben wir von Anfang an gedacht. Unsere Rechenkacheln und alle Chips wurden für die Anpassung an große Modelle entwickelt“, sagte Venkataramanan.

Software

Eine solche spezialisierte Rechenarchitektur erfordert einen spezialisierten Software-Stack. Venkataramanan und sein Team erkannten jedoch, dass die Programmierbarkeit Dojo entweder machen oder brechen würde.

„Einfache Programmierbarkeit für Software-Pendants ist von größter Bedeutung, wenn wir diese Systeme entwerfen“, sagte er. „Forscher warten nicht darauf, dass Ihre Software-Leute einen handgeschriebenen Kernel schreiben, um ihn an einen neuen Algorithmus anzupassen, den wir ausführen möchten.“

Um dies zu erreichen, verwarf Tesla die Idee, Kernel zu verwenden, und entwarf die Architektur von Dojo um Compiler herum.

„Wir haben PiTorch verwendet. Wir haben eine Zwischenschicht erstellt, die uns bei der Parallelisierung hilft, um die Hardware darunter zu skalieren. Unter allem befindet sich kompilierter Code“, sagte er. „Nur so lassen sich Software-Stacks erstellen, die an all diese zukünftigen Workloads anpassbar sind.“

Trotz der Betonung der Softwareflexibilität stellt Venkataramanan fest, dass die Plattform, die derzeit in ihren Labors läuft, vorerst auf die Nutzung durch Tesla beschränkt ist.

„Wir konzentrieren uns zuerst auf unsere internen Kunden“, sagte er. „Elon hat öffentlich gemacht, dass wir dies im Laufe der Zeit den Forschern zur Verfügung stellen werden, aber wir haben dafür keinen Zeitrahmen. ®