PyTorch 2.0 est là : tout ce que l’on sait

PyTorch est un framework de deep learning open source, porté par sa communauté, qui offre une approche flexible et performante pour créer des modèles de machine learning. Son interface est intuitive, la communauté est très active et l’intégration à l’écosystème Python est transparente.

Avec PyTorch 2.0, des changements fondamentaux ont été apportés aux opérations cœur du compilateur tout en préservant la familiarité et la simplicité d’usage pour les développeurs. Cette mise à jour promet des performances accrues ainsi qu’une prise en charge élargie des formes dynamiques (Dynamic Shapes) et du distribué.

Quoi de neuf dans PyTorch 2.0 ?

PyTorch rapatrie certaines pièces de C++ vers Python pour gagner en vitesse et en hackabilité. Avec la version 2.0, l’équipe a introduit torch.compile, qui modifie le fonctionnement de PyTorch au niveau du compilateur. Cette fonctionnalité est optionnelle et n’affecte pas votre ancien code.

Compilation avec PyTorch 2.0

Pour doter torch.compile d’une base solide, de nouvelles technologies ont été introduites :

• TorchDynamo. Un compilateur JIT au niveau Python, spécialement conçu pour accélérer PyTorch. En s’intégrant à l’API d’évaluation de frame de CPython, il modifie dynamiquement le bytecode Python à l’exécution afin d’accélérer l’exécution du code.
• AOTAutograd. Une boîte à outils qui aide les développeurs à accélérer l’entraînement des modèles sur PyTorch. Elle trace en amont les graphes avant et arrière. De plus, AOTAutograd propose des mécanismes simples pour compiler ces graphes extraits à l’aide de compilateurs de deep learning de dernière génération.
• PrimTorch. En réduisant drastiquement le nombre d’opérateurs PyTorch — de plus de 2000 à environ 250 opérateurs primitifs — PrimTorch a considérablement simplifié le développement de nouvelles fonctionnalités ou de backends PyTorch.
• TorchInductor. Un compilateur de deep learning natif PyTorch qui mappe automatiquement les modèles PyTorch vers du code généré pour plusieurs accélérateurs et backends. TorchInductor s’appuie sur OpenAI Triton comme brique de base pour l’accélération GPU.

Toutes ces nouveautés sont écrites en Python et prennent en charge les formes dynamiques. Le nouveau PyTorch exécute le code plus vite, offre davantage de flexibilité, se hacke facilement et abaisse la barrière d’entrée.

Exemples de code

Voyons une implémentation simple et rapide du compilateur PyTorch.

Sans torch.compile

import torch
model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")

Avec torch.compile

Pour booster les performances du modèle, il suffit d’ajouter le wrapper torch.compile autour du modèle pour obtenir une version compilée. C’est du plug-and-play.

import torch
model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")
compiled_model = torch.compile(model)

Apprenez à créer des modèles de deep learning avec la bibliothèque PyTorch en suivant le cours Deep Learning with PyTorch.

Vous pouvez entraîner votre modèle sans rien changer.

import torch
model = torch.compile(model)
for batch in dataloader:
    run_epoch(model, batch)

Ou exécuter simplement l’inférence.

model = torch.compile(model)
model(**input)

La fonction torch.compile() propose des paramètres supplémentaires :

1. mode : permet d’indiquer ce que le compilateur doit optimiser lors de la compilation.
2. dynamic : active le chemin de code pour les formes dynamiques.
3. fullgraph : compile le programme en un seul graphe.
4. backend : par défaut, TorchInductor est utilisé, mais vous pouvez spécifier d’autres backends de compilation disponibles.

def torch.compile(model: Callable,
  *,
  mode: Optional[str] = "default",
  dynamic: bool = False,
  fullgraph:bool = False,
  backend: Union[str, Callable] = "inductor",
  **kwargs
) -> torch._dynamo.NNOptimizedModule

Apprenez les fondamentaux de l’API PyTorch et créez un réseau de neurones simple de A à Z avec notre tutoriel PyTorch.

Benchmark

Pour démarrer avec PyTorch 2.0, les développeurs ont utilisé 163 projets open source (46 HuggingFace Transformers, 61 TIMM et 56 TorchBench) afin d’établir des benchmarks de performance de la nouvelle fonctionnalité de compilation. Le benchmark couvre notamment la classification d’images, la génération d’images, le modélisation de langage, les systèmes de recommandation et l’apprentissage par renforcement.

Les résultats montrent des gains de performance significatifs lors de l’entraînement sur des GPU NVIDIA A100.

Remarque : actuellement, le backend par défaut ne prend en charge que les CPU et les GPU Nvidia des séries Volta et Ampere.

Benchmark du compilateur PyTorch sur GPU NVIDIA A100

Ce n’est qu’un début ; les prochaines mises à jour apporteront encore plus de performances et de capacités de passage à l’échelle.

Comment installer PyTorch 2.0

Vous pouvez installer la nouvelle version de PyTorch simplement avec pip.

Copiez-collez la commande ci-dessous dans votre terminal.

Pour GPU : CUDA 11.8

Les versions récentes de GPU offrent des performances nettement supérieures.

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu118

Pour GPU : CUDA 11.7

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117

Pour CPU :

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cpu

Pour vérifier l’installation :

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

Accélérer Hugging Face avec PyTorch 2.0

Essayons la fonction torch.compile pour accélérer les transformers Hugging Face. Vous pouvez rendre votre code Hugging Face plus rapide avec un simple décorateur sur une ligne.

Remarque : avec torch.compile(), on observe un gain de 30 % à 200 % à l’entraînement – voir le tableau de bord des performances TorchDynamo.

Dans cet exemple, nous appliquons torch.compile à un grand modèle de langage « dolly-v2-3b » pour une inférence plus rapide. Pour exécuter le code dans Google Colab, nous devons d’abord installer les bibliothèques Python nécessaires.

%%capture
%pip install transformers accelerate xformers

Ensuite, nous téléchargeons et chargeons le tokenizer et le modèle de langage avec Hugging Face transformers. Puis nous passons le nn.Module à la fonction torch.compile().

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    GenerationConfig,
    pipeline,
)


tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-3b")
model = AutoModelForCausalLM.from_pretrained(
    "databricks/dolly-v2-3b", device_map="auto", torch_dtype=torch.bfloat16
)
model = torch.compile(model) #only line of code is required

Enfin, nous convertissons le texte en tokens avec le tokenizer, l’envoyons à model.generate, puis décodons la sortie générée en texte via tokenizer.batch_decode.

prompt = "I love you because"
inputs = tokenizer(prompt, return_tensors="pt").to(device="cuda:0")


# Generate
generate_ids = model.generate(inputs.input_ids, max_length=50)
tokenizer.batch_decode(
    generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
)[0]

Comme on le voit, « dolly-v2 » a complété la phrase par « you are a good person… »

"I love you because you are a good person.
You are kind, you help others, you are honest, you are loyal, you are humble, you are humble, you are humble.
You are a good person."

Cela fonctionne aussi avec le pipeline Hugging Face : il suffit d’indiquer le type de tâche, le modèle et le tokenizer.

generator = pipeline("text-generation", model= model,tokenizer=tokenizer)
generator("What is the name of Germany's Capital?")

Sortie :

[{'generated_text': "What is the name of Germany's Capital? 
The name of Germany's Capital is Berlin."}]

La fonction de compilation fonctionne immédiatement avec les bibliothèques Python transformers, accelerate et TIMM.

Si vous ne savez pas par où commencer pour lancer votre carrière en deep learning et en IA, le parcours Machine Learning Scientist with Python est un excellent point de départ. Il propose une vue d’ensemble complète des compétences Python clés pour décrocher un poste de machine learning scientist.