TPU-Konfigurationen

TPU v4-Konfigurationen

Ein TPU v4 Pod besteht aus 4.096 Chips, die über rekonfigurierbare Hochgeschwindigkeitsverbindungen miteinander verbunden sind. Mit dem flexiblen Netzwerk von TPU v4 können Sie die Chips in Podslices derselben Größe auf mehrere Arten verbinden. Wenn Sie einen TPU-Podslice erstellen, geben Sie die TPU-Version und die Anzahl der erforderlichen TPU-Ressourcen an. Wenn Sie eine TPU v4-Podslice erstellen, können Sie den Typ und die Größe auf zwei Arten angeben:

  1. AcceleratorType verwenden

    Diese Methode wird empfohlen, wenn Sie keine Topologie angeben.

  2. AcceleratorConfig verwenden

    Wir empfehlen die Verwendung von AcceleratorConfig, wenn Sie die physische Topologie anpassen möchten. Dies ist im Allgemeinen für die Leistungsoptimierung mit Podslices >= 256 Chips erforderlich.

Verwenden Sie zum Konfigurieren von v4-TPUs mit AcceleratorType das Flag --accelerator-type beim Erstellen der TPU-Podslices. Legen Sie für --accelerator-type einen String fest, der die TPU-Version und die Anzahl der zu verwendenden TensorCores enthält. Wenn Sie beispielsweise einen v4-Podslice mit 32 TensorCores erstellen möchten, verwenden Sie --accelerator-type=v4-32.

Verwenden Sie die Flags --version und --topology, um v4-TPUs mit AcceleratorConfig zu konfigurieren. Legen Sie --version auf die TPU-Version fest, die Sie verwenden möchten, und --topology auf die physische Anordnung der TPU-Chips in Podslice. Wenn Sie beispielsweise eine v4-Podslice mit 16 Chips erstellen möchten, verwenden Sie --version=v4 und --topology=2x2x4.

Sie geben eine TPU-Topologie mit einem 3-Tupel AxBxC an, wobei A<=B<=C und A, B, C entweder alle <= 4 oder ein ganzzahliges Vielfaches von 4 sind. Die Werte A, B und C sind die Anzahl der Chips in jeder der drei Dimensionen. Topologien mit 2A=B=C oder 2A=2B=C haben auch Topologievarianten, die für die gesamte Kommunikation optimiert sind, z. B. 4 × 4 × 8, 8 × 8 × 16 und 12 × 12 × 24. Diese werden als "Twisted Tori"-Topologien bezeichnet. Eine Twisted Tori-Konfiguration kann für jede Topologie festgelegt werden, die sie unterstützt. Eine Liste der unterstützten Twisted-Topologien finden Sie unter Variantenkonfigurationen.

Mit dem folgenden Befehl wird ein v4 TPU-Podslice mit 128 TPU-Chips erstellt, die in einem 4x4x8-Array angeordnet sind:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=4x4x8
    --version=tpu-vm-tf-2.13.0

Sie können mit --accelerator-type eine v4 TPU-Podslices mit der gleichen Anzahl von Chips/TensorCores erstellen:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v4-512
    --version=tpu-vm-tf-2.13.0

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

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

Image

Größere Podslices können aus einem oder mehreren 4x4x4-"Cubes" von Chips erstellt werden.

Topologievarianten

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

Die Standardtopologie, die einer bestimmten Chipanzahl zugeordnet ist, ist einem Cube am ähnlichsten (siehe Topologieformen). 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 der verwendeten Parallelität entspricht. Wenn Sie beispielsweise die 4-Wege-Modellparallelität auf die X-Dimension und die 256-Wege-Datenparallelität auf den Y- und Z-Dimensionen anwenden, entspricht eine 4x16x16-Topologie.

Modelle mit mehreren Parallelitätsdimensionen funktionieren am besten, wenn ihre Parallelitätsdimensionen den TPU-Topologiedimensionen zugeordnet sind. Dies sind in der Regel parallele Large Language Models (LLMs) aus Daten und Modellen. Für eine TPU v4-Podlice mit einer Topologie von 8 x 16 x 16 betragen die Abmessungen der TPU-Topologie beispielsweise 8, 16 und 16. Es ist leistungsfähiger, die 8-Wege- oder 16-Wege-Modellparallelität zu verwenden, die einer der physischen TPU-Topologiedimensionen 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 4x16x32-Topologie auf derselben Anzahl von Chips optimal.

TPU v4-Konfigurationen bestehen aus zwei Gruppen: Topologien mit weniger als 64 Chips (kleine Topologien) und Topologien mit mehr als 64 Chips (große Topologien).

Kleine v4-Topologien

Cloud TPU unterstützt die folgenden TPU v4-Podslices, die kleiner als 64 Chips sind, einen 4x4x4-Cube. 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-Podslices sind in Schritten von 64 Chips verfügbar, wobei die Formen für alle drei Dimensionen ein Vielfaches von 4 sind. Die Dimensionen müssen ebenfalls in aufsteigender Reihenfolge sein. In der folgenden Tabelle sind einige Beispiele aufgeführt. Einige dieser Topologien sind "benutzerdefinierte" Topologien, die nur mit den Flags --type und --topology erstellt werden können, da die Chips auf mehrere Arten angeordnet werden können.

Mit dem folgenden Befehl wird ein v4 TPU-Podslice mit 512 TPU-Chips erstellt, die in einem 8x8x8-Array angeordnet sind:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --type=v4
    --topology=8x8x8
    --version=tpu-vm-tf-2.13.0

Sie können mit --accelerator-type eine v4 TPU-Podslices mit der gleichen Anzahl von Chips/TensorCores erstellen:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v4-1024
    --version=tpu-vm-tf-2.13.0
Name (basierend auf der Anzahl der TensorCores) Anzahl der Chips Topologie
v4–128 64 4x4x4
v4–256 128 4x4x8
v4–512 256 4x8x8
– – muss die Flags --type und --topology verwenden 256 4x4x16
v4–1024 512 8x8x8
v4–1536 768 8x8x12
v4–2048 1.024 8x8x16
– – muss die Flags --type und --topology verwenden 1.024 4x16x16
v4–4096 2.048 8x16x16
… … …

TPU v3-Konfigurationen

Ein TPU v3 Pod besteht aus 1.024 Chips, die über Hochgeschwindigkeitsverbindungen miteinander verbunden sind. Verwenden Sie zum Erstellen eines TPU v3-Geräts oder Podslice das Flag --accelerator-type für den Befehl "gcloud compute tpus tpu-vm". Zum Angeben des Beschleunigertyps geben Sie die TPU-Version und die Anzahl der TPU-Kerne an. Verwenden Sie für eine einzelne v3-TPU --accelerator-type=v3-8. Verwenden Sie für eine v3-Podslice mit 128 TensorCores --accelerator-type=v3-128.

Der folgende Befehl zeigt, wie Sie eine v3 TPU-Podslice mit 128 TensorCores erstellen:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v3-128
    --version=tpu-vm-tf-2.13.0

In der folgenden Tabelle sind die unterstützten v3-TPU-Typen aufgeführt:

TPU-Version Support-Ende
v3-8 (Enddatum wurde noch nicht festgelegt)
v3-32 (Enddatum wurde noch nicht festgelegt)
v3–128 (Enddatum wurde noch nicht festgelegt)
v3–256 (Enddatum wurde noch nicht festgelegt)
v3–512 (Enddatum wurde noch nicht festgelegt)
v3–1024 (Enddatum wurde noch nicht festgelegt)
v3–2048 (Enddatum wurde noch nicht festgelegt)

Weitere Informationen zum Verwalten von TPUs finden Sie unter TPUs verwalten. Weitere Informationen zu den verschiedenen Versionen von Cloud TPU finden Sie unter Systemarchitektur.

TPU v2-Konfigurationen

Ein TPU v2 Pod besteht aus 512 Chips, die über rekonfigurierbare Hochgeschwindigkeitsverbindungen miteinander verbunden sind. Verwenden Sie zum Erstellen einer TPU v2-Podslice das Flag --accelerator-type für den Befehl gcloud compute tpus tpu-vm. Zum Angeben des Beschleunigertyps geben Sie die TPU-Version und die Anzahl der TPU-Kerne an. Verwenden Sie für eine einzelne v2-TPU --accelerator-type=v2-8. Verwenden Sie für eine v2-Podslice mit 128 TensorCores --accelerator-type=v2-128.

Der folgende Befehl zeigt, wie Sie eine v2 TPU-Podslice mit 128 TensorCores erstellen:

  $ gcloud compute tpus tpu-vm create tpu-name
    --zone=zone
    --accelerator-type=v2-128
    --version=tpu-vm-tf-2.13.0

Weitere Informationen zum Verwalten von TPUs finden Sie unter TPUs verwalten. Weitere Informationen zu den verschiedenen Versionen von Cloud TPU finden Sie unter Systemarchitektur.

In der folgenden Tabelle sind die unterstützten v2-TPU-Typen aufgeführt

TPU-Version Support-Ende
v2-8 (Enddatum wurde noch nicht festgelegt)
v2-32 (Enddatum wurde noch nicht festgelegt)
v2-128 (Enddatum wurde noch nicht festgelegt)
v2-256 (Enddatum wurde noch nicht festgelegt)
v2-512 (Enddatum wurde noch nicht festgelegt)

TPU-Kompatibilität

Sie können den TPU-Typ in einen anderen TPU-Typ mit der gleichen Anzahl von TensorCores oder Chips (z. B. v3-128 und v4-128) ändern und Ihr Trainingsskript ohne Codeänderungen ausführen. Wenn Sie jedoch zu einem TPU-Typ mit einer größeren oder kleineren Anzahl von Chips oder TensorCores wechseln, müssen Sie eine erhebliche Feinabstimmung und Optimierung vornehmen. Weitere Informationen finden Sie unter Auf TPU-Pods trainieren.

TPU-Softwareversionen

Welche Version der TPU-Software Sie verwenden sollten, hängt von der TPU-Architektur (TPU-VM oder TPU-Knoten) und dem verwendeten ML-Framework (TensorFlow, PyTorch oder JAX) ab.

TPU-VM

Für TensorFlow-, PyTorch- und JAX-Frameworks sind TPU-Softwareversionen verfügbar.

TensorFlow

Verwenden Sie die TPU-Softwareversion, die der Version von TensorFlow entspricht, mit der Ihr Modell geschrieben wurde. Wenn Sie beispielsweise TensorFlow 2.13.0 verwenden, nutzen Sie die TPU-Softwareversion tpu-vm-tf-2.13.0. Die aktuell unterstützten TensorFlow TPU-VM-Softwareversionen für TPUs sind:

  • tpu-vm-tf-2.13.0
  • tpu-vm-tf-2.12.1
  • tpu-vm-tf-2.12.0
  • tpu-vm-tf-2.11.1
  • tpu-vm-tf-2.11.0
  • tpu-vm-tf-2.10.1
  • tpu-vm-tf-2.10.0
  • tpu-vm-tf-2.9.3
  • tpu-vm-tf-2.9.1
  • tpu-vm-tf-2.8.4
  • tpu-vm-tf-2.8.3.
  • tpu-vm-tf-2.8.0
  • tpu-vm-tf-2.7.4
  • tpu-vm-tf-2.7.3

Wenn Sie ein Podslice verwenden, hängen Sie -pod an die TPU-Softwareversion an, die Sie verwenden möchten. Beispiel: tpu-vm-tf-2.13.0-pod.

TPU-VM mit TPU v4

Wenn Sie ein Modell auf TPU v4 mit TensorFlow trainieren, verwenden Sie eine der v4-Versionen aus der folgenden Tabelle.

TensorFlow-Version Version der TPU-Software
2.13.0 tpu-vm-tf-2.13.0, tpu-vm-tf-2.13.0-pod
2.12.1 tpu-vm-tf-2.12.1, tpu-vm-tf-2.12.1-pod
2.12.0 tpu-vm-tf-2.12.0, tpu-vm-tf-2.12.0-pod
2.11.1 tpu-vm-tf-2.11.1, tpu-vm-tf-2.11.1-pod
2.11.0 tpu-vm-tf-2.11.0, tpu-vm-tf-2.11.0-pod
2.10.0 tpu-vm-tf-2.10.0-v4, tpu-vm-tf-2.10.0-pod-v4
2.10.1 tpu-vm-tf-2.10.1, tpu-vm-tf-2.10.1-pod
2.9.3 tpu-vm-tf-2.9.3-v4, tpu-vm-tf-2.9.3-pod-v4
2.9.2 tpu-vm-tf-2.9.2-v4, tpu-vm-tf-2.9.2-pod-v4
2.9.1 tpu-vm-tf-2.9.1-v4, tpu-vm-tf-2.9.1-pod-v4

Weitere Informationen zu TensorFlow-Patchversionen finden Sie unter Unterstützte TensorFlow-Patchversionen.

TPU-VMs werden mit TensorFlow und der entsprechenden Libtpu-Bibliothek erstellt. Wenn Sie ein eigenes VM-Image erstellen, geben Sie die folgenden TensorFlow-TPU-Softwareversionen und entsprechende Libtpu-Versionen an:

TensorFlow-Version libtpu.so-Version
2.13.0 1.7.0
2.12.1 1.6.1
2.12.0 1.6.0
2.11.1 1.5.1
2.11.0 1.5.0
2.10.1 1.4.1
2.10.0 1.4.0
2.9.3 1.3.2
2.9.1 1.3.0
2.8.3 1.2.3
2,8.* 1.2.0
2.7.3 1.1.2

PyTorch

Verwenden Sie die TPU-Softwareversion, die der Version von PyTorch entspricht, mit der Ihr Modell geschrieben wurde. Wenn Sie beispielsweise PyTorch 1.13 und TPU v2 oder v3 verwenden, nutzen Sie die TPU-Softwareversion tpu-vm-pt-1.13. Wenn Sie TPU v4 verwenden, verwenden Sie die TPU-Softwareversion tpu-vm-v4-pt-1.13. Für TPU-Pods wird dieselbe TPU-Softwareversion verwendet (z.B. v2-32, v3-128, v4-32). Derzeit werden folgende TPU-Softwareversionen unterstützt:

TPU v2/v3:
  • tpu-vm-pt-2.0 (pytorch-2.0)
  • tpu-vm-pt-1.13 (pytorch-1.13)
  • tpu-vm-pt-1.12 (pytorch-1.12)
  • tpu-vm-pt-1.11 (pytorch-1.11)
  • tpu-vm-pt-1.10 (pytorch-1.10)
  • v2-alpha (pytorch-1.8.1)
TPU v4:
  • tpu-vm-v4-pt-2.0 (pytorch-2.0)
  • tpu-vm-v4-pt-1.13 (pytorch-1.13)

Wenn Sie eine TPU-VM erstellen, ist die neueste Version von PyTorch auf dieser vorinstalliert. Die richtige Version von libtpu.so wird bei der Installation von PyTorch automatisch installiert.

Informationen zum Ändern der aktuellen PyTorch-Softwareversion finden Sie unter PyTorch-Version ändern.

JAX

Sie müssen JAX manuell auf Ihrer TPU-VM installieren, da es keine JAX-spezifische TPU-Softwareversion gibt. Verwenden Sie für alle TPU-Versionen tpu-Ubuntu2204-base. Die richtige Version von libtpu.so wird bei der Installation von JAX automatisch installiert.

TPU-Knoten

TensorFlow

Verwenden Sie die TPU-Softwareversion, die der Version von TensorFlow entspricht, mit der Ihr Modell geschrieben wurde. Wenn Sie beispielsweise TF 2.12.0 nutzen, nutzen Sie die TPU-Softwareversion 2.12.0. Die TensorFlow-spezifischen TPU-Softwareversionen sind:

  • 2.12.1
  • 2.12.0
  • 2.11.1
  • 2.11.0
  • 2.10.1
  • 2.10.0
  • 2.9.3
  • 2.9.1
  • 2.8.4
  • 2.8.2
  • 2.7.3

Weitere Informationen zu TensorFlow-Patchversionen finden Sie unter Unterstützte TensorFlow-Patchversionen.

Wenn Sie einen TPU-Knoten erstellen, ist die neueste Version von TensorFlow auf dem TPU-Knoten vorinstalliert.

PyTorch

Verwenden Sie die TPU-Softwareversion, die der Version von PyTorch entspricht, mit der Ihr Modell geschrieben wurde. Wenn Sie beispielsweise PyTorch 1.9 verwenden, nutzen Sie die Softwareversion pytorch-1.9.

Die PyTorch-spezifischen TPU-Softwareversionen sind:

  • Pytorch-2.0
  • pytorch-1.13
  • pytorch-1.12
  • Pytorch-1.11
  • Pytorch-1.10
  • pytorch-1.9
  • pytorch-1.8
  • pytorch-1.7

  • pytorch-1.6

  • pytorch-nightly

Wenn Sie einen TPU-Knoten erstellen, ist die neueste Version von PyTorch auf dem TPU-Knoten vorinstalliert.

JAX

Sie müssen JAX manuell auf Ihrer TPU-VM installieren, sodass keine JAX-spezifische TPU-Softwareversion vorinstalliert ist. Sie können jede der für TensorFlow aufgeführten Softwareversionen verwenden.

Nächste Schritte