Maîtriser la compilation PyTorch pour des performances IA/ML optimales

Depuis son introduction avec PyTorch 2.0 en mars 2023, torch.compile est rapidement devenu un outil indispensable pour optimiser les performances des charges de travail d’IA et d’apprentissage automatique. PyTorch a initialement gagné une large popularité pour sa conception “Pythonic”, sa facilité d’utilisation et son exécution “eager” (ligne par ligne). L’adoption réussie d’un mode de compilation de graphe juste-à-temps (JIT) n’était donc pas acquise. Pourtant, un peu plus de deux ans plus tard, son importance dans l’amélioration des performances d’exécution ne peut être surestimée. Malgré sa puissance, torch.compile peut encore sembler être un art ésotérique ; bien que ses avantages soient clairs lorsqu’il fonctionne, le diagnostic des problèmes peut être difficile en raison de ses nombreux contrôles d’API et de sa documentation quelque peu décentralisée. Cet article vise à démystifier torch.compile, en expliquant ses mécanismes, en démontrant des stratégies d’application efficaces et en évaluant son impact sur les performances du modèle.

Le mode d’exécution “eager” par défaut de PyTorch, bien que convivial pour le débogage, sacrifie intrinsèquement les opportunités d’optimisation. Chaque ligne Python est traitée indépendamment, ce qui empêche des efficacités telles que la fusion d’opérateurs (combinant plusieurs opérations GPU en un seul noyau plus efficace) et la compilation anticipée (AOT) pour la disposition de la mémoire et l’ordre d’exécution. De plus, les transferts constants entre l’interpréteur Python et le backend CUDA introduisent une surcharge significative. torch.compile résout ces limitations en agissant comme un compilateur JIT. La première fois qu’une fonction compilée est appelée, elle trace le code Python dans une représentation de graphe intermédiaire, souvent appelée FX Graph, en utilisant TorchDynamo. Pour l’entraînement, AOTAutograd capture le passage arrière pour générer un graphe combiné avant et arrière. Ce graphe est ensuite passé à un backend de compilateur, typiquement TorchInductor, qui effectue des optimisations étendues comme la fusion de noyaux et l’exécution dans le désordre. Pour les GPU NVIDIA, TorchInductor exploite le compilateur Triton pour créer des noyaux GPU hautement optimisés et utilise CUDA Graphs lorsque cela est possible pour combiner plusieurs noyaux en séquences efficaces et répétables. Le graphe de calcul résultant, spécifique à la machine, est ensuite mis en cache et réutilisé pour toutes les invocations ultérieures, réduisant considérablement l’implication de l’interpréteur Python et maximisant l’optimisation du graphe.

Bien que torch.compile stimule généralement le débit du modèle, les développeurs rencontrent parfois des scénarios où les performances stagnent ou même se dégradent. Au-delà des goulots d’étranglement externes comme les pipelines d’entrée de données lents, deux “tueurs de compilation” principaux sont souvent responsables : les ruptures de graphe (graph-breaks) et les recompilations.

Les ruptures de graphe se produisent lorsque les bibliothèques de traçage, TorchDynamo ou AOTAutograd, rencontrent des opérations Python qu’elles ne peuvent pas convertir en une opération de graphe. Cela force le compilateur à segmenter le code, en compilant des portions séparément et en rendant le contrôle à l’interpréteur Python entre les segments. Cette fragmentation empêche les optimisations globales comme la fusion de noyaux et peut complètement annuler les avantages de torch.compile. Les coupables courants incluent les instructions print(), la logique conditionnelle complexe et les asserts. Frustrant, les ruptures de graphe retombent souvent silencieusement en exécution “eager”. Pour y remédier, les développeurs peuvent configurer le compilateur pour les signaler, par exemple en définissant TORCH_LOGS="graph_breaks" ou en utilisant fullgraph=True pour forcer l’échec de la compilation en cas de rupture. Les solutions impliquent généralement de remplacer les blocs conditionnels par des alternatives compatibles avec le graphe comme torch.where ou torch.cond, ou d’exécuter conditionnellement les instructions print/assert uniquement lorsque la compilation n’est pas en cours.

Le deuxième piège majeur est la recompilation du graphe. Lors de la compilation initiale, torch.compile fait des hypothèses, appelées “gardes”, sur les entrées, telles que les types de données et les formes des tenseurs. Si ces gardes sont violées lors d’une étape ultérieure – par exemple, si la forme d’un tenseur change –, le graphe actuel est invalidé, déclenchant une recompilation coûteuse. Des recompilations excessives peuvent annuler tous les gains de performance et peuvent même entraîner un retour au mode “eager” après une limite par défaut de huit recompilations. L’identification des recompilations peut être effectuée en définissant TORCH_LOGS="recompiles". Face aux formes dynamiques, plusieurs stratégies existent. Le comportement par défaut (dynamic=None) détecte automatiquement le dynamisme et recompilera de manière chirurgicale, mais cela peut atteindre la limite de recompilation. Marquer explicitement les tenseurs et les axes dynamiques avec torch._dynamo.mark_dynamic est souvent la meilleure approche lorsque les formes dynamiques sont connues à l’avance, car cela informe le compilateur de construire un graphe qui prend en charge le dynamisme sans recompilation. Alternativement, définir dynamic=True indique au compilateur de créer un graphe maximalement dynamique, bien que cela puisse désactiver certaines optimisations statiques comme les CUDA Graphs. Une approche plus contrôlée consiste à compiler un nombre fixe et limité de graphes statiques en complétant les tenseurs dynamiques à quelques longueurs prédéterminées, garantissant que toutes les variations de graphe sont créées pendant une phase d’échauffement.

Le débogage des problèmes de compilation, qui se présentent souvent avec des messages d’erreur longs et cryptiques, peut être décourageant. Les approches vont du “haut vers le bas”, où la compilation est appliquée à l’ensemble du modèle et les problèmes sont traités au fur et à mesure qu’ils surviennent (nécessitant un déchiffrage minutieux des journaux), au “bas vers le haut”, où les composants de bas niveau sont compilés de manière incrémentielle jusqu’à ce qu’une erreur soit identifiée (rendant le repérage plus facile et permettant des avantages d’optimisation partielle). Une combinaison de ces stratégies donne souvent les meilleurs résultats.

Une fois qu’un modèle est compilé avec succès, des gains de performance supplémentaires peuvent être obtenus grâce à diverses options de réglage, bien que celles-ci offrent généralement des améliorations plus modestes par rapport à la compilation initiale. Les modes de compilateur avancés comme “reduce-overhead” et “max-autotune” peuvent optimiser agressivement pour réduire la surcharge et évaluer plusieurs options de noyau, respectivement, bien qu’ils augmentent le temps de compilation. Différents backends de compilateur peuvent être spécifiés, TorchInductor étant le défaut pour les GPU NVIDIA, tandis que d’autres comme ipex pourraient être mieux adaptés aux CPU Intel. Pour les modèles avec des composants statiques et dynamiques distincts, la compilation modulaire – l’application de torch.compile à des sous-modules individuels – permet des paramètres d’optimisation adaptés à chaque partie. Au-delà du modèle lui-même, PyTorch 2.2 a introduit la capacité de compiler l’optimiseur, améliorant encore les performances des charges de travail d’entraînement. Par exemple, la compilation de l’optimiseur dans un modèle de légende d’image de démonstration a augmenté le débit de 5,17 à 5,54 étapes par seconde.

Bien que les temps de compilation initiale et d’échauffement puissent être longs, ils sont généralement négligeables par rapport à la durée de vie globale d’entraînement ou d’inférence d’un modèle. Cependant, pour les modèles extrêmement volumineux où la compilation pourrait prendre des heures, ou pour les serveurs d’inférence où le temps de démarrage a un impact sur l’expérience utilisateur, la réduction de cette durée devient critique. Deux techniques clés sont la mise en cache au moment de la compilation et la compilation régionale. La mise en cache au moment de la compilation implique la sauvegarde des artefacts de graphe compilés dans un stockage persistant (par exemple, Amazon S3) et leur rechargement pour les exécutions ultérieures, contournant ainsi la recompilation à partir de zéro. Dans une démonstration, cela a réduit l’échauffement de compilation de 196 secondes à 56 secondes, soit une accélération de 3,5x. La compilation régionale applique torch.compile aux blocs de calcul répétitifs au sein d’un grand modèle plutôt qu’à l’ensemble de la structure. Cela crée un graphe unique et plus petit qui est réutilisé sur toutes les instances de ce bloc. Pour le modèle de démonstration, cela a réduit l’échauffement de 196 secondes à 80 secondes (une accélération de 2,45x), bien que cela se soit accompagné d’une légère diminution du débit de 7,78 à 7,61 étapes par seconde. Bien que les gains sur un petit modèle de démonstration soient modestes, ces techniques peuvent être essentielles pour les déploiements à grande échelle dans le monde réel.

En fin de compte, à mesure que les modèles d’IA/ML continuent de croître en complexité et en échelle, l’optimisation de leurs performances d’exécution est primordiale. torch.compile se présente comme l’un des outils d’optimisation les plus puissants de PyTorch, capable de fournir des accélérations significatives – jusqu’à 78 % plus rapides dans certains scénarios de graphes statiques et 72 % plus rapides avec des graphes dynamiques dans les exemples présentés. Maîtriser ses nuances, de l’évitement des pièges courants comme les ruptures de graphe et les recompilations au réglage fin des paramètres et à l’exploitation des fonctionnalités avancées, est crucial pour libérer tout son potentiel.