Optimisation des réseaux de neurones : comment redresser un graphe de calcul emmêlé ✛
Contexte
Les réseaux de neurones profonds pour les systèmes de recommandation sont des monstres de complexité computationnelle. Plus le modèle est précis, plus son paysage interne est enchevêtré : des milliers d'opérations s'entrelacent en graphes à la topologie bizarre. Chaque détour superflu, chaque division et recombinaison répétée de tenseurs est comme un sentier non dégagé qui ralentit le voyageur. Au niveau matériel, cela se traduit par une avalanche de petits lancements de noyaux (kernel launch) sur le GPU, qui consomment du temps et de la mémoire. Les ingénieurs tentent de mettre de l'ordre manuellement, mais les modèles changent chaque semaine et les architectures deviennent de plus en plus capricieuses. Un besoin s'est fait sentir d'un « redresseur » automatique de sentiers de graphe, capable de trouver et de nettoyer les fragments de code inefficaces sans intervention humaine.
Méthodes
Le cœur d'Optimus est une collection de plus de trente motifs d'optimisation, chacun décrivant un schéma typique de « déchet » dans le graphe : par exemple, une chaîne insensée « découper un tenseur en morceaux puis le recoller immédiatement ». La recherche de tels schémas est menée par un algorithme glouton qui part d'« ancres » — des opérations clés comme split ou communes aux architectures de recommandation. À partir de l'ancre, l'algorithme parcourt les voisins proches, les compare au motif de référence, et en cas de correspondance, déclenche un remplacement : au lieu d'une construction encombrante, un équivalent compact est inséré. On distingue deux types de réparation : horizontale — lorsque plusieurs opérations identiques indépendantes (par exemple, une douzaine de couches linéaires) fusionnent en une seule, et verticale — lorsque des transitions redondantes sont supprimées, comme les fameux split-then-cat. Toute la procédure est intégrée au compilateur PyTorch 2.x et fonctionne sur n'importe laquelle de ses représentations intermédiaires (Torch-IR, Aten-IR), ce qui la rend universelle pour des modèles hétéroclites.
Résultats
Sur cinq modèles de recommandation industriels de Meta de différents types architecturaux (de type DHEN, de type Wukong, de type Interformer, etc.), Optimus a constamment augmenté le débit — jusqu'à 63 % d'exemples d'entraînement en plus par seconde. La contribution principale provient des optimisations verticales : l'élimination des cycles split-cat parasites a supprimé des centaines de noyaux Triton superflus. Le gain des fusions horizontales dépendait du nombre d'opérations identiques — plus elles sont nombreuses, plus le gain est important. La consommation mémoire de pointe sur les grands modèles avec quantification des activations a chuté de 2 à 6 %, et le temps de compilation à froid a tout simplement dégringolé : dans un test, il est passé de plus de 4000 secondes à environ 3600, soit une économie de plus de 400 secondes. Sur l'exemple illustratif ouvert de TorchBench, le nombre de noyaux (kernel) dans l'entraînement est passé de 12 772 à 24, et dans l'inférence de 6 320 à 6, ce qui a donné une accélération de l'entraînement de 11 fois et de l'inférence de 364 fois. La comparaison avec des analogues (TASO, TENSAT) a montré que ceux-ci n'ont pas du tout réussi à améliorer le graphe pour ce modèle, alors qu'Optimus a réussi grâce à son système de règles flexible.
Implications
Optimus n'est pas un outil jetable, mais un cadre modulaire intégré au code open source de PyTorch. Tout chercheur ou ingénieur peut ajouter un nouveau motif d'optimisation en l'enregistrant via l'interface standard, puis l'appliquer immédiatement à son graphe de réseau de neurones. Cette approche transforme ce qui était autrefois un travail manuel pénible sur le code en un nettoyage méthodique au niveau du compilateur. Pour l'industrie, cela signifie que les GPU coûteux seront utilisés plus efficacement sans réécrire les modèles, et que l'empreinte carbone de l'entraînement des systèmes de recommandation géants sera réduite.
Développements futurs
Actuellement, les règles principales fonctionnent sur Torch-IR et Aten-IR, mais les auteurs prévoient d'étendre la méthode à Inductor-IR — une représentation de plus bas niveau où l'on peut extraire encore plus de performances. Une autre direction est la sélection automatique de l'ordre et de l'ensemble des motifs pour un modèle donné : aujourd'hui, la séquence d'optimisations est choisie manuellement, et l'on souhaiterait confier cela à un algorithme qui analyse la structure du graphe.
Impact
La technologie accélérera directement l'entraînement et le déploiement des systèmes de recommandation — un maillon clé des plateformes Internet modernes. Elle peut également être adaptée à d'autres réseaux de neurones profonds, en particulier ceux avec des graphes de calcul ramifiés : traitement du langage naturel, vision par ordinateur, réseaux de neurones à graphes.
Prochaines étapes
Transférer les acquis au niveau Inductor-IR et mettre en place un planificateur automatique de la séquence de motifs. Il faudra également étendre la bibliothèque de motifs pour couvrir de nouveaux motifs architecturaux.
Problèmes ouverts clés
Le problème de l'exécution efficace des réseaux de neurones sur le matériel moderne se heurte à l'écart entre la complexité des architectures et la bande passante mémoire limitée des GPU. Optimus est un pas vers l'élimination automatique de cet écart au niveau du compilateur, sans exiger du développeur de modifier le code. Cela s'inscrit dans le contexte plus large de la recherche de méthodes évolutives d'optimisation numérique des calculs.
Le regard de la machine du savoir
Cette section s'adresse à l'auteur du travail et à ses collègues. Les travaux voisins ont été trouvés par une recherche sémantique dans notre archive, et non par correspondance de mots. Nous n'évaluons pas : nous montrons le champ autour.
Ce que fait ce travail
Dans cet article, les auteurs présentent Optimus, un framework modulaire pour l'optimisation automatique des graphes de calcul dans PyTorch 2.x, destiné aux modèles de recommandation. Il s'appuie sur une collection de plus de trente motifs de recherche et de remplacement de schémas inefficaces (par exemple, split-then-cat) à l'aide d'un algorithme glouton basé sur des ancres. Sur cinq modèles industriels de Meta, ils ont obtenu une accélération de l'entraînement allant jusqu'à 63 %, une réduction de la consommation mémoire de pointe de 2 à 6 % et une diminution significative du temps de compilation (de plus de 4000 à environ 3600 secondes dans un test).
Où est sa force
La force de ce travail réside dans son intégration pratique avec le compilateur PyTorch et dans son système de règles flexible qui permet de détecter et d'éliminer automatiquement les schémas inefficaces dans des graphes industriels réels, là où des outils comparables (TASO, TENSAT) échouent.
Vers où avancer
- Il pourrait être intéressant de se pencher sur l'expérience de SpecFANN pour accélérer les calculs grâce à des architectures spécialisées : on y utilise un émulateur rapide pour générer des spectres, ce qui fait écho à l'idée de remplacer des fragments lourds par des équivalents légers. Peut-être que les principes de remplacement basés sur des surrogates entraînés pourraient être adaptés pour enrichir l'ensemble de motifs d'Optimus.appui : SpecFANN : Esquisse éclair du spectre stellaire arXiv:2607.17348v1
- À proximité, il y a des travaux sur un réseau de neurones qui ressent la tension dans des pièces de formes variées : ils utilisent des méthodes d'analyse de graphes pour identifier les nœuds critiques. Cela pourrait suggérer comment améliorer l'heuristique de sélection des ancres dans Optimus afin de trouver des optimisations plus globales.appui : Un réseau de neurones qui sent les tensions dans les pièces de toutes formes arXiv:2606.08287
À essayer
- Envisager d'utiliser un surrogate entraîné pour prédire le gain apporté par l'application d'un motif, afin de prioriser les remplacements.
- Tester une analyse de la « tension » dans le graphe (par analogie avec les travaux sur les pièces) pour sélectionner des ancres qui ont plus de chances de déclencher des optimisations en cascade.
- Vérifier s'il est possible de générer automatiquement de nouveaux motifs en généralisant les schémas fréquents dans les graphes réels.
Travaux voisins dans notre archive
- Un réseau de neurones qui sent les tensions dans les pièces de toutes formesproximité 0.62
- SpecFANN : Esquisse éclair du spectre stellaireproximité 0.62
- Comment un jumeau numérique sauve un cerveau en quelques minutesproximité 0.62
- Un réseau de neurones apprend à s'écouler à travers un tamisproximité 0.60
- Étiquettes pour graphes : quand l'approche quantique cède le pas à la simplicitéproximité 0.60
- Les neurones livrent leurs secrets en quelques minutes, pas en heuresproximité 0.59
Le travail le plus proche de l'archive se situe à 0.62 par le sens — le niveau habituel de notre corpus.
Résumé original
tel que rédigé par les auteurs, en anglais arXiv:2607.02945
In large-scale industrial applications, deep learning models that power recommendation and ranking have complex and diverse model architectures. These models are continuously developed and refined by large teams of machine learning engineers, rendering manual optimization infeasible. Consequently, graph-based optimization techniques have become an industry standard for boosting performance, with PyTorch FX transformations leading the charge. These transformations typically rely on a set of human-engineered module-level rewrite rules which are not scalable to diverse model architectures. To address this limitation, we introduce Optimus, a general-purpose model transformation framework built in the PyTorch 2.x (PT2) machine learning compiler. With a concise set of predefined patterns, Optimus applies an efficient greedy search algorithm for pattern matching and replacement, while preserving model semantic. It is designed and implemented as a highly customizable and extensible framework integrated into the PT2 stack. Our evaluation shows that the framework can achieve up to 63% speedup, 6% peak memory reduction, and over 400 second compile time decrease for our industry-scale recommendation models compared to baselines. Optimus is open-sourced together with PyTorch 2.x as a customizable model transformation layer.
Réduire le temps de compilation de 400 secondes, c'est comme si, au lieu de faire du café, de vérifier ses e-mails et d'attendre que l'environnement de développement s'ouvre, vous vous mettiez directement au travail. Et le nombre de lancements de noyaux dans l'inférence du modèle illustratif est passé de 6320 à 6 — un tel compactage du graphe peut être comparé au remplacement d'une foule de coursiers par un seul camion.
Chiffres clés
- augmentation du débit (QPS): jusqu'à 63 %
- réduction de la mémoire de pointe: jusqu'à 6 %
- diminution du temps de compilation: plus de 400 secondes
- noyaux dans l'entraînement TorchBench avant/après: 12 772 -> 24
- noyaux dans l'inférence TorchBench avant/après: 6 320 -> 6