TPU-VM-Images
Beim Erstellen von TPU-Ressourcen ü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-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-tpuv5 (TPU v5p)
Installationsanweisungen finden Sie in den Kurzanleitungen für PyTorch/XLA und JAX.
TensorFlow
Es gibt TPU-VM-Images, die für jede Version von TensorFlow spezifisch sind. 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.
Bei TensorFlow-Versionen 2.15.0 und höher 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 nach der TensorFlow-Version pod ein, 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 von Ihnen verwendete Haupt- und Patchversion von TensorFlow. Wenn Sie eine TPU mit einem einzelnen Host verwenden, lassen Sie pod weg.
In v5e bereitstellen
Es gibt bereitstellende 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 mit v5p und v5e. Wenn Sie TensorFlow 2.10.0 oder eine frühere Version nutzen, 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 |
TPU v2 und v3
Verwenden Sie bei Verwendung von TPU v2 oder v3 das TPU-VM-Image, das der von Ihnen verwendeten Version von TensorFlow entspricht. Wenn Sie beispielsweise TensorFlow 2.14.1 verwenden, nutzen 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 einen TPU-Pod mit mehreren Hosts verwenden, hängen Sie an das Ende des TPU-Images an, z. B. tpu-vm-tf-2.14.1-pod.
Ab TensorFlow 2.15.0 müssen Sie außerdem 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 den Stream Executor.
Unterstützung für TensorFlow PJRT
Ab TensorFlow 2.15.0 können Sie die PJRT-Schnittstelle für TensorFlow auf TPU verwenden. PJRT ermöglicht eine 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.
| Beschleuniger | Funktion | PJRT-Unterstützung | Stream-Executor-Unterstützung |
|---|---|---|---|
| TPU v2–v4 | Dichtes Computing (keine TPU Embedding API) | Yes | Yes |
| TPU v2–v4 | Dense Compute API + TPU Embedding API | Nein | Yes |
| TPU v2–v4 | tf.summary/tf.print mit Soft-Device-Platzierung | Nein | Yes |
| TPU v5e | Dichtes Computing (keine TPU Embedding API) | Yes | Nein |
| TPU v5e | TPU Embedding API | – | Nein |
| TPU v5p | Dichtes Computing (keine TPU Embedding API) | Yes | Nein |
| TPU v5p | TPU Embedding 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 entsprechende 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 Verwendung von TPUs.