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() Peutessayez facilement différents backends de compilateur,Cela accélère l’exécution du code PyTorch. C'est comme torch.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/cu117Pour le processeur :
pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpuFacultatif : vérifier l’installation
git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.pyFacultatif : 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-nightlyPour 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/bashcommencer
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() et torche.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 est torche.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és La fusion s'est produite parce que les deux pé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 essayer aot_cudagraphs ou nvfuser .
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") et entrée codée PyTorch 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 !