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 éventuellement d'autres logiciels requis pour exécuter votre code sur les TPU. Ce document fournit des conseils sur la sélection de l'image de VM TPU appropriée lors de la création de Cloud TPU.

PyTorch et JAX

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

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

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. Les versions suivantes de TensorFlow sont compatibles avec les Cloud TPU:

  • 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 la page Versions de correctif TensorFlow compatibles.

Pour les versions 2.15.0 et ultérieures de TensorFlow, il existe des variantes d'image de VM TPU basées sur 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 les versions ultérieures. Spécifiez l'image de VM TPU au format suivant: tpu-vm-tf-x.y.z-{pod}-pjrt, où x est la version majeure de TensorFlow, y la version mineure et z la version du correctif TensorFlow. Si vous utilisez un TPU multi-hôtes, ajoutez pod après la version de TensorFlow. 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 les autres versions de TensorFlow, remplacez 2.16.0 par les versions majeures et correctives de TensorFlow que vous utilisez. Si vous utilisez un TPU à un seul hôte, omettez pod.

Diffusion dans la version v5e

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

TPU v4

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

Version de TensorFlow Version de l'image de 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 que vous utilisez. Si vous utilisez un TPU multi-hôtes, ajoutez le pod à la fin de l'image 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 l'image TPU tpu-vm-tf-2.16.1-se. 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 intègre une défragmentation automatique de la mémoire de l'appareil et simplifie l'intégration du matériel aux frameworks. Pour en savoir plus sur PJRT, consultez PJRT: simplification de l'intégration du matériel et du framework de ML.

Accélérateur Sélection Compatibilité 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 TPU Embedding Non disponible 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 de libtpu correspondantes:

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

Étapes suivantes