Google erklärt TPU v4: KI-Supercomputer mit flexiblem optischen Netz

Dynamisch anpassbare Netzwerkstruktur, verbesserte Beschleuniger, effizienter als Nvidia: Googles TPU-v4-Supercomputer bringt Verbesserungen für Sprach- und Empfehlungsmodelle.

Ein Bericht von Johannes Hiltscher veröffentlicht am 5. April 2023, 15:48 Uhr

Ein Mainboard mit vier TPU v4, die Stecker für OCS befinden sich auf der Unterseite, oben vier PCIe-Stecker. (Bild: Google)

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 im Juni, bei Arxiv 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 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).

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.