torch.compile pour l'inférence

torch.compile compile le code PyTorch en noyaux optimisés qui accélèrent considérablement l'inférence. Cette fonctionnalité s'appuie sur TorchDynamo pour compiler le code en graphes et sur TorchInductor pour compiler ensuite ces graphes en noyaux optimisés. C'est un outil d'optimisation puissant qui, dans de nombreux cas, ne nécessite l'ajout que d'une seule ligne de code.

Envelopper un modèle avec torch.compile pour le compiler et retourner un modèle optimisé.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model)

Le premier appel à torch.compile est lent car le modèle doit être compilé. Les appels suivants au modèle compilé sont beaucoup plus rapides, car il n’a pas besoin d’être recompilé.

Plusieurs paramètres permettent de personnaliser le processus de compilation. Deux des plus importants sont listés ci-dessous. Pour la liste complète des paramètres, reportez-vous à la documentation de torch.compile.

Modes

Le paramètre mode propose plusieurs options de performance pour la compilation. Essayez différents modes pour déterminer lequel convient le mieux à votre cas d’usage.

• default est une option équilibrée entre la vitesse et la mémoire.
• reduce-overhead réduit le surcoût Python au prix d’une légère augmentation de la mémoire, mais il peut être plus rapide.
• max-autotune offre la vitesse la plus rapide, mais la compilation prend plus de temps.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead")

Fullgraph

Fullgraph tente de compiler l’ensemble du modèle en un seul graphe afin de maximiser les performances. torch.compile lève une erreur s’il rencontre une coupure de graphe, ce qui signifie qu’il ne peut pas compiler le modèle en un seul graphe.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("google/gemma-2b", device_map="auto")
compiled_model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

Benchmarks

Reportez-vous au tableau ci-dessous pour des benchmarks de performance comparant le temps d’inférence moyen en millisecondes avec torch.compile activé et désactivé sur divers GPU et tailles de lot, pour la même image et différentes tâches de vision.

Sélectionnez Subset dans le tableau ci-dessous pour basculer entre différents GPU, ainsi que pour consulter les benchmarks sur PyTorch nightly 2.1.0dev et torch.compile avec le mode reduce-overhead activé.

Étapes suivantes

• Consultez le guide torch.compile for training pour apprendre à utiliser torch.compile pour l'entraînement.