Autonomes Fahren: Tesla erklärt Exaflops AI-Supercomputer Dojo
Tesla hat einen eigenen Supercomputer speziell für KI-Anwendungen entwickelt: Das Dojo genannte System soll anhand riesiger Mengen an Videodaten die Assistenzsysteme des Herstellers trainieren. Das Gesamtsystem erreicht beim 16-Bit-Gleitkommaformat BFP16 bis zu 1,1 Exaflops. Auf der Hot Chips sprachen Dojos Entwickler über Details der einzelnen Prozessoren und des Gesamtsystems. Erstmals vorgestellt wurde das System bereits 2021 .
Dojos Hardware wurde vollständig mit Blick auf die Anforderungen von AI-Anwendungen entwickelt. Dafür verzichtet der D1-Chip auf quasi alle Features anderer Prozessoren: Es gibt keinen virtuellen Speicher, keine Out-of-order-Execution, keine Caches, keine Kohärenzmechanismen. Die so eingesparte Chipfläche haben die Entwickler für große, schnelle Speicher und Verbindungen zwischen den einzelnen als Node bezeichneten Prozessoren genutzt. Das Ziel dabei: Die zwei Vektorrecheneinheiten permanent mit Daten zu versorgen, denn jeder Stillstand senkt die Leistung.
Möglich wird das durch vier parallele Hardware Threads, von denen Anwendungen üblicherweise ebenso viele zum Rechnen wie zur Kommunikation nutzen. Interessant ist, dass die Hardware pro Takt acht Befehle decodieren, aber maximal sechs verarbeiten kann: Den Threads stehen je zwei Address Generation Units (AGUs) und zwei Integer-ALUs zur Verfügung. Die Vektoreinheit besteht aus einem Block für SIMD-Berechnungen und einer Einheit für Matrixmultiplikationen. Viele Befehle arbeitet der Prozessor allerdings bereits beim Decodieren vollständig ab, so dass sie die Adress- und Recheneinheiten nicht erreichen.
Spezielle Recheneinheiten und viel schneller Speicher
Die Vektoreinheiten werden von einer sogenannten Gather-Einheit mit Daten versorgt, die einzelne Werte umsortieren und neu gruppieren kann. Damit können beispielsweise Matrizen für die Multiplikation transponiert werden. Daneben existiert ein Parser zur Abarbeitung von Listen. Er arbeitet ebenso wie die Gather-Engine direkt auf dem SRAM des Nodes.
Davon stehen jedem Node 1,25 MByte zur Verfügung, der Speicher kann mit 400 GByte/s gelesen werden, geschrieben werden Daten mit 270 GByte/s. Hier liegt der Programmcode, auch die benötigten Daten werden hineingeladen. Das passiert mittels Software Prefetching, idealerweise während die Vektoreinheiten rechnen. Jeder Node bekommt nur einen Teil der zu verarbeitenden Daten, die dann untereinander ausgetauscht werden. Dafür können Nodes über ein 2D-Mesh-Netzwerk gegenseitig auf ihren Speicher zugreifen.
Simples Netzwerk, hohe Bandbreite
An das chip-interne Netzwerk ist jeder Node über einen eigenen Funktionsblock angebunden. Der verbindet ihn mit seinen vier Nachbarn, in jede Richtung überträgt das Netzwerk pro Takt bidirektional 64 Byte. Diese 64-Byte-Blöcke bezeichnet Tesla als "Packet" , der Begriff wird später noch einmal auftauchen. Das Netzwerkmodul kann zudem pro Takt je ein Packet in den SRAM schreiben und aus ihm lesen. Auch beim Netzwerk haben sich die Architekten des Dojo für die simpelste Lösung entschieden: Das Ziel wird lediglich über seine X- und Y-Position adressiert. Den Weg durch die verschiedenen Nodes bestimmen einfache Routing-Tabellen.
Sie werden mittels Software verwaltet und dienen dazu, Daten an defekten Nodes vorbeizuleiten und die Kommunikationslast möglichst gleichmäßig zu verteilen. Das 2D-Gitter des Netzwerks führt Dojo auch außerhalb der einzelnen, 645 mm 2 großen und je 354 Nodes beinhaltenden Dies fort. Davon sind 25 in einem sogenannten Training Tile montiert. Zwischen zwei Dies werden Daten mit 2 TByte/s übertragen, im Die beträgt die Bisektionsbandbreite(öffnet im neuen Fenster) (Datentransferrate über die komplette Breite/Höhe) sogar 5 TByte/s.
Allerdings wurden auch ein paar Abkürzungen eingebaut. An jedes Training Tile können fünf sogenannte Dojo Interface Processors (DIP) angebunden werden. Auf denen sitzen nicht nur jeweils 32 GByte HBM-Speicher, sondern zusätzlich eine Ethernet-Schnittstelle. Sie kann als Ausweichroute genutzt werden, wenn die Distanz zwischen zwei Nodes zu groß wird.
Nicht nur Chip und Netzwerk sind selbst entwickelt, auch der Befehlssatz und einige Datentypen sind speziell an die Bedürfnisse von Anwendungen für maschinelles Lernen angepasst.
Eigener Befehlssatz, eigene Datenformate
Sowohl für die Vektor- als auch die SMT-Einheiten haben die Entwickler des Dojo einen eigenen, speziell an die Bedürfnisse der Anwendungen angepassten Befehlssatz entwickelt. Bei den SMT-Prozessoren habe man sich, so Emil Talpes, einer der leitenden Entwickler, an RISC-V orientiert. Allerdings habe man viele eigene Anweisungen ergänzt, sagte Talpes in einem der beiden Vorträge zu dem Thema auf der Hot Chips.
Dazu gehören neben Anweisungen zur Synchronisation zwischen einzelnen Nodes (Barrier) auch solche für den Datenaustausch. Soll ein Programm auf den Speicher eines anderen Nodes zugreifen, erfolgt das über spezielle Anweisungen. Auf den HBM-Speicher der DIPs wird mittels DMA zugegriffen, auch hierfür ist das Programm verantwortlich. Hier zeigt sich wieder die Einfachheit des Systems: Während andere Hochleistungsrechner - zumindest innerhalb eines Mainboards - versuchen, den Speicher als eine Einheit darzustellen, ist er bei Dojo extrem zerstückelt.
Die passenden Daten zur richtigen Zeit zur Verfügung zu haben, ist Sache der Softwareumgebung, also von Compiler und Laufzeitumgebung. Wie genau dies gelöst wurde, beantworteten die Entwickler Douglas Williams und Debjit Das Sarma im zweiten Vortrag zum Programmiermodell leider nicht. Besonders viele Befehle stehen Dojo für die Arbeit mit Matrizen zur Verfügung. Insgesamt 142 Anweisungen mit mehr als 1.000 Varianten decken diverse Spezialfälle ab.
Angepasste Datentypen
Weitere Eigenentwicklungen gibt es bei den Datentypen. Um mit möglichst wenigen Bits einen möglichst großen Zahlenbereich abdecken zu können, unterstützt Dojo neben den klassischen Gleitkommazahlen (FP32 und FP16) auch BFP16 und zwei als CFP8 bezeichnete und einen CFP16 benannten Datentypen. Sie unterscheiden sich in der Anzahl der für Exponent und Mantisse(öffnet im neuen Fenster) verwendeten Bits.
Während BFP16(öffnet im neuen Fenster) im Kontext des maschinellen Lernens geläufig ist, handelt es sich bei den CFP-Typen (das C steht für configurable) um Neuentwicklungen. Sie zeichnen sich durch einen verschiebbaren Wertebereich aus: Zwar bleiben Exponent und Mantisse stets gleich (4b/3b oder 5b/2b bei CFP8, 5b/10b bei CFP16), über einen sogenannten Bias kann aber deren Interpretation geändert werden. Dabei ist die Interpretation nicht fest, sondern kann per Software auf Ebene der Packets (64 Byte-Blöcke) angepasst werden.
Noch nicht voll in Betrieb
Etwas kurzsilbig antworteten die Entwickler auf die Frage, ob Dojo bereits in Betrieb sei. Hier hieß es lediglich, man habe "mehrere Systeme im Labor laufen" . Denkbar ist, dass die Software noch nicht wie gewünscht funktioniert. Denn je einfacher die Hardware ist, desto mehr Arbeit müssen Compiler, Laufzeitumgebung und - im schlimmsten Fall - Programmierer übernehmen. Speziell die Verteilung der Daten und deren Austausch zwischen den einzelnen Nodes dürfte eine Herausforderung sein.
Nichtsdestotrotz haben die Ingenieure ein extrem spannendes System entwickelt - das zeigt, dass durch volle Konzentration auf ein spezielles Problem extreme Leistungen möglich sind. Zwar bedient Dojo nur eine Nische, die ist für Tesla allerdings so wichtig, dass sie den enormen Aufwand eigener Hardwareentwicklung rechtfertigt.