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. Dieses Dokument enthält eine Anleitung zur Auswahl des geeigneten TPU-VM-Images, wenn erstellen Sie 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)
  • 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 folgende TensorFlow-Versionen werden auf Cloud TPUs unterstützt:

  • 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 2.15.0 und höher gibt es TPU-VM-Image-Varianten basierend auf der verwendeten Geräte-API (PJRT oder Stream-Executor).

Training zu v5p und v5e

TPU 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 z. B. TensorFlow verwenden, 2.16.0 auf einer TPU mit mehreren Hosts: Verwenden Sie das TPU-VM-Image tpu-vm-tf-2.16.0-pod-pjrt. Für anderen Versionen von TensorFlow, ersetzen Sie 2.16.0 durch die Hauptversion und den Patch. Version von TensorFlow, die Sie verwenden. Wenn Sie eine TPU mit einem einzelnen Host verwenden, lassen Sie pod weg.

Auslieferung über 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 auf v5p und v5e Wenn Sie TensorFlow 2.10.0 oder niedriger 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

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 z. B. TensorFlow verwenden, 2.14.1 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 eine TPU mit mehreren Hosts verwenden, hängen Sie den Pod an das Ende der TPU an. Bild, 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.

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