TPU v4

In diesem Dokument werden die Architektur und die unterstützten Konfigurationen von Cloud TPU v4 beschrieben.

Systemarchitektur

Jeder TPU v4-Chip enthält zwei TensorCores. Jeder TensorCore hat vier Matrix-Multiply-Einheiten (MXUs), eine Vektoreinheit und eine Skalareinheit. Die folgende Tabelle enthält die wichtigsten Spezifikationen für einen v4 TPU-Pod.

Wichtige Spezifikationen	v4 Pod-Werte
Spitzenberechnungen pro Chip	275 TeraFLOPS (bf16 oder int8)
HBM2-Kapazität und Bandbreite	32 GiB, 1.200 Gbit/s
Gemessene minimale/mittlere/maximale Leistung	90/170/192 W
TPU-Pod-Größe	4096 Chips
Interconnect-Topologie	3D-Mesh
Spitzenberechnungen pro Pod	1.1 Exaflops (bf16 oder int8)
Alle Reduzierung der Bandbreite pro Pod	1,1 PB/s
Bandbreite pro Pod aufteilen	24 TB/s

Das folgende Diagramm zeigt einen TPU v4-Chip.

3D-Mesh- und 3D-Torus

v4-TPUs haben eine direkte Verbindung zu den nächsten benachbarten Chips in drei Dimensionen, was zu einem 3D-Mesh-Netzwerk von Netzwerkverbindungen führt. Die Verbindungen können als 3D-Torus für Slices konfiguriert werden, bei denen die Topologie, AxBxC, entweder 2A=B=C oder 2A=2B=C ist, wobei jede Dimension ein Vielfaches von 4 ist. Beispiel: 4 × 4 × 8, 4 × 8 × 8 oder 12 × 12 × 24. Im Allgemeinen ist mit einer 3D-Torus-Konfiguration eine bessere Leistung als mit einer 3D-Mesh-Konfiguration. Weitere Informationen finden Sie unter Twisted-Tori-Topologien.

Leistungsvorteile von TPU v4 gegenüber v3

In diesem Abschnitt werden die Leistungsvorteile von TPU v4 beschrieben

Speichersystem

Non Uniform Memory Access (NUMA) ist eine Computerspeicherarchitektur für Maschinen mit mehreren CPUs. Jede CPU hat direkten Zugriff auf einen Block von Hochgeschwindigkeitsspeicher. Eine CPU und ihr Arbeitsspeicher werden als NUMA-Knoten bezeichnet. NUMA-Knoten sind mit anderen NUMA-Knoten verbunden, die direkt nebeneinander liegen. Eine CPU von einem NUMA-Knoten kann auf Arbeitsspeicher in einem anderen NUMA-Knoten zugreifen. Dieser Zugriff ist jedoch langsamer als der Zugriff auf den Arbeitsspeicher innerhalb eines NUMA-Knotens.

Software, die auf einer Maschine mit mehreren CPUs ausgeführt wird, kann von einer CPU benötigte Daten in ihrem NUMA-Knoten platzieren, was den Arbeitsspeicherdurchsatz erhöht. Weitere Informationen zu NUMA finden Sie unter Non Uniform Memory Access auf Wikipedia.

Sie können die Vorteile von NUMA-locality nutzen, indem Sie Ihr Trainingsskript an NUMA-Knoten 0 binden.

So aktivieren Sie die NUMA-Knotenbindung:

1. Installieren Sie das numactl-Befehlszeilentool.

    $ sudo apt-get update
    $ sudo apt-get install numactl
   

2. Binden Sie Ihren Skriptcode an NUMA Node 0.

    $ numactl --cpunodebind=0 python3 your-training-script
   

Aktivieren Sie die NUMA-Knotenbindung, wenn:

• Wenn Ihre Arbeitslast unabhängig vom Framework stark von CPU-Arbeitslasten abhängig ist (z. B. von Bildklassifizierung, Empfehlungsarbeitslasten).
• Wenn Sie eine TPU-Laufzeitversion ohne -pod-Suffix verwenden (z. B. tpu-vm-tf-2.10.0-v4).

Weitere Unterschiede beim Arbeitsspeichersystem:

• v4-TPU-Chips haben einen einheitlichen Speicherbereich von 32 GiB HBM auf dem gesamten Chip, was eine bessere Koordination zwischen den beiden TensorCores auf dem Chip ermöglicht.
• Verbesserte HBM-Leistung mit neuesten Speicherstandards und -geschwindigkeiten.
• Verbessertes DMA-Leistungsprofil mit integrierter Unterstützung für hohe Leistungsfortschritte bei 512 Milliarden Detaillierungsgraden.

TensorCores

• Die doppelte Anzahl von MXUs und eine höhere Taktrate für maximal 275 TFLOPS.
• 2-fache Transposition- und Permutationsbandbreite.
• Load-Store-Arbeitsspeicherzugriffsmodell für Common Memory (Cmem).
• Schnellere MXU-Gewichtungsbandbreite und Unterstützung im 8-Bit-Modus, um niedrigere Batchgrößen und eine verbesserte Inferenzlatenz zu ermöglichen.

Inter-Chip-Interconnect-Verbindung

Sechs Verbindungsverbindungen pro Chip, um Netzwerktopologien mit kleineren Netzwerkdurchmessern zu ermöglichen.

eine andere Person

• x16 PCIE gen3-Schnittstelle zum Host (Direct Connect).
• Verbessertes Sicherheitsmodell.
• Verbesserte Energieeffizienz.

Konfigurationen

Ein TPU v4 Pod besteht aus 4.096 Chips, die über neu konfigurierbare Hochgeschwindigkeitsverbindungen miteinander verbunden sind. Mit dem flexiblen Netzwerk von TPU v4 können Sie die Chips in einem Pod-Slice derselben Größe auf mehrere Arten verbinden. Beim Erstellen eines TPU-Pod-Slice geben Sie die TPU-Version und die Anzahl der erforderlichen TPU-Ressourcen an. Wenn Sie ein TPU v4 Pod-Slice erstellen, können Sie seinen Typ und seine Größe auf zwei Arten angeben: AcceleratorType und AccleratorConfig.

Hinweis: AcceleratorType und AccleratorConfig schließen sich gegenseitig aus. Sie können beim Erstellen einer TPU nur eins verwenden.

AcceleratorType verwenden

Verwenden Sie AcceleratorType, wenn Sie keine Topologie angeben. Wenn Sie v4-TPUs mit AcceleratorType konfigurieren möchten, verwenden Sie beim Erstellen des TPU-Pod-Slice das Flag --accelerator-type. Legen Sie für --accelerator-type einen String fest, der die TPU-Version und die Anzahl der TensorCores enthält, die Sie verwenden möchten. Um beispielsweise ein v4 Pod-Slice mit 32 TensorCores zu erstellen, verwenden Sie --accelerator-type=v4-32.

Mit dem folgenden Befehl wird mithilfe des Flags --accelerator-type ein v4 TPU-Pod-Slice mit 512 TensorCores erstellt:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v4-512 \
    --version=tpu-vm-tf-2.16.1-pod-pjrt

Die Zahl nach der TPU-Version (v4) gibt die Anzahl der TensorCores an. Eine v4 TPU enthält zwei TensorCores, sodass die Anzahl der TPU-Chips 512/2 = 256 beträgt.

AcceleratorConfig verwenden

Verwenden Sie AcceleratorConfig, wenn Sie die physische Topologie Ihres TPU-Slice anpassen möchten. Dies ist im Allgemeinen für die Leistungsoptimierung mit Pod-Slices mit mehr als 256 Chips erforderlich.

Verwenden Sie die Flags --version und --topology, um V4-TPUs mit AcceleratorConfig zu konfigurieren. Legen Sie --version auf die zu verwendende TPU-Version und --topology auf die physische Anordnung der TPU-Chips im Pod-Slice fest.

Sie geben eine TPU-Topologie mit einem 3-Tupel (AxBxC) an, wobei A<=B<=C und A, B, C entweder alle <= 4 oder alle ganzzahligen Vielfachen von 4 sind. Die Werte A, B und C sind die Anzahl der Chips in jeder der drei Dimensionen. Zum Erstellen eines v4-Pod-Slice mit 16 Chips würden Sie beispielsweise --version=v4 und --topology=2x2x4 festlegen.

Mit dem folgenden Befehl wird ein v4 TPU-Pod-Slice mit 128 TPU-Chips in einem 4x4x8-Array erstellt:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.16.1-pod-pjrt

Topologien mit 2A=B=C oder 2A=2B=C haben auch Topologievarianten, die für die Gesamtkommunikation optimiert sind, z. B. 4 × 4 × 8, 8 × 8 × 16 und 12 × 12 × 24. Diese werden als Twisted Tori-Topologien bezeichnet.

Die folgenden Abbildungen zeigen einige gängige TPU v4-Topologien.

Größere Pod-Slices können aus einem oder mehreren 4x4x4-Chips erstellt werden.

Twisted Tori-Topologien

Bei einigen 3D-Torus-Slice-Formen in v4 kann eine sogenannte Twisted Torus-Topologie verwendet werden. Zum Beispiel können zwei v4-Würfel als 4x4x8-Slice oder 4x4x8_twisted angeordnet werden. Twisted-Topologien bieten eine erheblich höhere Bandbreite durch Zweiabschnitte. Eine erhöhte Bandbreite der geteilten Abschnitte ist nützlich für Arbeitslasten, die globale Kommunikationsmuster verwenden. Twisted Topologien können die Leistung bei den meisten Modellen verbessern, wobei große Arbeitslasten von TPU-Einbettungen am meisten davon profitieren.

Bei Arbeitslasten, die Datenparallelität als einzige Parallelitätsstrategie verwenden, funktionieren verdrillte Topologien möglicherweise etwas besser. Bei LLMs kann die Leistung mit einer Twisted-Topologie je nach Art der Parallelität (DP, MP usw.) variieren. Als Best Practice wird empfohlen, das LLM mit und ohne verdrillte Topologie zu trainieren, um festzustellen, welche die beste Leistung für Ihr Modell bietet. Bei einigen Tests mit dem FSDP MaxText-Modell wurden mit einer verdrillten Topologie ein bis zwei MFU-Verbesserungen festgestellt.

Der Hauptvorteil von verdrillten Topologien besteht darin, dass sie eine asymmetrische Torus-Topologie (z. B. 4 × 4 × 8) in eine eng verwandte symmetrische Topologie umwandelt. Die symmetrische Topologie bietet viele Vorteile:

• Verbessertes Load-Balancing
• Höhere Zweisektionsbandbreite
• Kürzere Paketrouten

Diese Vorteile führen letztendlich zu einer verbesserten Leistung bei vielen globalen Kommunikationsmustern.

Die TPU-Software unterstützt ein verdrilltes Tori auf Segmenten, bei denen die Größe jeder Dimension entweder gleich oder doppelt so groß wie die kleinste Dimension ist. Beispiel: 4 × 4 × 8, 4 × 8 × 8 oder 12 × 12 × 24.

Sehen Sie sich als Beispiel diese 4 × 2-Torus-Topologie mit TPUs an, die mit ihren Koordinaten (X,Y) im Slice gekennzeichnet sind:

Die Kanten in diesem Topologiediagramm werden der Übersichtlichkeit halber als ungerichtete Kanten dargestellt. In der Praxis ist jede Edge eine bidirektionale Verbindung zwischen TPUs. Die Kanten zwischen einer Seite dieses Rasters und der gegenüberliegenden Seite werden als umlaufende Kanten bezeichnet, wie im Diagramm dargestellt.

Durch das Verdrehen dieser Topologie erhalten wir eine vollständig symmetrische 4 × 2-Twisted-Torus-Topologie:

Was sich zwischen diesem und dem vorherigen Diagramm geändert hat, sind nur die Y-Wickelkanten. Anstatt eine Verbindung zu einer anderen TPU mit derselben X-Koordinate herzustellen, wurden sie verschoben, um eine Verbindung zur TPU mit der Koordinate X+2 mod 4 herzustellen.

Dieselbe Idee wird auf verschiedene Dimensionsgrößen und eine unterschiedliche Anzahl von Dimensionen verallgemeinert. Das resultierende Netzwerk ist symmetrisch, solange jede Dimension gleich oder doppelt so groß wie die kleinste Dimension ist.

Weitere Informationen zum Angeben einer Twisted Tori-Konfiguration beim Erstellen einer Cloud TPU finden Sie unter AcceleratorConfig verwenden.

Die folgende Tabelle zeigt die unterstützten verdrillten Topologien und die theoretische Steigerung der Bandbreite in getrennten Teilen im Vergleich zu unverdrillten Topologien.

Twisted Topologie	Theoretische Steigerung der Zweisektion Bandbreite im Vergleich zu einem verdrehten Torus
4×4×8_twisted	~70%
8x8x16_twisted
12×12×24_twisted
4×8×8_twisted	~40%
8×16×16_twisted

Hinweis: Die verdrillte Topologie wird nur dann automatisch verwendet, wenn das Suffix _twisted mit dem Flag --topology angegeben ist. Für eine 4x4x8-Twisted-Topologie würden Sie beispielsweise das Flag topology so angeben: 4x4x8_twisted.

Varianten der TPU v4-Topologie

Einige Topologien mit derselben Anzahl von Chips können auf unterschiedliche Weise angeordnet werden. Ein TPU-Pod-Slice mit 512 Chips (1.024 TensorCores) kann beispielsweise mit den folgenden Topologien konfiguriert werden: 4x4x32, 4x8x16 oder 8x8x8. Ein TPU-Pod-Slice mit 2.048 Chips (4.096 TensorCores) bietet noch mehr Topologieoptionen: 4x4x128, 4x8x64, 4x16x32 und 8x16x16.

Die Standardtopologie, die einer bestimmten Chipanzahl zugeordnet ist, ist einem Cube am ähnlichsten. Diese Form ist wahrscheinlich die beste Wahl für datenparalleles ML-Training. Andere Topologien können für Arbeitslasten mit mehreren Arten von Parallelität nützlich sein (z. B. Modell- und Datenparallelität oder räumliche Partitionierung einer Simulation). Diese Arbeitslasten funktionieren am besten, wenn die Topologie an die verwendete Parallelität angepasst ist. Wenn Sie beispielsweise die 4-Wege-Modellparallelität für die X-Dimension und die 256-Wege-Datenparallelität auf den Y- und Z-Dimensionen festlegen, entspricht eine 4 × 16 × 16-Topologie.

Modelle mit mehreren Parallelitätsdimensionen funktionieren am besten, wenn ihre Parallelitätsdimensionen den TPU-Topologiedimensionen zugeordnet werden. Dabei handelt es sich in der Regel um Daten und parallele Large Language Models (LLMs). Für ein TPU v4-Pod-Slice mit einer Topologie von 8 x 16 x 16 betragen die Abmessungen der TPU-Topologie beispielsweise 8, 16 und 16. Leistungsstarker ist die 8-Wege- oder 16-Wege-Modellparallelität, die einer der Dimensionen der physischen TPU-Topologie zugeordnet ist. Eine 4-Wege-Modellparallelität wäre bei dieser Topologie nicht optimal, da sie nicht an eine der Dimensionen der TPU-Topologie ausgerichtet ist. Sie wäre jedoch mit einer 4 × 16 × 32-Topologie auf derselben Anzahl von Chips optimal.

TPU v4-Konfigurationen bestehen aus zwei Gruppen, von denen Topologien kleiner als 64 Chips (kleine Topologien) und Topologien mit mehr als 64 Chips (große Topologien) sind.

Kleine v4-Topologien

Cloud TPU unterstützt die folgenden TPU v4 Pod-Slices, die kleiner als 64 Chips sind, also einen 4x4x4-Würfel. Sie können diese kleinen v4-Topologien entweder mit ihrem TensorCore-basierten Namen (z. B. v4-32) oder ihrer Topologie (z. B. 2x2x4) erstellen:

Name (basierend auf der Anzahl der TensorCores)	Anzahl der Chips	Topologie
v4-8	4	2x2x1
v4-16	8	2x2x2
v4-32	16	2x2x4
v4-64	32	2x4x4

Große v4-Topologien

TPU v4 Pod-Slices sind in Schritten von 64 Chips verfügbar. Die Formen sind in allen drei Dimensionen ein Vielfaches von 4. Die Dimensionen müssen in aufsteigender Reihenfolge angegeben werden. Die folgende Tabelle enthält mehrere Beispiele. Einige dieser Topologien sind "benutzerdefinierte" Topologien, die nur mit den Flags --type und --topology erstellt werden können, da es mehr als eine Möglichkeit gibt, die Chips anzuordnen.

Name (basierend auf der Anzahl der TensorCores)	Anzahl der Chips	Topologie
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
Nicht zutreffend, muss die Flags --type und --topology verwenden	256	4x4x16
v4-1024	512	8×8×8
v4-1536	768	8×8×12
v4-2048	1.024	8×8×16
Nicht zutreffend, muss die Flags --type und --topology verwenden	1.024	4×16×16
v4-4096	2.048	8×16×16
…	…	…