Présentation de Cloud TPU

Les TPU (Tensor Processing Units) sont des circuits intégrés spécifiques aux applications (Application-Specific Integrated Circuit ou ASIC) et développés spécifiquement par Google qui permettent d'accélérer les charges de travail de machine learning. Pour en savoir plus sur le matériel TPU, consultez la page Architecture TPU. Cloud TPU est un service Web qui rend les TPU disponibles en tant que ressources de calcul évolutives sur Google Cloud.

Les TPU permettent d'entraîner vos modèles plus efficacement à l'aide du matériel conçu pour effectuer des opérations matricielles importantes souvent présentes dans les algorithmes de machine learning. Les TPU disposent d'une mémoire à haut débit (HBM) sur puce, ce qui vous permet d'utiliser des modèles et des tailles de lots plus importants. Les TPU peuvent être connectés dans des groupes appelés "slices" qui étendent vos charges de travail avec peu ou pas de modifications de code.

Le code qui s'exécute sur les TPU doit être compilé par le compilateur d'algèbre linéaire accélérée (XLA). XLA est un compilateur juste-à-temps qui prend le graphique émis par une application de framework ML et compile les composants d'algèbre linéaire, de perte et de gradient du graphique en code machine TPU. Le reste du programme s'exécute sur la machine hôte TPU. Le compilateur XLA fait partie de l'image de VM TPU exécutée sur une machine hôte TPU.

Quand utiliser les TPU ?

Les ressources Cloud TPU sont optimisées pour des charges de travail spécifiques. Dans certains cas, il est préférable d'utiliser des GPU ou des processeurs sur des instances Compute Engine afin d'exécuter des charges de travail de machine learning. En règle générale, vous pouvez choisir le matériel le mieux adapté à votre charge de travail en vous appuyant sur les descriptions suivantes :

Processeurs

  • Prototypage rapide nécessitant une flexibilité maximale
  • Modèles simples rapides à entraîner
  • Petits modèles comportant des lots efficaces de taille moindre
  • Modèles contenant de nombreuses opérations TensorFlow personnalisées écrites en C++
  • Modèles limités par la capacité d'E/S ou par la bande passante réseau du système hôte

GPU

  • Modèles avec un grand nombre d'opérations PyTorch/JAX personnalisées devant s'exécuter au moins partiellement sur des processeurs
  • Modèles avec des opérations TensorFlow non disponibles sur Cloud TPU (consultez la liste des opérations TensorFlow disponibles)
  • Modèles moyens à volumineux comportant des lots efficaces de taille plus importante

TPU

  • Modèles dominés par les calculs matriciels
  • Modèles sans opérations PyTorch/JAX personnalisées dans la boucle d'entraînement principale
  • Modèles dont l'entraînement prend des semaines ou des mois
  • Modèles volumineux avec lots efficaces de grande taille
  • Modèles avec des embeddings ultra-volumineux courants dans les charges de travail de classement et de recommandation avancées

Les ressources Cloud TPU ne sont pas adaptées aux charges de travail suivantes :

  • Les programmes algébriques linéaires qui nécessitent des ramifications fréquentes ou qui contiennent de nombreuses opérations algébriques au niveau des éléments
  • Les charges de travail qui nécessitent une arithmétique de haute précision
  • Les charges de travail de réseau de neurones qui contiennent des opérations personnalisées dans la boucle d'entraînement principale

TPU dans Google Cloud

Vous pouvez utiliser des TPU via les VM Cloud TPU, Google Kubernetes Engine et Vertex AI. Le tableau suivant liste les ressources pour chaque service Google Cloud.

ServiceGoogle Cloud Ressources
Cloud TPU Premiers pas avec les VM Cloud TPU
Google Kubernetes Engine

À propos des TPU dans GKE

Exécuter Ray sur GKE avec des TPU
Vertex AI

Entraînement sur Vertex AI avec des TPU

Utiliser des TPU pour la prédiction en ligne sur Vertex AI

Bonnes pratiques en matière de développement d'un modèle

Un programme dont le calcul est dominé par des opérations non matricielles (telles que des additions, des redimensionnements ou des concaténations) n'atteindra sans doute pas une utilisation élevée. Les instructions ci-dessous vous aideront à choisir et à concevoir des modèles adaptés à Cloud TPU.

Disposition

Le compilateur XLA réalise des transformations de code (telles que le tuilage d'un produit matriciel en plusieurs blocs) afin d'effectuer efficacement des calculs sur l'unité matricielle (MXU). Pour obtenir un tuilage efficace, le compilateur utilise la structure du matériel MXU, un tableau systolique de dimension 128 x 128, ainsi que la conception du sous-système de mémoire des TPU, qui préfère des dimensions correspondant à un multiple de 8. Par conséquent, certaines dispositions sont plus propices au tuilage, tandis que d'autres nécessitent des redimensionnements avant de pouvoir être tuilées. Sur une ressource Cloud TPU, les opérations de redimensionnement sont souvent liées à la mémoire.

Formes

Le compilateur XLA compile un graphe de ML de façon opportune pour le premier lot. Si des lots ultérieurs possèdent des formes différentes, le modèle ne fonctionne pas. (et recompiler le graphe à chaque modification de forme prend trop de temps). C'est pourquoi les modèles comportant des Tensors dotés de formes dynamiques ne sont pas adaptés aux TPU.

Remplissage

Un programme Cloud TPU hautes performances est un programme dans lequel les calculs denses peuvent être tuilés en blocs de dimension 128 x 128. Lorsqu'un calcul matriciel ne peut pas occuper l'ensemble d'une unité matricielle, le compilateur remplit les Tensors avec des zéros. Cette méthode présente toutefois deux inconvénients :

  • Les Tensors remplis de zéros sous-utilisent le cœur du TPU.
  • Le remplissage augmente la quantité de mémoire sur puce requise pour un Tensor. Cette technique peut également entraîner une erreur de mémoire insuffisante dans certains cas.

Bien que le remplissage soit automatiquement effectué par le compilateur XLA en cas de besoin, il est possible de définir la proportion de remplissage à l'aide de l'outil op_profile. Vous pouvez éviter le remplissage en choisissant des dimensions de Tensors adaptées aux TPU.

Dimensions

Le choix des dimensions appropriées pour les Tensors aide considérablement à tirer les meilleures performances du matériel TPU, en particulier de l'unité matricielle. Le compilateur XLA tente d'utiliser la taille des lots ou une dimension de fonctionnalité pour optimiser l'utilisation de l'unité matricielle. L'une de ces valeurs doit donc correspondre à un multiple de 128. Dans le cas contraire, le compilateur remplit l'un de ces deux éléments afin d'atteindre 128. Idéalement, la taille des lots et les dimensions des caractéristiques doivent correspondre à des multiples de 8, ce qui permet d'exploiter le sous-système de mémoire au maximum de ses capacités.

Premiers pas avec Cloud TPU

Demander de l'aide

Contactez l'assistance Cloud TPU. Si vous avez un projet Google Cloud actif, soyez prêt à fournir les informations suivantes:

  • Votre Google Cloud ID de projet
  • Nom de votre TPU, le cas échéant
  • Autres informations que vous souhaitez fournir

Étape suivante

Vous souhaitez en savoir plus sur Cloud TPU ? Les ressources suivantes peuvent vous y aider: