Avec OpenXLA, Google veut booster les performances des modèles ML
Google Cloud est à l’initiative d’un nouveau projet open source intitulé OpenXLA. Le géant du cloud veut réunir les grands noms du matériel et du logiciel pour améliorer les performances des modèles de machine learning, quelle que soit l’architecture matérielle utilisée pour les déployer.
La multiplication des équipements, les différences d’architecture logicielle et la pluralité des configurations matérielles complexifient le développement de l’IA. Surtout, l’intelligence artificielle est très gourmande en puissance de calcul, donc chère et énergivore.
OpenXLA, un héritage de TensorFlow
En ce sens, Google a dévoilé OpenXLA, un projet open source auquel devraient participer d’autres fournisseurs tels qu’AMD, Meta, Nvidia, Intel, ARM, AWS et Apple. Ce sont les principaux contributeurs réunis lors du don de PyTorch par Meta à la fondation Linux.
Avec OpenXLA, il s’agit d’augmenter considérablement les performances des modèles ML développés avec les frameworks TensorFlow, PyTorch et JAX.
En l’occurrence, ces librairies étaient souvent pensées avec une architecture matérielle en tête. Les modèles développés avec ces outils héritent de leurs avantages et limites. Par exemple, GPT-3, le modèle NLG d’OpenAI a été développé pour s’exécuter sur des clusters de GPU Nvidia. Un modèle de ce type sera donc, par défaut, plus complexe à exécuter sur des TPUv4, les ASICS de Google Cloud.
L’initiative OpenXLA vise à créer un compilateur modulaire compatible avec les CPU, GPU, et les ASICS présents sur le marché. Elle débute par le dépaysement d’un sous-projet rattaché à TensorFlow, XLA (Accelerated Linear Algebric). Originellement, ce compilateur d’algèbre linéaire doté d’un langage dédié doit optimiser les modèles TensorFlow « sans en changer le code ».
Avec XLA, les ingénieurs de GCP ont déjà pu multiplier par sept les performances du modèle NLP BERT exécuté sur un système Nvidia regroupant huit GPU V100.
« XLA propose un mode alternatif d’exécution des modèles : il compile le graphe TensorFlow en une séquence de kernels de calcul générés spécifiquement pour le modèle donné », explique la documentation du projet. « Comme ces noyaux sont uniques au modèle, ils peuvent exploiter les informations spécifiques pour son optimisation ».
La méthode utilisée peut être étendue à d’autres frameworks. En effet, la plupart d’entre eux utilisent un modèle graphe pour représenter l’enchaînement des transferts et des transformations de données. Ces graphes sont ensuite traduits en instructions pour les puces. En l’occurrence, les ingénieurs de GCP ont déjà ouvert StableHLO, « une couche de portabilité entre frameworks et compilateurs ». Il s’agit d’une évolution de HLO (High Level Operations), le langage dédié d’XLA. De plus, XLA prend en charge les architectures x86-64, ARM64, et celles développées par Nvidia.
Une tendance à la standardisation de l’IA
« Nous nous engageons en faveur d’écosystèmes ouverts de toutes sortes, et […] nous croyons fermement qu’aucune entreprise ne devrait posséder l’innovation en matière d’IA/ML. »
Sachin GuptaV-P et DG, Infrastructures, Google cloud
L’approche est complémentaire à celle de DeepSpeed, la librairie de Microsoft, et de la technique « mixture of experts ».Ces deux technologies visent plus particulièrement à améliorer les performances des mégamodèles. OpenXLA semble avoir une portée élargie. Pour autant, les collaborations évoquées par Google Cloud concernent davantage les fameux LLM (Large Language Models). Par exemple, GCP entend aider Hugging Face à exécuter les modèles développés avec le framework JAX sur ses TPU, dont BLOOM, un mégamodèle multilingue entraîné sur les infrastructures du CNRS.
Le projet, encore rattaché à TensorFlow, devra d’abord gagner son indépendance. Une gouvernance ouverte sera mise en place au cours de l’année 2023. Pour l’instant, une responsable de programme open source et un ingénieur chez Google Cloud (Thea Lamkin et Mehdi Amini, un alumni de l’Université de Strasbourg) sont les gardiens d’OpenXLA.
L’initiative s’inscrit dans un contexte d’ouverture des « backends », à savoir des équipements et des solutions logicielles qui prennent en charge différents cas d’usage. Outre des notions de standardisation (comme veut le faire l’Open Programmable Infrastructure avec les DPU) et d’interopérabilité, il s’agit d’éviter l’enfermement propriétaire.
« Nous nous engageons en faveur d’écosystèmes ouverts de toutes sortes, et […] nous croyons fermement qu’aucune entreprise ne devrait posséder l’innovation en matière d’IA/ML », écrit Sachin Gupta, Vice-président et directeur général, Infrastructure chez Google Cloud, dans un billet de blog.
