TVM est un framework compilateur open-source conçu pour optimiser et exécuter des modèles de apprentissage automatique sur une large gamme de plateformes matérielles. Développé initialement par des chercheurs du laboratoire Berkeley AI Research et de la communauté de l'Université Carnegie Mellon, TVM répond au défi du déploiement de modèles de apprentissage profond sur divers appareils, des téléphones mobiles aux accélérateurs de centres de données à grande échelle. Il fournit une représentation intermédiaire (IR) unifiée et un ensemble de passes d'optimisation qui transforment des descriptions de modèles de haut niveau en code efficace de bas niveau, permettant des gains de performance significatifs par rapport aux frameworks standards.
Le projet a été annoncé pour la première fois en 2017 et est depuis devenu une pierre angulaire de l'écosystème d'infrastructure d'apprentissage automatique. TVM est hébergé sous la Apache Software Foundation, garantissant une licence open-source permissive et un modèle de développement collaboratif. Son architecture sépare le frontend, qui ingère des modèles provenant de frameworks populaires comme PyTorch et TensorFlow, du backend, qui cible des matériels spécifiques tels que les GPU AMD, les CPU Intel, les processeurs ARM et des accélérateurs spécialisés comme les systèmes AWS Trainium et Cerebras.
Architecture principale
La conception de TVM repose sur une IR au niveau du graphe et une IR au niveau des tenseurs. L'IR du graphe capture la structure globale d'un réseau de neurones, permettant des optimisations de haut niveau comme la fusion d'opérateurs, le pliage de constantes et la planification mémoire. L'IR au niveau des tenseurs, connue sous le nom de script TVM, permet aux développeurs d'écrire des planifications personnalisées qui contrôlent l'ordre des boucles, la vectorisation et la parallélisation pour un matériel spécifique. Cette approche à deux niveaux équilibre flexibilité et performance, permettant à la fois un réglage automatisé et manuel.
Un composant clé est constitué par les modules AutoTVM et Ansor, qui automatisent la recherche de planifications optimales. AutoTVM, introduit en 2018, utilise un modèle de coût et une recherche évolutive pour explorer les configurations de planification. Ansor, ajouté en 2020, améliore cela en générant des planifications à partir de zéro grâce à une approche sans modèle prédéfini, atteignant souvent des performances comparables aux bibliothèques réglées à la main. Ces outils réduisent l'expertise requise pour optimiser des modèles sur de nouveaux matériels.
Support matériel et écosystème
TVM prend en charge un large éventail de backends matériels grâce à ses interfaces de génération de code et d'exécution. Il peut cibler les CPU traditionnels d'Intel et d'AMD, les processeurs mobiles et embarqués d'ARM et de Qualcomm, ainsi que les GPU de NVIDIA (bien que non listé dans les slugs fournis, c'est un backend courant) et d'AMD. Il prend également en charge les accélérateurs émergents de sociétés comme Groq, SambaNova et Graphcore, ainsi que les puces cloud spécifiques telles que AWS Trainium et les TPU Google Cloud.
L'environnement d'exécution, appelé TVM Runtime, est léger et peut être intégré dans des applications mobiles ou déployé dans des environnements serveur. Il prend en charge plusieurs langages de programmation, notamment Python, C++, Java et Rust, le rendant accessible à une large base de développeurs. Le projet s'intègre également à la communauté Apache TVM, qui inclut des contributions d'Amazon Web Services, d'Alibaba Cloud et de Samsung Electronics.
Techniques d'optimisation des performances
TVM emploie plusieurs techniques d'optimisation avancées pour améliorer les performances d'inférence et d'entraînement. La fusion d'opérateurs combine plusieurs opérations en un seul noyau, réduisant l'utilisation de la bande passante mémoire et les frais généraux de lancement. Par exemple, une convolution suivie d'une normalisation par lots et d'une activation ReLU peut être fusionnée en un seul noyau. TVM effectue également des transformations de mise en page automatiques, choisissant des formats de données optimaux comme NHWC ou NCHW en fonction du matériel cible.
Une autre caractéristique importante est l'utilisation de modèles de coût basés sur l'apprentissage automatique pour guider l'optimisation. Ces modèles prédisent le temps d'exécution des planifications candidates, permettant à la recherche de se concentrer sur les configurations prometteuses. TVM prend également en charge la quantification et l'élagage, permettant le déploiement de modèles sur des appareils aux ressources limitées. Ces techniques ont démontré des accélérations de 2 à 10 fois par rapport aux frameworks standards comme TensorFlow ou PyTorch sur diverses charges de travail.
Applications et impact
TVM a été adopté dans des environnements de production par de grandes entreprises technologiques. Amazon Web Services utilise TVM pour optimiser des modèles pour ses puces AWS Trainium et Inferentia, offrant des services d'inférence haute performance. Alibaba Cloud intègre TVM dans sa plateforme d'apprentissage automatique, tandis que Samsung Electronics l'utilise pour l'IA sur appareil dans les smartphones et les wearables. Le framework est également utilisé dans la recherche, permettant des expériences avec de nouvelles architectures et de nouveaux matériels.
Dans le domaine des grands modèles de langage, TVM a été étendu pour gérer efficacement les modèles basés sur des transformeurs. Des techniques comme l'optimisation du cache KV et des noyaux d'attention fusionnés ont été implémentées, réduisant la latence et l'utilisation mémoire pour des modèles tels que la série GPT d'OpenAI et Claude d'Anthropic. Cela fait de TVM un outil pertinent pour les applications d'IA générative, où le coût d'inférence est une préoccupation majeure.
Communauté et développement
TVM est développé ouvertement sur GitHub, avec des contributions d'une communauté mondiale de chercheurs et d'ingénieurs. Le projet organise des réunions bihebdomadaires régulières et une conférence annuelle, TVMCon, qui rassemble utilisateurs et développeurs. Le modèle de gouvernance comprend un comité de gestion de projet (PMC) élu parmi les contributeurs actifs, garantissant une durabilité à long terme. En 2024, le projet compte plus de 1 000 contributeurs et a été cité dans de nombreux articles académiques.
Le framework continue d'évoluer, avec des travaux en cours sur l'amélioration du temps de compilation, l'expansion du support matériel et l'intégration avec des frameworks d'IA comme Triton d'OpenAI. La nature open-source et la communauté active de TVM en font une pièce critique de l'infrastructure de l'apprentissage automatique, comblant le fossé entre le développement de modèles et leur déploiement.
Orientations futures
À l'avenir, TVM vise à améliorer le support des modèles transformeurs et des LLM, en se concentrant sur l'inférence et l'entraînement distribués. Le projet explore également la différenciation automatique et la compilation pour les charges de travail de apprentissage profond, ce qui pourrait simplifier l'optimisation des pipelines d'entraînement. Avec l'essor de l'IA en périphérie et du matériel spécialisé, le rôle de TVM en tant que compilateur universel est susceptible de s'étendre, en faisant un outil essentiel pour les praticiens de l'IA.
Malgré ses forces, TVM fait face à des défis tels que la complexité de sa base de code et la nécessité d'une adaptation continue aux nouveaux matériels. Cependant, ses principes de conception et le soutien de sa communauté le positionnent bien pour relever ces défis, garantissant sa pertinence dans le domaine en évolution rapide de l'infrastructure d'apprentissage automatique.