Prise En Main De PyTorch 2.0 : Accélération Des Modèles HuggingFace Et TIMM !

PyTorch 2.0 peut augmenter la vitesse d'entraînement du modèle de 30% à 200% avec une simple ligne de torch.compile(). Ce tutoriel montrera comment reproduire réellement cette accélération.

torch.compile() **Peut****essayez facilement différents backends de compilateur,Cela accélère l’exécution du code PyTorch. C'est commetorch.jit.script()**Un remplacement direct pour , qui peut être exécuté directement sur nn.Module sans modifier le code source.

Dans l'article précédent, nous avons présenté que torch.compile prend en charge le code PyTorch arbitraire, le flux de contrôle, la mutation et, dans une certaine mesure, prend en charge les formes dynamiques.

En testant 163 modèles open source, nous avons constaté que torch.compile() peut apporter une accélération de 30%-200%.

opt_module = torch.compile(module)

Les résultats des tests sont détaillés dans :

Ce tutoriel vous montrera comment utiliser**torch.compile()**Accélérez la formation des modèles.

Exigences et paramètres

Pour les GPU (les GPU plus récents présentent des améliorations de performances plus significatives) :

pip3 install numpy --pre torch[dynamo] --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

Pour le processeur :

pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu

Facultatif : vérifier l’installation

git clone https://github.com/pytorch/pytorch  
cd tools/dynamo  
python verify_dynamo.py

Facultatif : installation de Docker

Toutes les dépendances nécessaires sont fournies dans le fichier PyTorch Nightly Binaries, qui peut être téléchargé via :

docker pull ghcr.io/pytorch/pytorch-nightly

**Pour les expériences ad hoc,**Assurez-vous simplement que le conteneur peut accéder à tous les GPU :

docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

commencer

Exemple simple

Prenons d’abord un exemple simple et remarquons que l’accélération est plus prononcée avec les GPU plus récents.

import torch  
   def fn(x, y):  
       a = torch.sin(x).cuda()  
       b = torch.sin(y).cuda()  
       return a + b  
   new_fn = torch.compile(fn, backend="inductor")  
   input_tensor = torch.randn(10000).to(device="cuda:0")  
   a = new_fn()

Cet exemple n'augmentera pas réellement la vitesse, mais il peut être utilisé pour commencer.

Dans cet exemple,torche.cos()ettorche.sin()sont des exemples d'opérations ponctuelles, elles opèrent sur des vecteurs élément par élément. Une opération ponctuelle plus célèbre esttorche.relu().

Les opérations point par point en mode impatient ne sont pas optimales car chaque opérateur doit lire un tenseur depuis la mémoire, apporter des modifications, puis réécrire ces modifications.

L’une des optimisations les plus importantes de PyTorch 2.0 est la fusion.

Ainsi, dans ce cas, nous pouvons transformer 2 lectures et 2 écritures en 1 lecture et 1 écriture, ce qui est essentiel sur les GPU plus récents où le goulot d'étranglement est la bande passante mémoire (la vitesse à laquelle les données peuvent être envoyées au GPU) plutôt que le calcul (la vitesse à laquelle le GPU peut effectuer des opérations en virgule flottante).

La deuxième optimisation importante de PyTorch 2.0 concerne les graphiques CUDA.

Les graphiques CUDA aident à éliminer la surcharge liée au lancement de noyaux individuels à partir de programmes Python.

torch.compile() prend en charge de nombreux backends différents, dont le plus notable est Inductor, qui peut générer des noyaux Triton.

Ces noyaux sont écrits en Python.**Mais c'est mieux que la plupart des noyaux CUDA écrits à la main.**En supposant que l'exemple ci-dessus s'appelle trig.py, vous pouvez réellement inspecter le code qui génère le noyau triton en exécutant

TORCHINDUCTOR_TRACE=1 python trig.py

@pointwise(size_hints=[16384], filename=__file__, meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'constants': {}, 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})  
   @triton.jit  
   def kernel(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):  
       xnumel = 10000  
       xoffset = tl.program_id(0) * XBLOCK  
       xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK])  
       xmask = xindex < xnumel  
       x0 = xindex  
       tmp0 = tl.load(in_ptr0 + (x0), xmask)  
       tmp1 = tl.sin(tmp0)  
       tmp2 = tl.sin(tmp1)  
       tl.store(out_ptr0 + (x0 + tl.zeros([XBLOCK], tl.int32)), tmp2, xmask)

À partir du code ci-dessus, nous pouvons voir que :péchésLa fusion s'est produite parce que les deuxpéchéLes opérations se déroulent dans un noyau Triton, et les variables temporaires sont stockées dans des registres, qui sont très rapides d'accès.

Exemple de modèle réel

Prenons l'exemple de resnet50 dans PyTorch Hub :

import torch  
   model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)  
   opt_model = torch.compile(model, backend="inductor")  
   model(torch.randn(1,3,64,64))

En fonctionnement réel, vous constaterez que la première exécution est très lente car le modèle est en cours de compilation. La vitesse de fonctionnement ultérieure sera plus rapide.Ainsi, avant de commencer l’analyse comparative, il est courant de réchauffer le modèle.

Comme vous pouvez le voir, nous utilisons ici « inductor » pour représenter le nom du compilateur, mais ce n'est pas le seul backend disponible. Vous pouvez l'exécuter dans le REPL**torch._dynamo.list_backends()**pour voir la liste complète des backends disponibles.

Vous pouvez également essayeraot_cudagraphsounvfuser.

Exemple de modèle de visage enlacé

La communauté PyTorch utilise souvent des modèles pré-entraînés de transformateurs ou de TIMM :

L’un des objectifs de conception de PyTorch 2.0 est que toute pile de compilation doit pouvoir être utilisée prête à l’emploi dans la grande majorité des modèles réellement exécutés.

Ici, nous téléchargeons un modèle pré-entraîné directement depuis le hub HuggingFace et l'optimisons :

import torch  
   from transformers import BertTokenizer, BertModel  
   # Copy pasted from here https://huggingface.co/bert-base-uncased  
   tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')  
   model = BertModel.from_pretrained("bert-base-uncased").to(device="cuda:0")  
   model = torch.compile(model) # This is the only line of code that we changed  
   text = "Replace me by any text you'd like."  
   encoded_input = tokenizer(text, return_tensors='pt').to(device="cuda:0")  
   output = model(**encoded_input)

Si vous le retirez du modèleà(appareil="cuda:0")etentrée codéePyTorch 2.0 générera des noyaux C++ optimisés pour fonctionner sur le CPU.

Vous pouvez consulter les noyaux Triton ou C++ de BERT, qui sont évidemment plus complexes que les exemples trigonométriques ci-dessus. Mais si vous connaissez PyTorch, vous pouvez l'ignorer.

Le même code peut être utilisé avec ce qui suit pour obtenir de meilleurs résultats :

* https://github.com/huggingface/accelerate

* DDP

Encore une fois, essayez l’exemple TIMM :

import timm  
   import torch  
   model = timm.create_model('resnext101_32x8d', pretrained=True, num_classes=2)  
   opt_model = torch.compile(model, backend="inductor")  
   opt_model(torch.randn(64,3,7,7))

**L’objectif de PyTorch est de créer un compilateur capable de s’adapter à davantage de modèles et d’accélérer le fonctionnement de la plupart des modèles open source.Visitez HuggingFace Hub maintenant,**Accélérez le modèle TIMM avec PyTorch 2.0 !