Utiliser des TPU pour entraîner un modèle

Les TPU (Tensor Processing Unit) sont des ASIC spécifiquement développés par Google pour accélérer les charges de travail de machine learning. Vous pouvez exécuter vos tâches d'entraînement sur AI Platform Training à l'aide de Cloud TPU. AI Platform Training fournit une interface de gestion des tâches qui vous évite de gérer vous-même le TPU. À la place, vous pouvez utiliser l'API jobs d'AI Platform Training de la même façon que pour un entraînement sur un processeur ou un GPU.

Les API TensorFlow de haut niveau vous permettent d'exécuter vos modèles sur le matériel Cloud TPU.

Configurer votre environnement Google Cloud

Configurez votre environnement Google Cloud en parcourant la section de configuration du guide de démarrage.

Autoriser Cloud TPU à accéder à votre projet

Pour autoriser le nom de compte de service Cloud TPU associé à votre projet Google Cloud, procédez comme suit :

1. Obtenez le nom de votre compte de service Cloud TPU en appelant projects.getConfig. Exemple :

   PROJECT_ID=PROJECT_ID
   
   curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
       https://ml.googleapis.com/v1/projects/$PROJECT_ID:getConfig
   

2. Enregistrez les valeurs des champs serviceAccountProject et tpuServiceAccount renvoyées par l'API.

3. Initialisez le compte de service Cloud TPU :

   curl -H "Authorization: Bearer $(gcloud auth print-access-token)"  \
     -H "Content-Type: application/json" -d '{}'  \
     https://serviceusage.googleapis.com/v1beta1/projects/<serviceAccountProject>/services/tpu.googleapis.com:generateServiceIdentity
   

Ajoutez ensuite le compte de service Cloud TPU en tant que membre de votre projet, avec le rôle Agent de service Cloud ML. Effectuez les étapes suivantes dans Google Cloud Console ou à l'aide de l'outil de ligne de commande gcloud :

Pour en savoir plus sur l'attribution de rôles aux comptes de service, consultez la documentation IAM.

Exemple : Entraîner un exemple de modèle MNIST

Cette section explique comment entraîner un exemple de modèle MNIST à l'aide d'un TPU et de la version d'exécution 2.11. L'exemple de tâche utilise le niveau d'évolutivité prédéfini BASIC_TPU pour la configuration de votre machine. Les sections suivantes du guide vous expliquent comment mettre en place une configuration personnalisée.

Cet exemple suppose que vous utilisez un shell Bash avec la CLI gcloud installée. Exécutez les commandes suivantes pour obtenir le code et envoyer votre tâche d'entraînement à AI Platform Training :

1. Téléchargez le code des modèles de référence de TensorFlow et accédez au répertoire contenant l'exemple de code :

   git clone https://github.com/tensorflow/models.git \
     --branch=v2.11.0 \
     --depth=1
   
   cd models
   

2. Créez un fichier setup.py dans le répertoire models. Cela permet de s'assurer que la commande gcloud ai-platform jobs submit training inclut tous les sous-packages nécessaires dans le répertoire models/official lors de la création d'un package tarball de votre code d'entraînement, et qu'AI Platform Training installe les ensembles de données TensorFlow en tant que dépendance lors de l'exécution de la tâche d'entraînement. Ce code d'entraînement s'appuie sur des ensembles de données TensorFlow pour charger les données MNIST.

   Pour créer le fichier setup.py, exécutez la commande suivante dans votre interface système :

   cat << END > setup.py
   from setuptools import find_packages
   from setuptools import setup
   
   setup(
       name='official',
       install_requires=[
          'tensorflow-datasets~=3.1',
          'tensorflow-model-optimization>=0.4.1'
      ],
       packages=find_packages()
   )
   END
   

3. Envoyez votre tâche d'entraînement à l'aide de la commande gcloud ai-platform jobs submit training :

   gcloud ai-platform jobs submit training tpu_mnist_1 \
     --staging-bucket=gs://BUCKET_NAME \
     --package-path=official \
     --module-name=official.vision.image_classification.mnist_main \
     --runtime-version=2.11 \
     --python-version=3.7 \
     --scale-tier=BASIC_TPU \
     --region=us-central1 \
     -- \
     --distribution_strategy=tpu \
     --data_dir=gs://tfds-data/datasets \
     --model_dir=gs://BUCKET_NAME/tpu_mnist_1_output
   

   Remplacez BUCKET_NAME par le nom d'un bucket Cloud Storage dans votre projet Google Cloud. La CLI gcloud importe votre code d'entraînement empaqueté dans ce bucket et AI Platform Training enregistre les résultats d'entraînement dans le bucket.

4. Surveillez votre tâche d'entraînement. Une fois la tâche terminée, vous pouvez afficher son résultat dans le répertoire gs://BUCKET_NAME/tpu_mnist_1_output.

En savoir plus sur l'entraînement d'un modèle sur Cloud TPU

Cette section explique comment configurer une tâche et entraîner un modèle sur AI Platform Training avec Cloud TPU.

Spécifier une région qui offre un accès aux TPU

Vous devez exécuter votre tâche dans une région où les TPU sont disponibles. Les régions suivantes offrent actuellement un accès aux TPU :

• us-central1
• europe-west4

Pour bien comprendre quelles régions sont disponibles pour les services AI Platform Training, y compris l'entraînement de modèles et la prédiction en ligne ou par lot, consultez le guide des régions.

Gérer les versions de TensorFlow et AI Platform Training

Les versions d'exécution 1.15, 2.1, 2.2, 2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9 et 2.11 de AI Platform Training sont disponibles pour l'entraînement de vos modèles sur Cloud TPU. Informez-vous sur les versions d'exécution d'AI Platform Training et les versions correspondantes de TensorFlow.

Le règlement de gestion des versions est le même que pour Cloud TPU. Dans la requête de tâche d'entraînement, veillez à spécifier une version d'exécution disponible pour les TPU et correspondant à la version de TensorFlow utilisée dans votre code d'entraînement.

Se connecter au serveur TPU gRPC

Dans votre programme TensorFlow, utilisez TPUClusterResolver pour vous connecter au serveur TPU gRPC, qui s'exécute sur la VM TPU.

Le guide d'utilisation des TPU TensorFlow explique comment utiliser TPUClusterResolver avec la stratégie de distribution TPUStrategy.

Cependant, vous devez apporter une modification importante lorsque vous utilisez TPUClusterResolver pour le code exécuté sur AI Platform Training. Ne fournissez aucun argument lorsque vous construisez l'instance TPUClusterResolver. Lorsque les arguments de mot clé tpu, zone et project sont tous définis sur la valeur par défaut None, AI Platform Training fournit automatiquement au résolveur de cluster les informations de connexion nécessaires via des variables d'environnement.

L'exemple suivant de TensorFlow 2 montre comment initialiser un résolveur de cluster et une stratégie de distribution pour l'entraînement sur AI Platform Training :

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.experimental.TPUStrategy(resolver)

Utiliser des TPU dans le code TensorFlow

Pour utiliser les TPU sur une machine, utilisez l'API TPUStrategy de TensorFlow 2. Le guide d'utilisation des TPU TensorFlow explique comment procéder.

Pour effectuer un entraînement avec des TPU dans TensorFlow 1, vous pouvez utiliser l'API TPUEstimator. Le guide Cloud TPU concernant l'API TPUEstimator explique comment procéder.

La documentation Cloud TPU fournit également une liste des opérations TensorFlow de bas niveau disponibles sur Cloud TPU.

Utiliser des TPU dans le code PyTorch

Pour utiliser un TPU lorsque vous utilisez un conteneur PyTorch prédéfini, utilisez le package torch_xla. Découvrez comment utiliser torch_xla pour les TPU lors des entraînements dans la documentation de PyTorch. Pour plus d'exemples de l'utilisation de torch_xla, consultez les tutoriels dans le dépôt GitHub de PyTorch XLA.

Notez que lorsque vous effectuez un entraînement avec un TPU sur AI Platform Training, vous utilisez un seul appareil XLA, et non plusieurs appareils XLA.

Consultez également la section suivante sur cette page concernant la configuration de votre tâche d'entraînement pour PyTorch et les TPU.

Configurer une machine TPU personnalisée

Une tâche d'entraînement TPU s'exécute sur une configuration à deux machines virtuelles. Une VM (maîtresse) exécute votre code Python. La VM maître pilote le serveur TensorFlow en cours d'exécution sur un nœud de calcul TPU.

Pour utiliser un TPU avec AI Platform Training, choisissez l'une des options suivantes pour configurer votre tâche d'entraînement afin d'accéder à une machine compatible TPU :

• Utilisez le niveau d'évolutivité BASIC_TPU. Vous pouvez employer cette méthode pour accéder aux accélérateurs TPU v2.
• Utilisez un nœud de calcul cloud_tpu et un ancien type de machine pour la VM maître. Vous pouvez employer cette méthode pour accéder aux accélérateurs TPU v2.
• Utilisez un nœud de calcul cloud_tpu et un type de machine Compute Engine pour la VM maître. Vous pouvez employer cette méthode pour accéder aux accélérateurs TPU v2 ou TPU v3. Les accélérateurs TPU v3 sont disponibles en version bêta.

Machine compatible TPU de base

Définissez le niveau d'évolutivité sur BASIC_TPU pour obtenir une VM maître ainsi qu'une VM TPU comprenant un TPU doté de huit cœurs TPU v2, de la même manière que pour l'exécution de l'exemple précédent.

Travailleur TPU dans une configuration d'ancien type de machine

Vous pouvez également établir une configuration de machine personnalisée si vous avez besoin de davantage de ressources informatiques sur la VM maître :

• Définissez le niveau d'évolutivité sur CUSTOM.
• Configurez la VM maître pour utiliser un ancien type de machine adapté à votre tâche.
• Définissez workerType sur cloud_tpu pour obtenir une VM TPU comprenant un Cloud TPU doté de huit cœurs TPU v2.
• Définissez workerCount sur 1.
• Ne spécifiez pas de serveur de paramètres lorsque vous utilisez un Cloud TPU. Le service rejette la requête de tâche si la valeur parameterServerCount est supérieure à zéro.

L'exemple suivant décrit un fichier config.yaml qui utilise ce type de configuration :

trainingInput:
  scaleTier: CUSTOM
  masterType: complex_model_m
  workerType: cloud_tpu
  workerCount: 1

Travailleur TPU dans une configuration de type de machine Compute Engine

Vous pouvez également établir une configuration de machine personnalisée avec un type de machine Compute Engine pour votre VM maître et un élément acceleratorConfig associé à votre VM TPU.

Vous pouvez utiliser ce type de configuration pour configurer un travailleur TPU doté de huit cœurs TPU v2 (ce qui est semblable à une configuration sans acceleratorConfig) ou un travailleur TPU doté de huit cœurs TPU v3 (bêta). Découvrez la différence entre les accélérateurs TPU v2 et TPU v3.

L'utilisation d'un type de machine Compute Engine offre également plus de flexibilité pour la configuration de votre VM maître :

• Définissez le niveau d'évolutivité sur CUSTOM.
• Configurez la VM maître de façon qu'elle utilise un type de machine Compute Engine répondant aux exigences de la tâche.
• Définissez workerType sur cloud_tpu.
• Ajoutez un élément workerConfig avec un champ acceleratorConfig. Dans ce champ acceleratorConfig, définissez type sur TPU_V2 ou TPU_V3 et count sur 8. Vous ne pouvez indiquer aucun autre nombre de cœurs de TPU.
• Définissez workerCount sur 1.
• Ne spécifiez pas de serveur de paramètres lorsque vous utilisez un Cloud TPU. Le service rejette la requête de tâche si la valeur parameterServerCount est supérieure à zéro.

L'exemple suivant décrit un fichier config.yaml qui utilise ce type de configuration :

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2
      count: 8

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3
      count: 8

Utiliser des pods TPU

Un pod TPU est un ensemble d'appareils TPU connectés ensemble par des interfaces réseau haut débit dédiées. Un pod TPU peut comporter jusqu'à 2048 cœurs de TPU, ce qui vous permet de répartir la charge de traitement sur plusieurs TPU.

Pour utiliser des pods TPU, vous devez d'abord déposer une demande d'augmentation de quota.

Les exemples de fichiers config.yaml suivants montrent comment utiliser les pods TPU :

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V2_POD
      count: 128

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    acceleratorConfig:
      type: TPU_V3_POD
      count: 32

Le nombre de cœurs de pod pouvant être utilisés pour chaque type de TPU est limité. Configurations disponibles :

Pour savoir comment utiliser pleinement les cœurs de pods TPU, consultez la documentation Cloud TPU sur les pods TPU.

Utiliser un conteneur PyTorch prédéfini sur un travailleur TPU

Si vous souhaitez effectuer un entraînement PyTorch avec un TPU, vous devez spécifier le champ tpuTfVersion dans le paramètre trainingInput de votre tâche d'entraînement. Définissez tpuTfVersion pour qu'il corresponde à la version du conteneur PyTorch prédéfini que vous utilisez pour l'entraînement.

AI Platform Training est compatible avec l'entraînement à l'aide de TPU pour les conteneurs PyTorch préconfigurés suivants :

Par exemple, pour effectuer l'entraînement à l'aide du conteneur prédéfini PyTorch 1.11, vous pouvez utiliser le fichier config.yaml suivant pour configurer l'entraînement :

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/cloud-ml-public/training/pytorch-xla.1-11
    tpuTfVersion: pytorch-1.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

Consultez également la section précédente sur cette page concernant l'utilisation des TPU dans le code PyTorch.

Utiliser un conteneur personnalisé sur un travailleur TPU

Si vous souhaitez exécuter un conteneur personnalisé sur votre travailleur TPU plutôt que d'utiliser l'une des versions d'exécution d'AI Platform Training compatibles avec les TPU, vous devez spécifier un champ de configuration supplémentaire lorsque vous envoyez votre tâche d'entraînement. Définissez tpuTfVersion sur une version d'exécution incluant la version de TensorFlow utilisée par votre conteneur. Vous devez spécifier une version d'exécution actuellement compatible avec l'entraînement à l'aide de TPU.

Étant donné que vous configurez votre tâche pour qu'elle fasse appel à un conteneur personnalisé, AI Platform Training n'utilise pas l'environnement de cette version d'exécution lorsqu'il exécute votre tâche d'entraînement. Néanmoins, il a besoin de ce champ afin de préparer correctement le travailleur TPU pour la version de TensorFlow utilisée par votre conteneur personnalisé.

L'exemple suivant montre un fichier config.yaml avec une configuration TPU semblable à celle de la section précédente, à la différence près que la VM maître et le travailleur TPU exécutent ici chacun des conteneurs personnalisés distincts :

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V2
      count: 8

trainingInput:
  scaleTier: CUSTOM
  masterType: n1-highcpu-16
  masterConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-master-image-name:your-master-tag-name
  workerType: cloud_tpu
  workerCount: 1
  workerConfig:
    imageUri: gcr.io/YOUR_PROJECT_ID/your-worker-image-name:your-worker-tag-name
    tpuTfVersion: 2.11
    acceleratorConfig:
      type: TPU_V3
      count: 8

Si vous exécutez la commande gcloud beta ai-platform jobs submit training pour envoyer votre tâche d'entraînement, vous pouvez spécifier le champ d'API tpuTfVersion avec l'option --tpu-tf-version au lieu d'un fichier config.yaml.

Utiliser TPUClusterResolver après le provisionnement du TPU

Lorsque vous utilisez un conteneur personnalisé, vous devez attendre que le TPU soit provisionné avant de pouvoir appeler TPUClusterResolver pour l'utiliser. L'exemple de code suivant montre comment gérer la logique TPUClusterResolver :

def wait_for_tpu_cluster_resolver_ready():
  """Waits for `TPUClusterResolver` to be ready and return it.

  Returns:
    A TPUClusterResolver if there is TPU machine (in TPU_CONFIG). Otherwise,
    return None.
  Raises:
    RuntimeError: if failed to schedule TPU.
  """
  tpu_config_env = os.environ.get('TPU_CONFIG')
  if not tpu_config_env:
    tf.logging.info('Missing TPU_CONFIG, use CPU/GPU for training.')
    return None

  tpu_node = json.loads(tpu_config_env)
  tf.logging.info('Waiting for TPU to be ready: \n%s.', tpu_node)

  num_retries = 40
  for i in range(num_retries):
    try:
      tpu_cluster_resolver = (
          tf.contrib.cluster_resolver.TPUClusterResolver(
              tpu=[tpu_node['tpu_node_name']],
              zone=tpu_node['zone'],
              project=tpu_node['project'],
              job_name='worker'))
      tpu_cluster_resolver_dict = tpu_cluster_resolver.cluster_spec().as_dict()
      if 'worker' in tpu_cluster_resolver_dict:
        tf.logging.info('Found TPU worker: %s', tpu_cluster_resolver_dict)
        return tpu_cluster_resolver
    except Exception as e:
      if i < num_retries - 1:
        tf.logging.info('Still waiting for provisioning of TPU VM instance.')
      else:
        # Preserves the traceback.
        raise RuntimeError('Failed to schedule TPU: {}'.format(e))
    time.sleep(10)

  # Raise error when failed to get TPUClusterResolver after retry.
  raise RuntimeError('Failed to schedule TPU.')

Apprenez-en plus sur l'entraînement distribué avec des conteneurs personnalisés.

Étape suivante

• Documentez-vous sur les modèles d'entraînement sur AI Platform Training.
• Obtenez plus d'informations sur les réglages d'hyperparamètres sur AI Platform Training, en particulier les réglages d'hyperparamètres avec Cloud TPU.
• Explorez des modèles de référence supplémentaires pour Cloud TPU.
• Optimisez vos modèles pour Cloud TPU en suivant les bonnes pratiques relatives à Cloud TPU.
• Consultez la section relative au dépannage et aux questions fréquentes pour Cloud TPU afin d'obtenir de l'aide pour le diagnostic et la résolution des problèmes.