L'équipe PyTorch a réécrit le modèle « tout diviser », qui est 8 fois plus rapide que l'implémentation originale

Comment devrions-nous optimiser le modèle « tout diviser » de Meta ? Ce blog écrit par l'équipe PyTorch vous aidera à y répondre du simple au plus profond.

Depuis le début de l'année jusqu'à aujourd'hui, l'IA générative s'est développée rapidement. Mais souvent, nous sommes confrontés à un problème difficile : comment accélérer la formation, le raisonnement, etc. de l’IA générative, notamment lors de l’utilisation de PyTorch.

Dans cet article, les chercheurs de l'équipe PyTorch nous apportent une solution. L'article se concentre sur la façon d'utiliser PyTorch natif pur pour accélérer les modèles d'IA génératifs. Il présente également de nouvelles fonctionnalités de PyTorch et des exemples pratiques sur la façon de les combiner.

Quel a été le résultat ? L'équipe PyTorch a déclaré avoir réécrit le modèle « Split Everything » (SAM) de Meta, ce qui a abouti à un code 8 fois plus rapide que l'implémentation d'origine sans perte de précision, le tout optimisé à l'aide de PyTorch natif.

Adresse du blog : https://pytorch.org/blog/accelerating-generative-ai/

Après avoir lu cet article, vous apprendrez :

• Torch.compile : PyTorch Compilateur de modèles, PyTorch 2.0 ajoute une nouvelle fonction appelée torch.compile (), qui peut accélérer les modèles existants avec une seule ligne de code
• Quantification GPU : accélère le modèle en réduisant la précision du calcul
• SDPA (Scaled ; Dot Product Attention) : implémentation d'une attention efficace en termes de mémoire ;
• Sparsité semi-structurée (2:4) : un format de mémoire clairsemé optimisé pour le GPU
• Nested Tensor : Nested Tensor regroupe {tensor, mask} ensemble pour regrouper des données de taille non uniforme dans un seul tenseur, telles que des images de différentes tailles ;
• Opérations personnalisées Triton : utilisez Triton Python DSL pour écrire des opérations GPU et les personnaliser. L'enregistrement de l'opérateur facilite son intégration dans divers composants de PyTorch. .

                                                                                      ^{Augmentation du débit et réduction de la surcharge de mémoire grâce aux fonctionnalités natives de PyTorch.}

SAM a été proposé par Meta. Pour plus d'informations sur cette recherche, veuillez vous référer à "CV n'existe plus ? Meta publie un modèle d'IA "tout diviser", CV peut inaugurer le moment GPT-3" .

Ensuite, l'article présente le processus d'optimisation SAM, y compris l'analyse des performances, l'identification des goulots d'étranglement et comment intégrer ces nouvelles fonctionnalités dans PyTorch pour résoudre ces problèmes rencontrés par SAM. De plus, cet article présente également quelques nouvelles fonctionnalités de PyTorch : torch.compile, SDPA, noyaux Triton, Nested Tensor et parcimonie semi-structurée (sparsité semi-structurée).

Le contenu de cet article est détaillé étape par étape. À la fin de l'article, la version rapide de SAM sera présentée. Les amis intéressés peuvent la télécharger sur GitHub. données via l'interface utilisateur Perfetto pour expliquer PyTorch La valeur d'application de chaque caractéristique.

Adresse GitHub : https://github.com/pytorch-labs/segment-anything-fast

Réécriture du modèle de segmentation n'importe quoi SAM

Cette étude a déclaré que cet article Le type de données de base SAM utilisé est le type float32 et la taille du lot est de 1. Les résultats de l'utilisation de PyTorch Profiler pour afficher la trace du noyau sont les suivants :

Cet article a révélé qu'il existe deux endroits où SAM peut être optimisé :

Le premier est un long appel à aten::index, qui est provoqué par l'appel sous-jacent généré par les opérations d'index tensoriel (telles que []) . Cependant, le temps réel que le GPU passe sur aten::index est relativement faible. La raison en est que pendant le processus de démarrage de deux cœurs, aten::index bloque cudaStreamSynchronize entre eux. Cela signifie que le CPU attend que le GPU termine le traitement jusqu'à ce que le deuxième cœur soit lancé. Par conséquent, afin d'optimiser SAM, cet article estime qu'il faut s'efforcer d'éliminer le blocage de la synchronisation GPU qui provoque des temps d'inactivité.

La seconde est que SAM passe beaucoup de temps GPU dans la multiplication matricielle (vert foncé dans l'image ci-dessus), ce qui est courant dans Transformers. Si nous pouvons réduire le temps GPU qu'un modèle SAM consacre aux multiplications matricielles, nous pouvons considérablement accélérer SAM.

Ensuite, cet article utilise le débit (img/s) et la surcharge mémoire (GiB) de SAM pour établir une référence. Vient ensuite le processus d’optimisation.

Bfloat16 demi-précision (plus la synchronisation GPU et le traitement par lots)

Pour résoudre le problème ci-dessus, c'est-à-dire faire en sorte que la multiplication matricielle prenne moins de temps, cet article se tourne vers bfloat16. Bfloat16 est un type de demi-précision couramment utilisé qui peut économiser beaucoup de temps de calcul et de mémoire en réduisant la précision de chaque paramètre et activation. 1 Utilisez BFLOAT16 pour remplacer le type de remplissage. De plus, afin de supprimer la synchronisation GPU, cet article constate qu'il existe deux positions qui peuvent être optimisées.

Plus précisément (c'est plus facile à comprendre en référence à l'image ci-dessus, les noms de variables qui apparaissent sont tous dans le code), l'étude a révélé que dans l'encodeur d'image SAM, il existe des variables de mise à l'échelle de coordonnées q_coords. et k_coords, ces variables sont allouées et traitées sur le CPU. Cependant, une fois ces variables utilisées pour indexer dans rel_pos_resized, ces opérations d'indexation déplaceront automatiquement ces variables vers le GPU, et cette copie provoquera une synchronisation GPU. Pour résoudre le problème ci-dessus, l'étude a noté que cette partie peut être résolue en la réécrivant à l'aide de torch.where comme indiqué ci-dessus.

Kernel Trace

Après avoir appliqué ces modifications, cet article a remarqué qu'il existe un écart de temps important entre les appels individuels du noyau, en particulier pour les petits lots (ici 1). Pour mieux comprendre ce phénomène, cet article commence par une analyse des performances de l'inférence SAM avec une taille de lot de 8 :

En regardant le temps passé par cœur, cet article observe que SAM dépense l'essentiel de son GPU time Sur les noyaux élément par élément et les opérations softmax.

Vous pouvez maintenant voir que le coût relatif de la multiplication matricielle est beaucoup plus faible.

En combinant la synchronisation GPU et l'optimisation bfloat16, les performances SAM sont améliorées de 3 fois.

Torch.compile (+ ruptures de graphiques et graphiques CUDA)

Cet article a révélé qu'il existe de nombreuses petites opérations en cours d'étude approfondie de SAM. Ils pensent que l'utilisation d'un compilateur pour fusionner des opérations est formidable. avantages, donc PyTorch Les optimisations suivantes ont été apportées à torch.compile :

• Fusion des séquences d'opérations telles que nn.LayerNorm ou nn.GELU en un seul noyau GPU
• Fusion des opérations immédiatement suivantes ; le noyau de multiplication matricielle, pour réduire le nombre d'appels au noyau GPU.

Grâce à ces optimisations, la recherche a réduit le nombre d'allers-retours dans la mémoire globale du GPU, accélérant ainsi l'inférence. Nous pouvons maintenant essayer torch.compile sur l’encodeur d’image de SAM. Pour maximiser les performances, cet article utilise des techniques de compilation avancées :

Traçage du noyau

Les résultats montrent que torch.compile fonctionne très bien.

On peut observer que softmax occupe une grande partie du temps, suivi de diverses variantes GEMM. Les mesures suivantes concernent des lots de 8 et plus.

SDPA : scaled_dot_product_attention

Ensuite, cet article a mené des expériences sur SDPA (scaled_dot_product_attention), en se concentrant sur le mécanisme d'attention. En général, les mécanismes d’attention natifs évoluent quadratiquement avec la longueur de la séquence en temps et en mémoire. Les opérations SDPA de PyTorch reposent sur les principes d'attention économes en mémoire de Flash Attention, FlashAttentionV2 et xFormer, qui peuvent accélérer considérablement l'attention du GPU. Combinée avec torch.compile, cette opération permet d'exprimer et de fusionner un motif commun dans des variantes de MultiheadAttention. Après un petit changement, le modèle peut désormais utiliser scaled_dot_product_attention.

Traçage du noyau

Vous pouvez maintenant voir que le noyau d'attention efficace en mémoire prend beaucoup de temps de calcul sur le GPU :

En utilisant le scaled_dot_product_attention natif de PyTorch, vous pouvez augmenter considérablement la taille du lot. Le graphique ci-dessous montre les changements pour les tailles de lots de 32 et plus.

Après cela, la recherche a également expérimenté Triton, NestedTensor, le traitement par lots Predict_torch, la quantification int8, la parcimonie semi-structurée (2:4) et d'autres opérations.

Par exemple, cet article utilise un noyau Triton positionnel personnalisé et observe les résultats de mesure avec une taille de lot de 32.

En utilisant Nested Tensor, la taille des lots varie de 32 et plus.

Mesures pour des tailles de lots de 32 et plus après l'ajout de la quantification.

La fin de l'article est une parcimonie semi-structurée. L’étude montre que la multiplication matricielle reste un goulot d’étranglement auquel il faut faire face. La solution consiste à utiliser la sparsification pour approximer la multiplication matricielle. Grâce à des matrices clairsemées (c'est-à-dire en mettant à zéro les valeurs), moins de bits peuvent être utilisés pour stocker les poids et les tenseurs d'activation. Le processus de définition des poids dans un tenseur qui sont mis à zéro est appelé élagage. La suppression de poids plus petits peut potentiellement réduire la taille du modèle sans perte significative de précision.

Il existe de nombreuses façons de tailler, de complètement non structuré à très structuré. Bien que l'élagage non structuré ait théoriquement un impact minime sur la précision, les GPU, bien que très efficaces pour effectuer de grandes multiplications matricielles denses, peuvent subir une dégradation significative des performances dans les cas clairsemés. Une méthode d'élagage récemment prise en charge par PyTorch vise à trouver un équilibre appelé parcimonie semi-structurée (ou 2:4). Ce stockage clairsemé réduit le tenseur d'origine de 50 % tout en produisant une sortie de tenseur dense. Voir l'illustration ci-dessous.

Afin d'utiliser ce format de stockage clairsemé et les noyaux rapides associés, la prochaine chose à faire est d'élaguer les poids. Cet article sélectionne les deux plus petits poids pour l'élagage avec une parcimonie de 2:4. Il est facile de changer les poids de la disposition PyTorch (« stried ») par défaut vers cette nouvelle disposition clairsemée semi-structurée. Pour implémenter apply_sparse (modèle), seules 32 lignes de code Python sont nécessaires :

Avec une parcimonie de 2:4, cet article observe les performances maximales de SAM lorsque vit_b et la taille du lot est de 32 :

Enfin, Pour résumer cet article en une phrase : Cet article présente l'implémentation de Segment Anything la plus rapide sur PyTorch à ce jour. Avec une série de nouvelles fonctionnalités officiellement publiées, cet article réécrit le SAM original dans PyTorch pur sans perdre en précision.

Les lecteurs intéressés peuvent consulter le blog original pour plus d'informations.

^{Lien de référence : https://pytorch.org/blog/accelerating-generative-ai/}

Ce qui précède est le contenu détaillé de. pour plus d'informations, suivez d'autres articles connexes sur le site Web de PHP en chinois!