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, Docker et d'autres logiciels requis pour exécuter votre code sur les TPU. Ce document fournit des conseils pour sélectionner l'image de VM TPU appropriée lorsque vous créez des 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-tpuv6e (TPU v6e)
  • 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

Il existe des images de VM TPU spécifiques à chaque version de TensorFlow. Les versions de TensorFlow suivantes sont compatibles avec les Cloud TPU:

  • 2.18.0
  • 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 de correctifs TensorFlow, consultez la section Versions de correctifs TensorFlow compatibles.

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

Entraînement sur les versions v6e, v5p et v5e

Les TPU v6e, v5e et v5p sont compatibles avec TensorFlow 2.15.0 et versions ultérieures. Vous spécifiez l'image de VM TPU à l'aide du format tpu-vm-tf-x.y.z-{pod}-pjrt, où x correspond à la version majeure de TensorFlow, y à la version mineure et z à la version de correctif de TensorFlow. Ajoutez pod après la version de TensorFlow si vous utilisez un TPU multi-hôte. Par exemple, si vous utilisez TensorFlow 2.16.0 sur un TPU multi-hôte, 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 seul hôte TPU, omettez pod.

Diffusion sur v6e et 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 la présentation de l'inférence Cloud TPU v5e.

TPU v4

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

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 TPU 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ôte, ajoutez un pod à la fin de l'image TPU, par exemple tpu-vm-tf-2.14.1-pod.

À partir de TensorFlow 2.15.0, vous devez également spécifier une API d'appareil dans le nom de l'image. 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 de l'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 à 2.15.0 ne sont compatibles qu'avec l'exécuteur de flux.

Compatibilité avec TensorFlow PJRT

À partir de TensorFlow 2.15.0, vous pouvez utiliser l'interface PJRT pour TensorFlow sur TPU. PJRT offre 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: simplifier l'intégration du matériel et du framework ML.

Accélérateur Caractéristique Compatibilité avec PJRT Prise en charge de l'exécuteur de flux
TPU v2 à v4 Calcul dense (aucune API d'encapsulation TPU) Oui Oui
TPU v2 à v4 API de calcul dense + API d'embedding TPU Non Oui
TPU v2 à v4 tf.summary/tf.print avec placement d'appareil souple Non Oui
TPU v5e Calcul dense (aucune API d'encapsulation TPU) Oui Non
TPU v5e API d'embedding TPU N/A Non
TPU v5p Calcul dense (aucune API d'encapsulation TPU) Oui Non
TPU v5p API d'embedding TPU Oui Non

Versions de 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 de TPU TensorFlow et les versions libtpu correspondantes suivantes:

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