Google erklärt TPU v4: KI-Supercomputer mit flexiblem optischen Netz
Google verrät Details zur vierten Version seiner Tensor Processing Unit (TPU v4), die seit 2021 im Einsatz ist. Offiziell vorgestellt werden sie zwar erst auf dem International Symposium on Computer Architecture(öffnet im neuen Fenster) im Juni, bei Arxiv(öffnet im neuen Fenster) findet sich allerdings bereits das eingereichte Paper. Neben architektonischen Verbesserungen ist das rekonfigurierbare optische Netzwerk besonders interessant.
Bei Google bezeichnet TPU nicht nur den eigenen Prozessor, sondern auch einen damit aufgebauten Supercomputer für KI-Anwendungen. Bei dessen vierter Version sollte die Anzahl an TPUs vervierfacht werden - 4.096 anstelle von 1.024 Chips. Das stellte die Entwickler vor eine Herausforderung: Das Netzwerk sollte, um eine hohe Bisektionsbandbreite(öffnet im neuen Fenster) zu erreichen, ein Torus sein. Allerdings passen in ein Rack nur 64 TPUs, angeordnet als 4x4x4-Würfel. Dadurch wurden aber einige Strecken zu lang für elektrische Verbindungen.
Aufgrund der deutlich höheren Kosten ist Redundanz zum Umschiffen ausgefallener Knoten hier aufwendiger zu schaffen. Daher entschied sich das Google-Team, einen optischen Switch zu entwickeln. Der enthält einen MEMS-Chip mit mikromechanischen Spiegeln, der schlicht als Optical Circuit Switch (OCS) bezeichnet wird. So kann jeder Eingang auf einen beliebigen Ausgang geschaltet werden.
Netzwerkstruktur lässt sich ändern
Neben einer Erhöhung der Fehlertoleranz hat der OCS noch weitere Vorteile: Durch die änderbare Netzwerkstruktur lassen sich neben dem normalen 3D-Torus auch verdrehte Varianten (Twisted Torus) erzeugen. Die können, je nach Anwendung, den Netzwerkdurchsatz verbessern (PDF)(öffnet im neuen Fenster) .
Auch lassen sich beliebige Racks miteinander verbinden, was die Auslastung des Clusters verbessert. Beim Vorgänger musste eine Anfrage noch auf nebeneinanderliegende Racks verteilt werden, diese Einschränkung fällt weg. Ein weiterer Vorteil: Das Netzwerk lässt sich an die Struktur des berechneten KI-Modells anpassen. Das soll die Kommunikation effizienter gestalten können.
Die Sicherheit soll OCS ebenfalls erhöhen, da Racks, die an verschiedenen Aufgaben arbeiten, auf Netzwerkebene voneinander getrennt werden können. Als letzten Punkt führt das Paper an, dass der Supercomputer dank OCS wesentlich schneller in Betrieb genommen werden konnte, da nicht erst alle Racks aufgebaut und verkabelt werden mussten.
Trotz der vielen Möglichkeiten soll OCS verhältnismäßig günstig sein: Googles Ingenieure geben an, dass die Kosten bei unter fünf Prozent der Gesamtkosten des Supercomputers liegen. Bei der Leistungsaufnahme schlägt das optische Netz mit unter drei Prozent ebenfalls kaum zu Buche.
Gegenüber der TPU v3 wurde aber nicht nur das Netzwerk überarbeitet. Auch die Berechnungseinheiten sind und können mehr, um Veränderungen im Aufbau neuronaler Netze Rechnung zu tragen.
Anforderungen neuronaler Netze ändern sich
Als erste Änderung wird eine direkte Unterstützung für sogenannte Embeddings(öffnet im neuen Fenster) genannt. Dabei handelt es sich um eine kompakte Darstellung großer, aber dünn besetzter (sparse) Vektoren. Genutzt werden sie etwa bei Sprach- und Empfehlungsmodellen, implementiert sind sie in den Sparsecores der TPU.
Die existieren zwar bereits seit Version zwei der TPU, im Paper geht Google erstmals näher auf sie ein. Sie sind speziell auf sogenannte Scatter-Gather-Operationen ausgelegt, bei denen einzelne Daten aus dem Speicher eingesammelt oder an viele Stellen verteilt werden müssen. Dafür verfügen die Sparsecores über eine spezielle DMA-Einheit, die auf den HBM-Speicher des gesamten Supercomputers in einem globalen Adressraum zugreifen kann.
Embeddings sind dabei eine besondere Herausforderung: Üblicherweise können Scatter-Gather-Einheiten nur strukturiert auf den Speicher zugreifen. Daten müssen dann feste Abstände haben. Embeddings hingegen erfordern, dass die Daten an beliebigen Adressen liegen können - hier dürfte die Weiterentwicklung gegenüber der TPU v3 liegen.
Da die Bedeutung von Netzen mit dünn besetzten Gewichtsmatrizen zunehmen, wurde die Anzahl der Sparsecores verdoppelt, die TPU v4 verfügt jetzt über vier. Gegenüber dem Vorgänger gibt es eine weitere Neuerung: Jeder TPU-Chip verfügt nun über 128 MByte schnellen Scratchpad-Speicher, der in einigen Anwendungen deutliche Geschwindigkeitsvorteile bringt. Durch einen Umstieg auf 7-nm-Fertigung - die TPU v3 entstand in einem 16-nm-Prozess - ist die aktuelle Generation trotz mehr als doppelt so hoher Transistorzahl (22 Milliarden) und Rechenleistung (275 TFLOPS bei BF16 und Int8) mit 600 mm 2 kleiner und maximal 192 Watt Leistungsaufnahme sparsamer.
KI hilft bei Parametrisierung
Die vielen Konfigurationsmöglichkeiten haben allerdings auch einen Nachteil: Ein neuronales Netz optimal zu konfigurieren und seine Berechnung auf die Hardware zu verteilen, wird zu einem komplexen Optimierungsproblem. Um das zu lösen, setzt Google wieder KI ein - die Vorschläge für die Aufteilung der Berechnungen einer anderen KI macht.
Durch eine Änderung der Anzahl verwendeter TPUs, Netzwerkstruktur und Embeddings soll die KI selbst Experten um bis zu 20 Prozent schlagen - und deutlich schneller sein. Durch zusätzliche Embeddings erreicht sie zudem eine bessere Verteilung der Last auf Tensor- und Sparsecores. So lässt sich die theoretische Leistung des Chips besser ausreizen.
Schneller als Nvidias A100 bei geringerer Leistungsaufnahme
Abschließend vergleicht das Paper noch die Leistung des TPU-v4-Clusters gegen Nvidias A100 Rechen-GPU. Während Googles System im Benchmark MLPerf nur fünf Prozent vorne liegt, sind es bei anderen Netzen bis zu 87 Prozent Vorsprung (Retinanet). Gleichzeitig soll die Leistungsaufnahme - je nach Benchmark - bis zu 50 Prozent niedriger sein. Einen großen Anteil daran hat das Netzwerk, weshalb laut Paper bereits überlegt werde, die kommende TPU-Generation vollständig optisch zu vernetzen.
Mit Nvidias H100 werden zwar die Karten neu gemischt, aber auch Google hat bereits einen Nachfolger: Die TPU v5 soll noch 2023 fertig werden, dank 4-nm-Fertigung wird auch hier die Rechenleistung ordentlich steigen. Der Wettlauf um die leistungsfähigste KI-Hardware bleibt also spannend.