TPU-VM-Images

Wenn Sie TPU-Ressourcen erstellen, übergeben Sie die --version oder --runtime-version Parameter, 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 Anleitungen zur Auswahl des geeigneten TPU-VM-Images beim Erstellen Cloud TPUs

PyTorch und JAX

Verwenden Sie die folgenden gängigen TPU-VM-Basis-Images für PyTorch und JAX und installieren Sie sie dann das Framework, das Sie verwenden möchten.

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

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 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. Sie geben die TPU-VM an Bild mit dem Format tpu-vm-tf-x.y.z-{pod}-pjrt, wobei x der wichtigste TensorFlow-Version, y ist die Nebenversion und z die TensorFlow-Patchversion. pod nach der TensorFlow-Version hinzufügen 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.

Bereitstellung in Version 5e

Es werden Docker-Images bereitgestellt, die alle erforderlichen Softwareanforderungen für mit TensorFlow, PyTorch und JAX. 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 für das Training auf v5p und v5e Wenn Sie Verwenden Sie ab TensorFlow 2.10.0 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 Version von TensorFlow, das Sie verwenden. Wenn Sie z. B. TensorFlow verwenden, 2.14.1 verwenden Sie das TPU-Image tpu-vm-tf-2.14.1. Für andere Versionen von TensorFlow, ersetzen Sie 2.14.1 durch die TensorFlow-Version, die Sie verwenden. 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 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. Bei Verwendung der Stream Executor API mit derselben Version von TensorFlow, verwenden Sie die tpu-vm-tf-2.16.1-se-TPU Bild. 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 PJRT-Funktion zur automatischen Defragmentierung des Gerätespeichers und vereinfacht die Integration von Hardware in Frameworks. Weitere Informationen zu PJRT, siehe PJRT: Vereinfachung der Integration von ML-Hardware und ‐Frameworks

Beschleuniger Feature PJRT-Unterstützung Unterstützung für Stream Executor
TPU v2–v4 Dichtes Computing (keine TPU-Einbettungs-API) Ja Ja
TPU v2–v4 Dense Compute API + TPU Embedding API Nein Ja
TPU v2–v4 tf.summary/tf.print mit Soft Device Placement Nein Ja
TPU v5e Dichtes Computing (keine TPU-Einbettungs-API) Ja Nein
TPU v5e TPU-Einbettungs-API – Nein
TPU v5p Dichtes Computing (keine TPU-Einbettungs-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 ein eigenes VM-Image erstellen, die folgenden TensorFlow TPU-Softwareversionen und entsprechende 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