TPU-Softwareversionen

In diesem Dokument finden Sie eine Anleitung zur Auswahl der geeigneten TPU-Softwareversion beim Erstellen von Cloud TPUs.

Wenn Sie TPU-Ressourcen erstellen, geben Sie die auch als Laufzeitversion bezeichnete Softwareversion an. Sie bezieht sich auf die Softwareumgebung, die auf Ihrer TPU-VM vorinstalliert ist. Dazu gehören das Ubuntu-Betriebssystem, Docker und andere Software, die zum Ausführen Ihres Codes auf TPUs erforderlich ist.

Wenn Sie die Google Cloud CLI verwenden, geben Sie die TPU-Softwareversion mit dem Parameter --version oder --runtime-version an. Wenn Sie die Google Cloud Console verwenden, wählen Sie eine TPU-Softwareversion aus der Liste Version der TPU-Software aus.

PyTorch und JAX

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

TPU-Version TPU-Softwareversion
Trillium (v6e) v2-alpha-tpuv6e
v5p v2-alpha-tpuv5
v5e v2-alpha-tpuv5-lite
v4 und älter tpu-ubuntu2204-base

Weitere Informationen zur Installation und zum Einstieg in PyTorch oder JAX finden Sie unter Berechnung mit PyTorch auf einer Cloud TPU-VM ausführen und Berechnung mit JAX auf einer Cloud TPU-VM ausführen.

TensorFlow

TPU-Softwareversionen für TensorFlow folgen einer bestimmten Namenskonvention:

tpu-vm-tf-x.y.z[-{pod}][-{device_api}]

  • x.y.z: stellt die Haupt-, Neben- und Patchversionen von TensorFlow dar
  • -pod (optional): gibt an, dass Sie einen TPU-Slice mit mehreren Hosts verwenden
  • -{device_api} (optional): gibt die Geräte-API an, z. B. -pjrt (wenn Sie die PJRT API verwenden)

Weitere Informationen zum Angeben einer TPU-Softwareversion finden Sie in den folgenden Abschnitten.

Für jede TensorFlow-Version gibt es spezifische TPU-Softwareversionen. In der folgenden Tabelle sind die unterstützten TensorFlow-Versionen und die zugehörigen libtpu-Versionen aufgeführt:

TensorFlow-Version libtpu.so-Version
2.18.0 1.12.0
2.17.1 1.11.1
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

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

TPU v6e, v5p und v5e

TPU v6e, v5e und v5p unterstützen TensorFlow 2.15.0 und höher. Sie geben die TPU-Softwareversion im Format tpu-vm-tf-x.y.z-{pod}-pjrt an, wobei x die Hauptversion von TensorFlow, y die Nebenversion und z die TensorFlow-Patchversion ist. Fügen Sie pod nach der TensorFlow-Version 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 die TPU-Softwareversion tpu-vm-tf-2.16.0-pod-pjrt. Ersetzen Sie bei anderen TensorFlow-Versionen 2.16.0 durch die Haupt- und Patchversionen von TensorFlow, die Sie verwenden. Wenn Sie eine TPU mit einem einzelnen Host verwenden, lassen Sie pod weg.

TPU v4

Wenn Sie TPU v4 und TensorFlow 2.10.1 oder höher verwenden, folgen Sie der Anleitung für TPU v2 und v3. Wenn Sie TensorFlow 2.10.0 oder früher verwenden, verwenden Sie eine v4-spezifische TPU-Softwareversion:

TensorFlow-Version TPU-Softwareversion
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

TPU v2 und v3

Wenn Sie TPU v2 oder v3 verwenden, verwenden Sie die TPU-Softwareversion, die der von Ihnen verwendeten TensorFlow-Version entspricht. Wenn Sie beispielsweise TensorFlow 2.14.1 verwenden, verwenden Sie die TPU-Softwareversion tpu-vm-tf-2.14.1. Ersetzen Sie bei anderen TensorFlow-Versionen 2.14.1 durch die von Ihnen verwendete TensorFlow-Version. Wenn Sie eine TPU mit mehreren Hosts verwenden, hängen Sie „pod“ an das Ende der TPU-Softwareversion an, z. B. tpu-vm-tf-2.14.1-pod.

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

TensorFlow-Unterstützung für PJRT

Ab TensorFlow 2.15.0 können Sie die PJRT-Schnittstelle für TensorFlow auf TPUs verwenden. PJRT bietet automatische Defragmentierung des Gerätespeichers und vereinfacht die Integration von Hardware in Frameworks. Weitere Informationen zu PJRT finden Sie unter PJRT: Simplifying ML Hardware and Framework Integration.

Beschleuniger Funktion PJRT-Unterstützung Unterstützung für Stream Executor
TPU v2 bis v4 Dense Compute (keine TPU-Einbettungs-API) Ja Ja
TPU v2 bis v4 Dense-Compute-API + TPU-Einbettungs-API Nein Ja
TPU v2 bis v4 tf.summary/tf.print mit Soft Device Placement Nein Ja
TPU v5e Dense Compute (keine TPU-Einbettungs-API) Ja Nein
TPU v5e TPU-Einbettungs-API Nein
TPU v5p Dense Compute (keine TPU-Einbettungs-API) Ja Nein
TPU v5p TPU-Einbettungs-API Ja Nein

Nächste Schritte