OpenXLA est un écosystème de compilateurs ML open source co-développé
Récemment, les plus grandes entreprises spécialisées dans le développement dans le domaine de l'apprentissage automatique ont présenté le projet OpenXLA , visant au développement conjoint d'outils pour compiler et optimiser les modèles pour les systèmes d'apprentissage automatique.
Le projet a pris en charge le développement d'outils permettant d'unifier la compilation de modèles préparés dans les frameworks TensorFlow, PyTorch et JAX pour un entraînement et une exécution efficaces sur différents GPU et accélérateurs spécialisés. Des entreprises telles que Google, NVIDIA, AMD, Intel, Meta, Apple, Arm, Alibaba et Amazon ont rejoint le travail conjoint du projet.
Le projet OpenXLA fournit un compilateur ML de pointe qui peut s'adapter à la complexité de l'infrastructure ML. Ses piliers fondamentaux sont la performance, l'évolutivité, la portabilité, la flexibilité et l'extensibilité pour les utilisateurs. Avec OpenXLA, nous aspirons à libérer le véritable potentiel de l'IA en accélérant son développement et sa livraison.
OpenXLA permet aux développeurs de compiler et d'optimiser des modèles à partir de tous les principaux frameworks ML pour une formation et une maintenance efficaces sur une grande variété de matériels. Les développeurs utilisant OpenXLA verront des améliorations significatives en termes de temps de formation, de performances, de latence de service et, finalement, de délais de commercialisation et de coûts de calcul.
On espère qu'en unissant les efforts des principales équipes de recherche et des représentants de la communauté, il sera possible de stimuler le développement des systèmes d'apprentissage automatique et de résoudre le problème de la fragmentation des infrastructures pour les différents cadres et équipes.
OpenXLA permet une prise en charge efficace de divers matériels, quel que soit le framework sur lequel repose le modèle d'apprentissage automatique. OpenXLA devrait réduire le temps d'entraînement des modèles, améliorer leurs performances, diminuer la latence et la charge de calcul, et accélérer leur mise sur le marché.
OpenXLA se compose de trois composants principaux , dont le code est distribué sous la licence Apache 2.0 :
- XLA (algèbre linéaire accélérée) est un compilateur qui vous permet d'optimiser les modèles d'apprentissage automatique pour une exécution haute performance sur différentes plates-formes matérielles, y compris les GPU, les CPU et les accélérateurs spécialisés de divers fabricants.
- StableHLO est une spécification de base et l'implémentation d'un ensemble d'opérations de haut niveau (HLO) à utiliser dans les modèles de système d'apprentissage automatique. Il agit comme une couche entre les frameworks d'apprentissage automatique et les compilateurs qui transforment le modèle pour qu'il s'exécute sur du matériel spécifique. Les couches sont préparées pour générer des modèles au format StableHLO pour les frameworks PyTorch, TensorFlow et JAX. La suite MHLO est utilisée comme base pour StableHLO, qui est étendu avec la prise en charge de la sérialisation et du contrôle de version.
- IREE (Intermediate Representation Execution Environment) est un compilateur et un runtime qui convertit les modèles d'apprentissage automatique en une représentation intermédiaire universelle basée sur le format MLIR (Intermediate Multi-Level Representation) du projet LLVM. Parmi les fonctionnalités, la possibilité de précompilation (à l'avance), la prise en charge du contrôle de flux, la possibilité d'utiliser des éléments dynamiques dans les modèles, l'optimisation pour différents CPU et GPU et une faible surcharge sont mises en évidence.
Parmi les principaux avantages d'OpenXLA, on peut citer l' obtention de performances optimales sans avoir à écrire de code spécifique au périphérique, ainsi que des optimisations prêtes à l'emploi, notamment la simplification des expressions algébriques, l'allocation efficace de mémoire, la planification de l'exécution, en tenant compte de la réduction de la consommation maximale de mémoire et de la surcharge.
Un autre avantage réside dans la simplification de la mise à l'échelle et de la parallélisation des calculs. Il suffit au développeur d'ajouter des annotations pour un sous-ensemble de tenseurs critiques, et le compilateur peut alors générer automatiquement le code nécessaire au calcul parallèle.
Il est également souligné que la portabilité est assurée grâce à la prise en charge de plusieurs plateformes matérielles, telles que les GPU AMD et NVIDIA, les processeurs x86 et ARM, les accélérateurs ML Google TPU, les IPU AWS Trainium Inferentia, Graphcore et le moteur Cerebras Wafer-Scale.
Prise en charge de la connexion d'extensions avec l'implémentation de fonctions supplémentaires, telles que la prise en charge de l'écriture de primitives d'apprentissage automatique profond utilisant CUDA, HIP, SYCL, Triton et d'autres langages pour le calcul parallèle, ainsi que la possibilité d'ajuster manuellement les goulots d'étranglement dans les modèles.
Enfin, si vous souhaitez en savoir plus, vous trouverez les détails en suivant ce lien.