Introduction aux unités de traitement tensoriel de Google

Écrit par Vincent Piumi, le 11 juillet 2019

 

Elles ont été annoncées en mai 2016 lors de la conférence Google I/O. Les Tensor Processing Units (TPU) de Google sont une nouvelle architecture prometteuse orientée vers l'IA. Elles ont été conçues à l'origine pour aider à réaliser des percées commerciales et de recherche dans le domaine de l'apprentissage automatique. Elles sont accessibles au public via la plateforme cloud de Google. Elles sont également compatibles avec Tensorflow et leur propre compilateur XLA ("Accelerated Linear Algebra").

Il est de notoriété publique que les modèles d'apprentissage automatique effectuent des opérations lourdes sur les matrices. Cependant, si l'on considère les choses sous un autre angle, il apparaît que, sur la base de leurs composants, les Tensor Processing Units pourraient s'avérer utiles dans d'autres domaines.

Chez Aldwin, nous nous concentrons principalement sur les problèmes liés au calcul intensif qui traitent de structures de données similaires, telles que la simulation physique.

L'état actuel des architectures HPC

Tout d'abord, le calcul intensif peut être considéré comme un ensemble de concepts ou de méthodes visant à réduire le temps d'exécution ou à améliorer la consommation d'énergie des algorithmes de calcul. Les améliorations informatiques peuvent être liées au logiciel ou au matériel : le calcul intensif vise les deux.

De nos jours, le matériel est un élément clé du calcul intensif. Ses progrès et ses capacités peuvent constituer un facteur limitant pour le développement de logiciels. L'utilisateur n'a que la possibilité de choisir l'implémentation logicielle qui s'adaptera le mieux à son architecture. D'un ordinateur à cœur unique à un superordinateur complexe utilisant des milliers d'unités centrales de traitement (CPU), l'éventail des performances peut être large.

CPU contre GPU

Les unités centrales de traitement sont les circuits électroniques de base qui effectuent les opérations arithmétiques, logiques et de contrôle spécifié par un programme informatique. De nos jours, certains ordinateurs utilisent un processeur multicœur, c'est-à-dire une puce unique contenant deux unités centrales de traitement ou plus. OpenMP est un exemple d'API qui prend en charge la programmation multiprocessus.

Les unités de traitement graphique (GPU) sont un autre type de circuits électroniques. Ils sont conçus pour être performants dans la manipulation des graphiques informatiques et le traitement des images. Contrairement aux processeurs multicœurs modernes, les GPU sont constitués de milliers d'unités relativement petites qui fonctionnent en parallèle. Ils se prêtent bien à l'informatique massivement parallèle. L'apprentissage profond a contribué à la démocratisation des GPU à usage général pour travailler avec des opérations graphiques complexes.

Le meilleur indicateur possible de l'utilisation des CPU et des GPU pour le calcul intensif serait le projet de classement Top500 [1].

Produit depuis juin 1993, il classe les 500 systèmes de supercalculateurs les plus puissants au monde, et en juin 2018, Summit [2] est le plus puissant. Sans surprise, il repose entièrement sur une architecture hybride de CPU (9 216 IBM POWER9) et de GPU (27 648 Nvidia Tesla). Comme il est désormais facile d'utiliser des langages tels que CUDA (développé par Nvidia) pour travailler directement avec les deux architectures, c'est la combinaison la plus courante que l'on puisse trouver. Par exemple, GROMACS, utilisé en dynamique moléculaire, peut fonctionner sur les deux architectures pour de meilleures performances.

Les Tensor Processing Units de Google, une réponse aux limitations matérielles

Les paragraphes précédents soulèvent la question des limites du matériel, qui doit faire face à des opérations, des ensembles de données et des architectures plus complexes. Sans mentionner la stagnation de la loi de Moore ou l'énergie et la puissance comme principaux facteurs limitants.

On peut envisager deux solutions à ce problème dans un avenir proche : opter pour des supercalculateurs encore plus grands, comme le Sunway TaihuLight et ses 10 649 600 cœurs de processeur [3], ou opter pour des architectures spécifiques. C'est cette dernière option qui a été choisie par Google en 2015, lorsqu'il a lancé les Tensor Processing Units.

L'article "In-Datacenter Performance Analysis of a Tensor Processing Unit" (Norman Jouppi et al.), donne la meilleure définition d'une Tensor Processing Unit. Google a été confronté à des limitations matérielles en 2013. À cette époque, une projection a montré que l'utilisation de leurs réseaux neuronaux profonds pourrait doubler la demande de calcul de leurs propres centres de données. Ils ont décidé de concevoir un ASIC personnalisé [4] pour le prendre en charge, et en 15 mois, les TPU ont été déployées.

Une TPU est conçue pour être un coprocesseur sur le bus d'E/S PCIe. Son esprit est plus proche de celui d'un coprocesseur à virgule flottante que de celui d'un GPU. Son principal composant est une unité de multiplication matricielle, contenant 256×256 multiplicateurs-accumulateurs, qui peuvent effectuer des multiplications et des additions de 8 bits sur des nombres entiers signés ou non signés. Ce composant met en évidence la possibilité d'effectuer des opérations à forte intensité matricielle.

Dans l'article "Quantifying the performance of the Tensor Processing Units, our first machine learning chip", Google donne des chiffres de production et de consommation d'énergie liés à ses propres charges de travail d'IA de production. Les TPU seraient 15 à 30 fois plus rapides que les GPU et les CPU (pour les modèles d'inférence des réseaux neuronaux). En outre, elles améliorent de 30 à 80 fois le nombre de téra-opérations par watt d'énergie consommée.

Explorer de nouvelles possibilités

Ces deux derniers paragraphes rappellent l'objectif fondamental du calcul intensif. Chez Aneo, nous pensons que le HPC et l'IA partagent une base commune : tirer parti des architectures informatiques avancées pour faire des découvertes révolutionnaires. Par conséquent, nous allons de l'avant pour combler le fossé entre ces domaines en explorant l'utilisation de TPU pour exécuter une application non-ML : SeWaS.

Aneo a développé un simulateur d'ondes sismiques (SeWaS) suivant un modèle entièrement basé sur les tâches [5]. Comme beaucoup d'autres applications de simulation, il s'appuie sur des grilles de calcul, représentées sous forme de tableaux multidimensionnels.

De plus, il utilise un modèle assez similaire à Tensorflow [6], qui est la principale API utilisée par Google pour écrire du code sur les Tensor Processing Units. En effet, SeWaS utilise les nœuds comme des opérations et les arêtes comme des flux de données. En utilisant l'API C++ de Tensorflow, nous réécrivons le code original de SeWaS et essayons de le porter sur les TPUs ; en bénéficiant de la Matrix Multiply Unit et des spécifications de cette architecture.

[1] https://www.top500.org/ – Top500
[2] https://www.ibm.com/thought-leadership/summit-supercomputer/ – Summit supercomputer
[3]https://www.top500.org/ – Sunway TaihuLight supercomputer
[4]https://fr.wikipedia.org/wiki/Application-specific_integrated_circuit – ASIC
[5]https://github.com/aneoconsulting/SeWaS – SeWaS
[6]https://github.com/tensorflow/tensorflow – Tensorflow