Un Aperçu Complet Des Outils FX Utilisés Par Meta : Optimisation Des Modèles PyTorch Avec La Transformation Graphique
Le mode graphique dans PyTorch est plus performant. Cet article présente Torch.FX, un outil puissant capable de capturer et d'optimiser le graphique des programmes PyTorch.
1. Introduction
PyTorch prend en charge deux modes d'exécution : le mode impatient et le mode graphique.
En mode impatient, les opérateurs du modèle sont exécutés immédiatement lorsqu'ils sont lus. Il est facile à utiliser et plus convivial pour les praticiens de l'apprentissage automatique, il est donc défini comme mode d'exécution par défaut.
En mode graphique, les opérateurs sont d'abord synthétisés dans un graphique, puis compilés et exécutés dans leur ensemble. Il a des performances supérieures et est donc largement utilisé dans la production réelle.
Plus précisément, le mode graphique prend en charge la fusion d’opérateurs. En fusionnant deux opérateurs, la surcharge totale des lectures de mémoire et des lancements de noyau peut être réduite ou localisée.
La fusion peut être horizontale :Prend une seule opération (telle que BatchNorm) appliquée à plusieurs opérandes et les fusionne dans un seul tableau.
La fusion peut également être verticale :Combinez un noyau avec un autre noyau qui nécessite la sortie du premier noyau (comme ReLU suivi d'une convolution).
Torch.FX (abrégé en FX) est une boîte à outils accessible au public qui prend en charge l'exécution en mode graphique dans le cadre du package PyTorch. Ça peut:
1. Obtenez le graphique du programme PyTorch
2. Permettre aux développeurs d'écrire des transformations sur le graphique obtenu
Meta a déjà utilisé FX pour optimiser le débit de formation des modèles de production. Cet article présentera l'optimisation basée sur FX développée par Meta pour montrer comment optimiser les performances des modèles déployés par PyTorch à l'aide de la transformation graphique.
II. Arrière-plan
Les tables d’intégration sont largement utilisées dans les systèmes de recommandation.Cette section présentera les connaissances de base sur les effets spéciaux et la table d'intégration.
2.1. Effets spéciaux
La figure 1 est un exemple simple montrant comment convertir un programme PyTorch à l’aide de FX.Elle se déroule en trois étapes :
- Obtenir le graphique du programme
- Modifier le graphique (dans ce cas, nous utilisons GELU au lieu de RELU)
- Générer un nouveau programme à partir du graphique modifié
L'API FX fournit de nombreuses autres fonctions pour inspecter et transformer les graphiques du programme PyTorch.
2.2. Tableau d'intégration
Dans le système de recommandation,Les fonctionnalités éparses (par exemple, l'ID utilisateur, l'ID d'histoire) sont représentées par une table d'intégration.
La table d'intégration E est une matrice HxD, où H est la taille du hachage et D est la dimension du vecteur d'intégration. Chaque ligne de E est un vecteur de nombres à virgule flottante.
La fonction du hachage de caractéristiques est de mapper une caractéristique clairsemée à une liste d'index de E, telle que [S1, S2, …, Sk], où 0 ≤ Si
Pour utiliser pleinement le GPU, les fonctionnalités éparses sont généralement traitées par lots.Chaque entité du lot possède sa propre liste d’index. Si un lot contient B entités, il peut être simplement compris comme une représentation contenant B listes d'index.
Une représentation plus rigoureuse consisterait à fusionner les listes d’index B en une seule liste d’index et à ajouter une liste de longueurs d’index (une pour chaque entité du lot).
Par exemple, si un lot contient 3 entités, sa liste d'index est la suivante :
- Entité 1 : indices = [10, 20]
- Entité 2 : indices = [5, 9, 77, 81]
- Entité 3 : indices = [15, 20, 45]
L'indice et la longueur de la taille totale du lot seront alors :
- Indices = [10, 20, 5, 9, 77, 81, 15, 20, 45]
- Longueurs = [2, 4, 3]
La sortie de la requête de table d'intégration pour l'ensemble du lot est une matrice BxD.
3. 3 Transformations FX
PyTorch a mis à jour trois transformations FX pour accélérer l'accès à la table d'intégration, qui seront introduites une par une dans cette section.
Vous trouverez ci-dessous 3.1 sur la transformation de la combinaison de plusieurs petits tenseurs d’entrée en un seul grand tenseur ; 3.2 sur la transformation de la fusion de plusieurs chaînes de calcul parallèles en une seule chaîne de calcul ; et 3.3 sur la transformation de la communication et de l'informatique qui se chevauchent.
3.1 Combinaison de caractéristiques éparses en entrée
Chaque fonction éparse d'entrée dans un lot peut être représentée sous forme de deux listes : une liste d'index et une liste de longueur B, où B représente la taille du lot.
Dans PyTorch, les deux listes peuvent exister sous forme de tenseurs.Lorsqu'un modèle PyTorch s'exécute sur un GPU, la table d'intégration est généralement stockée dans la mémoire du GPU (qui est plus proche du GPU et a une bande passante de lecture et d'écriture plus élevée que la mémoire du CPU).
Lorsque les fonctionnalités éparses d'entrée doivent être utilisées, les deux tenseurs doivent d'abord être copiés du CPU vers le GPU. Cependant, chaque copie de mémoire de l'hôte vers l'appareil nécessite le lancement d'un noyau, ce qui prend plus de temps que le transfert de données réel.
Si un modèle utilise de nombreuses fonctionnalités éparses en entrée, cette copie peut devenir un goulot d'étranglement en termes de performances (par exemple, 1 000 fonctionnalités éparses en entrée nécessiteront la copie de 2 000 tenseurs de l'hôte vers le périphérique).
Une optimisation pour réduire le nombre de mémoires hôte-appareil consiste à combiner plusieurs fonctionnalités éparses d'entrée avant de les envoyer à l'appareil.
Par exemple, étant donné les trois caractéristiques d’entrée suivantes :
- Caractéristique_A : indices = [106, 211, 7], longueurs = [2, 1]
- Caractéristique_B : indices = [52, 498, 616, 870, 1013], longueurs = [3, 2]
- Feature_C : indices = [2011, 19, 351, 790], longueurs = [1, 3]
La forme combinée est :
Caractéristiques_A_B_C : indices = [106, 211, 7, 52, 498, 616, 870, 1013, 2011, 19, 351, 790], longueurs = [2, 1, 3, 2, 1, 3]
Ainsi, au lieu de copier 3×2=6 tenseurs de l’hôte vers l’appareil, seuls 2 tenseurs doivent être copiés.
La figure 3(b) illustre la mise en œuvre de cette optimisation, qui se compose de deux éléments :
- Côté CPU :Le pipeline d'entrée est modifié pour combiner les indices de toutes les fonctionnalités clairsemées dans un tenseur et toutes les longueurs dans un autre tenseur. Ces deux tenseurs sont ensuite copiés sur le GPU.
- Côté GPU :À l'aide de FX, insérez un opérateur Permute_and_Split dans le graphique du modèle pour récupérer les indices de caractéristiques individuels et les tenseurs de longueur à partir des tenseurs fusionnés et les envoyer aux nœuds correspondants en aval.
3.2 Fusion horizontale des chaînes de calcul à partir de l'accès à la table d'intégration
Dans un modèle de production, il est courant d’avoir 10 tables d’intégration par GPU. Pour des raisons de performances,Les requêtes sur ces tables sont regroupées de manière à ce que leurs sorties soient concaténées en un seul grand tenseur.(Voir la partie rouge dans la figure 4(a)).
Pour calculer une sortie de fonctionnalité unique,Utilisez l'opérateur Split pour diviser un grand tenseur en N petits tenseurs(où N est le nombre de fonctionnalités) puis applique le calcul souhaité à chaque tenseur.
Comme le montre la figure 4(a), le calcul appliqué à chaque sortie de fonction O est Tanh(LayerNorm(O)). Tous les résultats de calcul sont concaténés dans un grand tenseur puis transmis à l'opérateur en aval (Op1 dans la figure 4(a)).
Le principal coût d’exécution ici est la surcharge de lancement du noyau GPU.Par exemple, le nombre de lancements du noyau GPU dans la figure 4(a) est de 2*N+3 (chaque ellipse dans la figure représente un noyau GPU). Cela affecte les performances puisque le temps d'exécution de LayerNorm et Tanh sur GPU est très court par rapport à leur temps de lancement du noyau.
De plus, l'opérateur Split peut créer une copie supplémentaire du tenseur de sortie du vecteur d'intégration, consommant ainsi de la mémoire GPU supplémentaire.
L'utilisation de FX pour implémenter une optimisation appelée fusion horizontale peut réduire considérablement le nombre de lancements du noyau GPU(Dans cet exemple, le nombre de lancements du noyau GPU après optimisation est de 5, voir Figure 4(b)).
Utilisez l'opérateur Add_middle_dim au lieu d'un Split explicite pour remodeler le tenseur d'intégration 2D de forme (B, NxD) en un tenseur 3D de forme (B, N, D). Ensuite, une seule LayerNorm est appliquée à sa dernière dimension. Appliquez un Tanh au résultat de LayerNorm. Enfin, l'opérateur Remove_middle_dim est utilisé pour restaurer le résultat Tanh en un tenseur 2D.
Étant donné que Add_middle_dim et Remove_middle_dim remodèlent simplement le tenseur,Aucune copie supplémentaire n'est créée, la consommation de mémoire GPU peut donc également être réduite.
3.3 Chevauchement entre calcul et communication
La formation des modèles de recommandation pour la production est généralement effectuée sur des systèmes GPU distribués.Étant donné que la capacité de mémoire de chaque GPU n'est pas suffisante pour contenir toutes les tables d'intégration dans le modèle, elles doivent être réparties sur plusieurs GPU.
Au cours de l'étape de formation, un GPU doit lire/écrire les valeurs des fonctionnalités de la table d'intégration sur d'autres GPU.C'est ce qu'on appelle une communication de tous à tous et cela peut entraîner une baisse significative des performances.
En mettant en œuvre une transformation via FX, il est possible de superposer le calcul avec la communication de tous à tous.La figure 5(a) montre un exemple de graphique de modèle avec accès à une table de vecteurs d'intégration (EmbeddingAllToAll) et d'autres opérateurs. Comme le montre la figure 5(b), sans aucune optimisation, ils sont exécutés séquentiellement sur un flux GPU.
Utilisez FX pour diviser EmbeddingAllToAll en EmbeddingAllToAll_Request et EmbeddingAllToAll_Wait, et organiser des opérateurs indépendants entre eux.
3.4 Résumé
Pour découvrir quels modèles bénéficieraient de ces transformations, les développeurs ont analysé les données de performance collectées par MAIProf pour les modèles exécutés dans le Meta Data Center.Nous montrons que ces transformations permettent d'obtenir une accélération de 2 à 3 fois supérieure sur un ensemble de modèles de production par rapport au mode impatient.
IV. Conclusion
D'un point de vue des performances, le mode graphique dans PyTorch est préféré au mode impatient utilisé dans les environnements de production. FX est un outil puissant pour capturer et optimiser les graphiques de programmes PyTorch. Cet article présente trois transformations FX pour optimiser les modèles de recommandation de production dans Meta.
Enfin, j’espère que davantage de développeurs PyTorch pourront utiliser la transformation graphique pour améliorer les performances du modèle.
—— Fin ——