Changement de version du logiciel sur votre Cloud TPU

Présentation

La version logicielle du framework exécutée sur votre TPU doit correspondre à la version exécutée sur votre VM locale. Cette version logicielle peut désormais être remplacée par un Cloud TPU en cours d'exécution, sans supprimer ni recréer le TPU. Cela permet également de configurer le Cloud TPU avec des versions automatiques spécifiques des frameworks logiciels. Il est toujours recommandé de sélectionner une version compatible de ces frameworks.

Utilisation

La méthode recommandée pour changer de version consiste à utiliser la bibliothèque Python cloud-tpu-client.

Exemple d'utilisation de TensorFlow.

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

Le Cloud TPU est configuré de manière à correspondre à la version de TensorFlow exécutée sur votre VM locale, y compris les versions officielles et les builds automatiques datés.

Le paramètre restart_type de l'API configure_tpu_version définit le comportement de redémarrage du TPU lors du changement de version. Les options sont "always" (par défaut) et "ifNeeded".

  • Vous pouvez utiliser l'option "always" pour corriger un TPU avec, par exemple, l'état UNHEALTHY_TENSORFLOW, ou qui renvoie des erreurs de mémoire saturée (OOM, Out Of Memory) en raison de fuites de ressources lors d'une exécution précédente. Si cette option est définie, le TPU est redémarré, même lorsqu'une nouvelle version du framework n'est pas installée.

  • L'option "ifNeeded" peut être utile, car elle ne redémarre pas l'environnement d'exécution s'il est déjà dans la bonne version. Elle n'ajoute donc aucun temps de démarrage significatif à un script d'entraînement. Lorsque cette option est définie, le TPU ne redémarre que si la version de framework appropriée n'est pas installée.

La bibliothèque communique directement avec Cloud TPU. Ce code doit donc être exécuté sur une VM du même réseau. Nous vous recommandons de l'exécuter dans le code du reste de votre modèle.

Options logicielles supplémentaires

Tensorflow inclut une chaîne tf.__version__ qui constitue le moyen le plus simple de configurer la version appropriée. Les autres options logicielles sont les suivantes :

  • PyTorch - pytorch-1.13, pytorch-nightly-dev20220930, pytorch-nightly
  • Jax - tpu_driver, tpu_driver0.1-dev20200320, tpu_driver_nightly

Par exemple, pour configurer un TPU à exécuter avec la dernière version nocturne de PyTorch.

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()