TPU-VM-Images

Wenn Sie TPU-Ressourcen erstellen, übergeben Sie den Parameter --version oder --runtime-version, der ein TPU-VM-Image angibt. TPU-VM-Images enthalten das Ubuntu-Betriebssystem, Docker und andere Software, die zum Ausführen Ihres Codes auf TPUs erforderlich ist. In diesem Dokument finden Sie eine Anleitung zum Auswählen des geeigneten TPU-VM-Images beim Erstellen von Cloud TPUs.

PyTorch und JAX

Verwenden Sie die folgenden gängigen TPU-VM-Basis-Images für PyTorch und JAX und installieren Sie dann das gewünschte Framework.

• tpu-ubuntu2204-base (Standard)
• v2-alpha-tpuv6e (TPU v6e)
• v2-alpha-tpuv5 (TPU v5p)
• v2-alpha-tpuv5-lite (TPU v5e)

Eine Anleitung zur Installation finden Sie in den Kurzanleitungen für PyTorch/XLA und JAX.

TensorFlow

Für jede Version von TensorFlow gibt es spezifische TPU-VM-Images. Die folgenden TensorFlow-Versionen werden auf Cloud TPUs unterstützt:

• 2.18.0
• 2.17.0
• 2.16.2
• 2.16.1
• 2.15.1
• 2.15.0
• 2.14.1
• 2.14.0
• 2.13.1
• 2.13.0
• 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.3
• 2.8.0
• 2.7.4
• 2.7.3

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

Für TensorFlow-Versionen ab 2.15.0 gibt es TPU-VM-Imagevarianten, die auf der von Ihnen verwendeten Geräte-API (PJRT oder Stream-Executor) basieren.

Training auf TPU v6e, v5p und v5e

TPU v6e, v5e und v5p unterstützen TensorFlow 2.15.0 und höher. Geben Sie das TPU-VM-Image im folgenden Format an: tpu-vm-tf-x.y.z-{pod}-pjrt. Dabei ist x die Hauptversion von TensorFlow, y die Nebenversion und z die TensorFlow-Patchversion. Fügen Sie nach der TensorFlow-Version pod hinzu, wenn Sie eine TPU mit mehreren Hosts verwenden. Wenn Sie beispielsweise TensorFlow 2.16.0 auf einer TPU mit mehreren Hosts verwenden, verwenden Sie das TPU-VM-Image tpu-vm-tf-2.16.0-pod-pjrt. Ersetzen Sie bei anderen TensorFlow-Versionen 2.16.0 durch die von Ihnen verwendete Haupt- und Patchversion von TensorFlow. Wenn Sie eine TPU mit einem einzelnen Host verwenden, lassen Sie pod weg.

Auslieferung über v6e und v5e

Es gibt Docker-Images für das Bereitstellen, die alle erforderlichen Softwareanforderungen für das Bereitstellen mit TensorFlow, PyTorch und JAX enthalten. Weitere Informationen finden Sie unter Einführung in die Cloud TPU v5e-Inferenz.

TPU v4

Wenn Sie TPU v4 und TensorFlow 2.15.0 oder höher verwenden, folgen Sie der Anleitung für das Training mit v6e, v5p und v5e. Wenn Sie TensorFlow 2.10.0 oder niedriger verwenden, verwenden Sie ein v4-spezifisches TPU-VM-Image:

TensorFlow-Version	Version des TPU-VM-Images
2.10.0	tpu-vm-tf-2.10.0-v4 tpu-vm-tf-2.10.0-pod-v4
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

Hinweis :Wenn Sie einen Pod-Speicherplatz verwenden, fügen Sie nach der TensorFlow-Versionsnummer -pod an. Beispiel: tpu-vm-tf-2.17.0-pod-pjrt.

TPU v2 und v3

Wenn Sie TPU v2 oder v3 verwenden, verwenden Sie das TPU-VM-Image, das der von Ihnen verwendeten TensorFlow-Version entspricht. Wenn Sie beispielsweise TensorFlow 2.14.1 verwenden, verwenden Sie das TPU-Image tpu-vm-tf-2.14.1. Bei anderen TensorFlow-Versionen ersetzen Sie 2.14.1 durch die von Ihnen verwendete TensorFlow-Version. Wenn Sie einen TPU-Pod mit mehreren Hosts verwenden, fügen Sie am Ende des TPU-Images einen Pod an, z. B. tpu-vm-tf-2.14.1-pod.

Ab TensorFlow 2.15.0 müssen Sie auch eine Geräte-API als Teil des Image-Namens angeben. Wenn Sie beispielsweise TensorFlow 2.16.1 mit der PJRT API verwenden, verwenden Sie das TPU-Image tpu-vm-tf-2.16.1-pjrt. Wenn Sie die Stream Executor API mit derselben Version von TensorFlow verwenden, verwenden Sie das TPU-Image tpu-vm-tf-2.16.1-se. TensorFlow-Versionen älter als 2.15.0 unterstützen nur den Stream-Executor.

Unterstützung für TensorFlow-PJRT

Ab TensorFlow 2.15.0 können Sie die PJRT-Oberfläche für TensorFlow auf TPU verwenden. PJRT bietet eine automatische Defragmentierung des Gerätespeichers und vereinfacht die Einbindung von Hardware in Frameworks. Weitere Informationen zu PJRT finden Sie unter PJRT: Vereinfachte Integration von ML-Hardware und ‑Frameworks.

Wichtig :TPU v5e und zukünftige TPU-Generationen unterstützen nur PJRT. Nicht alle Funktionen von TPU v2, v3 und v4 wurden in die PJRT API migriert. In der folgenden Tabelle wird beschrieben, welche Funktionen auf PJRT oder Stream-Executor unterstützt werden.

Beschleuniger	Funktion	PJRT-Unterstützung	Unterstützung für Stream-Executor
TPU v2 bis v4	Dichte Berechnung (keine TPU-Embedding API)	Ja	Ja
TPU v2 bis v4	Dense Compute API + TPU Embedding API	Nein	Ja
TPU v2 bis v4	tf.summary/tf.print mit flexibler Geräteplatzierung	Nein	Ja
TPU v5e	Dichte Berechnung (keine TPU-Embedding API)	Ja	Nein
TPU v5e	TPU-Embedding API	–	Nein
TPU v5p	Dichte Berechnung (keine TPU-Embedding API)	Ja	Nein
TPU v5p	TPU-Embedding API	Ja	Nein

Libtpu-Versionen

TensorFlow-Images für TPU-VMs enthalten eine bestimmte TensorFlow-Version und die entsprechende libtpu-Bibliothek. Wenn Sie Ihr eigenes VM-Image erstellen, verwenden Sie die folgenden TensorFlow TPU-Softwareversionen und die entsprechenden libtpu-Versionen:

TensorFlow-Version	libtpu.so-Version
2.18.0	1.12.0
2.17.0	1.11.0
2.16.2	1.10.1
2.16.1	1.10.1
2.15.1	1.9.0
2.15.0	1.9.0
2.14.1	1.8.1
2.14.0	1.8.0
2.13.1	1.7.1
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.0	1.2.0
2.7.3	1.1.2

Nächste Schritte

• Weitere Informationen zur TPU-Architektur finden Sie auf der Seite Systemarchitektur.
• Weitere Informationen zu den für Cloud TPU geeigneten Modellen finden Sie unter Verwendung von TPUs.