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 Betriebssystem (Ubuntu) und optional weitere Software, die zum Ausführen Ihres Codes auf TPUs erforderlich ist. Dieses Dokument enthält eine Anleitung zur Auswahl des geeigneten TPU-VM-Image 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 Framework, das Sie verwenden möchten.

• tpu-Ubuntu2204-base (Standard)
• v2-alpha-tpuv5 (TPU v5p)

Hinweis: Diese Images sind in der Google Cloud Console nicht verfügbar. Sie müssen mit dem gcloud-Befehl verwendet werden.

Installationsanweisungen finden Sie in den Schnellstartdokumenten 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.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.

Ab TensorFlow-Version 2.15.0 gibt es TPU-VM-Image-Varianten, die auf der von Ihnen verwendeten Geräte-API (PJRT oder Stream-Executor) basieren.

Training zu v5p und v5e

TPU v5e und v5p unterstützen TensorFlow 2.15.0 und höher. Sie geben das TPU-VM-Image im folgenden Format an: tpu-vm-tf-x.y.z-{pod}-pjrt, wobei x die TensorFlow-Hauptversion, 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 das TPU-VM-Image tpu-vm-tf-2.16.0-pod-pjrt. Ersetzen Sie bei anderen Versionen von TensorFlow 2.16.0 durch die verwendete Haupt- und Patchversion von TensorFlow. Wenn Sie eine TPU mit einem einzelnen Host verwenden, lassen Sie pod weg.

Bereitstellung in Version 5e

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

TPU v4

Wenn Sie TPU v4 und TensorFlow 2.15.0 oder höher verwenden, folgen Sie der Anleitung zum Training auf v5p und v5e. Wenn Sie TensorFlow 2.10.0 oder eine ältere Version verwenden, verwenden Sie ein v4-spezifisches TPU-VM-Image:

TensorFlow-Version	TPU-VM-Image-Version
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 ein Pod-Slice verwenden, hängen Sie -pod nach der TensorFlow-Versionsnummer an. Beispiel: tpu-vm-tf-2.16.1-pod-pjrt.

TPU v2 und v3

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

Ab TensorFlow 2.15.0 müssen Sie auch eine Geräte-API 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, die älter als 2.15.0 sind, unterstützen nur Stream Executor.

PJRT-Unterstützung für TensorFlow

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

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

Beschleuniger	Feature	PJRT-Unterstützung	Unterstützung für Stream Executor
TPU v2–v4	Dichtes Computing (keine TPU-Einbettungs-API)	Yes	Yes
TPU v2–v4	Dense Compute API + TPU Embedding API	Nein	Yes
TPU v2–v4	tf.summary/tf.print mit Soft Device Placement	Nein	Yes
TPU v5e	Dichtes Computing (keine TPU-Einbettungs-API)	Yes	Nein
TPU v5e	TPU-Einbettungs-API	–	Nein
TPU v5p	Dichtes Computing (keine TPU-Einbettungs-API)	Yes	Nein
TPU v5p	TPU-Einbettungs-API	Yes	Nein

Libtpu-Versionen

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

TensorFlow-Version	libtpu.so-Version
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 Wann TPUs verwenden.