Images de VM TPU

Lorsque vous créez des ressources TPU, vous transmettez le paramètre --version ou --runtime-version, qui spécifie une image de VM TPU. Les images de VM TPU contiennent le système d'exploitation (Ubuntu) et et éventuellement d'autres logiciels nécessaires pour exécuter votre code sur des TPU. Ce document fournit des conseils sur la sélection de l'image de VM TPU appropriée lors de la création Cloud TPU.

PyTorch et JAX

Utilisez les images de base de VM TPU courantes suivantes pour PyTorch et JAX, puis installez le framework que vous souhaitez utiliser.

  • tpu-Ubuntu2204-base (par défaut)
  • v2-alpha-tpuv5 (TPU v5p)
  • v2-alpha-tpuv5-lite (TPU v5e)

Pour obtenir des instructions d'installation, consultez les documents de démarrage rapide pour PyTorch/XLA et JAX.

TensorFlow

Des images de VM TPU sont propres à chaque version de TensorFlow. La Les versions suivantes de TensorFlow sont compatibles avec les Cloud TPU:

  • 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

Pour en savoir plus sur les versions du correctif TensorFlow, consultez Versions de correctif TensorFlow compatibles

Pour TensorFlow version 2.15.0 et ultérieure, il existe des variantes d'image de VM TPU en fonction de l'API de l'appareil (PJRT ou exécuteur de flux) que vous utilisez.

Entraînement sur v5p et v5e

Les TPU v5e et v5p sont compatibles avec TensorFlow 2.15.0 et versions ultérieures. Vous spécifiez la VM TPU image au format suivant: tpu-vm-tf-x.y.z-{pod}-pjrt, où x est le principal La version de TensorFlow, y est la version mineure et z est la version Version du correctif TensorFlow. Ajouter pod après la version de TensorFlow si vous utilisez un TPU multi-hôtes. Par exemple, si vous utilisez TensorFlow 2.16.0 sur un TPU multi-hôtes, utilisez l'image de VM TPU tpu-vm-tf-2.16.0-pod-pjrt. Pour autres versions de TensorFlow, remplacez 2.16.0 par le numéro majeur et effectuez un correctif. de TensorFlow que vous utilisez. Si vous utilisez un TPU à un seul hôte, omettez pod.

Diffusion sur la version v5e

Il existe des images Docker de diffusion qui contiennent toutes les exigences logicielles nécessaires pour la diffusion avec TensorFlow, PyTorch et JAX. Pour en savoir plus, consultez Présentation de l'inférence Cloud TPU v5e

TPU v4

Si vous utilisez TPU v4 et TensorFlow version 2.15.0 ou ultérieure, suivez les instructions instructions pour l'entraînement sur les versions v5p et v5e. Si vous utilisez TensorFlow 2.10.0 ou version antérieure, utilisez une image de VM TPU spécifique à la version v4:

Version de TensorFlow Version de l'image de la VM TPU
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 et v3

Si vous utilisez TPU v2 ou v3, utilisez l'image de VM TPU correspondant à la version de TensorFlow que vous utilisez. Par exemple, si vous utilisez TensorFlow 2.14.1, utilisez l'image TPU tpu-vm-tf-2.14.1. Pour les autres versions de TensorFlow, remplacez 2.14.1 par la version de TensorFlow utilisent. Si vous utilisez un TPU multi-hôte, ajoutez le pod à la fin du TPU (par exemple, tpu-vm-tf-2.14.1-pod).

À partir de la version 2.15.0 de TensorFlow, vous devez également spécifier une API d'appareil. Par exemple, si vous utilisez TensorFlow 2.16.1 avec l'API PJRT, utilisez l'image TPU tpu-vm-tf-2.16.1-pjrt. Si vous utilisez l'API d'exécuteur de flux avec la même version de TensorFlow, utilisez le TPU tpu-vm-tf-2.16.1-se l'image. Les versions de TensorFlow antérieures à la version 2.15.0 ne sont compatibles qu'avec l'exécuteur de flux.

Compatibilité avec TensorFlow PJRT

À partir de la version 2.15.0 de TensorFlow, vous pouvez utiliser l'interface PJRT pour TensorFlow sur TPU. PJRT propose une défragmentation automatique de la mémoire de l'appareil et simplifie l'intégration du matériel aux frameworks. Pour plus d'informations PJRT, consultez PJRT : "Simplifier l'intégration du matériel et des frameworks de ML"

Accélérateur Caractéristique Compatibilité avec PJRT Prise en charge de l'exécuteur de flux
TPU v2 – v4 Calcul dense (pas d'API de représentation vectorielle continue de TPU) Oui Oui
TPU v2 – v4 API de calcul dense + API d'intégration de TPU Non Oui
TPU v2 – v4 tf.summary/tf.print avec emplacement d'appareil virtuel Non Oui
TPU v5e Calcul dense (pas d'API de représentation vectorielle continue de TPU) Oui Non
TPU v5e API d'embedding TPU N/A Non
TPU v5p Calcul dense (pas d'API de représentation vectorielle continue de TPU) Oui Non
TPU v5p API TPU Embedding Oui Non

Versions Libtpu

Les images TensorFlow des VM TPU contiennent une version spécifique de TensorFlow et la bibliothèque libtpu correspondante. Si vous créez votre propre image de VM, utilisez les versions logicielles des TPU TensorFlow suivantes et les versions correspondantes Versions libtpu:

Version de TensorFlow Version de libtpu.so
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

Étape suivante