Anforderungen neuronaler Netze ändern sich

Als erste Änderung wird eine direkte Unterstützung für sogenannte Embeddings 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² 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.