Débroussailler les graphes : comment Meta a accéléré l'apprentissage des réseaux de neurones de 63 % ✛

‹›

Les systèmes de recommandation qui décident quel message vous verrez ensuite fonctionnent sur des réseaux de neurones profonds composés de milliers d'opérations — des graphes de calcul complexes (des schémas d'opérations sur les données). Ces systèmes sont des systèmes dynamiques vivants, en constante évolution : la moindre modification de l'architecture se propage en cascade à travers tout le graphe. À mesure que la précision des modèles augmente, les sentiers s'emmêlent : des embranchements superflus comme « découper un tenseur (tableau multidimensionnel de nombres) en morceaux puis le recoller aussitôt » créent de véritables embouteillages sur les processeurs graphiques. Chaque méandre déclenche un noyau de calcul distinct, et une journée d'apprentissage peut accumuler des billions de lancements inutiles. Les ingénieurs tentent de déblayer ces sentiers à la main, mais les modèles changent chaque semaine, et remettre de l'ordre dans ce chaos entropique sans l'aide d'un automate est devenu impossible.

Optimus est un module du compilateur PyTorch 2.x qui agit comme un cartographe avide. Il est équipé d'une bibliothèque de plus de trente motifs d'optimisation — des schémas de gaspillage typiques comme la chaîne split-then-cat. L'algorithme part de points d'ancrage — des opérations critiques autour desquelles l'enchevêtrement se forme — et parcourt les nœuds voisins, en vérifiant chaque étape par rapport à un motif de référence. Dès qu'il trouve une correspondance, il découpe la construction encombrante et insère un équivalent compact. Le nettoyage horizontal fusionne une dizaine de couches identiques en une seule ; le nettoyage vertical élimine les divisions et les collages inutiles. Toute la procédure ne change pas le résultat final, mais simplifie radicalement le graphe interne. C'est précisément ce type d'optimisation qui s'applique aux graphes de n'importe quelle topologie, car les règles s'appliquent sur la représentation intermédiaire universelle Torch-IR.

L'effet sur les systèmes industriels de Meta s'est avéré tangible. Cinq modèles de apprentissage automatique de recommandation, d'architectures différentes, ont montré une augmentation du débit allant jusqu'à 63 % — plus d'exemples d'apprentissage par seconde sans changer de matériel. Les nettoyages verticaux ont éliminé des centaines de noyaux Triton parasites, tandis que les nettoyages horizontaux ont fusionné une multitude d'opérations identiques en une seule ; plus il y a de clones, plus le gain est visible. La consommation de mémoire de pointe a diminué de 2 à 6 %, ce qui est crucial pour les modèles avec quantification des activations, où chaque mégaoctet compte. Dans un test illustratif sur TorchBench, le nombre de noyaux lors de l'apprentissage est passé de 12 772 à 24, et lors de l'inférence de 6 320 à 6, donnant respectivement une accélération de 11 et 364 fois. Ce triomphe numérique a été rendu possible parce qu'Optimus opère au niveau du compilateur, voyant le graphe dans son ensemble.

Ce framework transforme un réseau de sentiers chaotiques en autoroute pour les données, et surtout, tous ses outils sont ouverts à tous. Le code est intégré à PyTorch, et tout développeur peut non seulement utiliser les motifs prêts à l'emploi, mais aussi ajouter les siens, apprenant au système de nouvelles astuces de redressement. À terme, les auteurs prévoient de déplacer l'optimisation au niveau encore plus bas d'Inductor-IR et d'apprendre à l'algorithme à choisir lui-même la meilleure séquence de règles — sans intervention humaine. Ce n'est pas seulement une accélération de l'apprentissage ; c'est un pas vers des réseaux de neurones auto-optimisés, qui s'adaptent eux-mêmes au matériel, économisant électricité et temps. Les sentiers enchevêtrés des graphes ne dévoreront plus les ressources inutilement — le pilote automatique les redressera.