Machine Learning: PyTorch 2.0 wird beschleunigen und sich auf Python besinnen

Die anstehende Version führt einen Python-basierten Kompiliermodus für eine deutliche Beschleunigung ein. Als Nightly-Feature steht er experimentell bereit.

In Pocket speichern vorlesen Druckansicht
Schlange, Python, Paradies

(Bild: Michael Schwarzenberger, gemeinfrei (Creative Commons CC0))

Lesezeit: 4 Min.
Von
  • Maika Möbus
Inhaltsverzeichnis

Im Rahmen der PyTorch Conference hat das Entwicklungsteam Details zur anstehenden Version PyTorch 2.0 bekanntgegeben. Sie liegt in den aktuellen Nightly Builds vor und fokussiert sich auf das neue Feature torch.compile, das das Kompilierverhalten zugunsten einer erhöhten Geschwindigkeit verändern soll und dessen Komponenten in Python geschrieben sind. Dabei soll das quelloffene Machine-Learning-Framework trotz des Sprungs in der Versionsnummer abwärtskompatibel bleiben.

PyTorch 2.0 bringt in der Nightly-Version den neuen experimentellen Kompiliermodus torch.compile mit, der die Performance deutlich beschleunigen soll und den ersten Schritt der Migration von Bestandteilen des Frameworks von C++ zurück zu Python markiert. Aufgrund dieser richtungsweisenden Änderung wird somit Version 2.0 auf die aktuelle Version 1.13 folgen. Neben dem optional einsetzbaren neuen Kompiliermodus wird in PyTorch 2.0 auch weiterhin die bekannte Eager Execution verwendbar sein, weshalb die Version vollständig abwärtskompatibel sein soll.

Hinter torch.compile stecken die in Python geschriebenen Technologien TorchDynamo, AOTAutograd, PrimTorch und TorchInductor. TorchDynamo dient dem Generieren von FX-Graphen anhand von Bytecode-Analyse, während AOTAutograd Backward-Graphen ahead-of-time generiert. PrimTorch bringt ein kleines Set an Operatoren zum Vereinfachen von Backends und TorchInductor ist ein Deep-Learning-Compiler, hinter dem Triton steckt – eine von OpenAI entwickelte Programmiersprache, die auf Deep Learning ausgelegt und für GPUs optimiert ist.

Auf seiner Website geht das PyTorch-Team auf die Gründe der Migration zunächst zu C++ und nun wieder zu Python ein. So seien seit der Veröffentlichung von PyTorch im Jahr 2017 Hardwarebeschleuniger wie GPUs etwa 15-mal schneller im Rechen- und etwa doppelt so schnell im Speicherzugriff geworden. Um eine hohe Performance der Eager Execution beizubehalten, mussten wichtige Teile von PyTorch zu C++ wechseln, was allerdings eine höhere Hürde für das Beitragen am Code darstelle.

Messungen des PyTorch-Teams zufolge funktioniert torch.compile bei 93 Prozent von 163 getesteten Open-Source-Modellen – aus verschiedenen Bereichen wie Vision und Natural Language Processing (NLP) – und läuft auf einer Nvidia-A100-GPU beim Training zwischen 38 und 76 Prozent schneller als der Eager Mode. Dabei handelt es sich um 46 Modelle von HuggingFace Transformers, 61 Modelle von TIMM und 56 Modelle von TorchBench.

Speedup-Zeiten für torch.compile verglichen mit dem Eager Mode auf einer Nvidia-A100-GPU

(Bild: The Linux Foundation)

Allerdings fiel dem PyTorch-Team auf, dass die Speed-ups bei Desktop-Class-GPUs wie Nvidia 3090 geringer waren als bei Server-Class-GPUs wie Nvidia A100. Das Standard-Backend TorchInductor kann derzeit mit CPUs und Nvidias Volta- und Ampere-GPUs umgehen. Weitere GPUs, xPUs oder ältere Nvidia-GPUs werden (noch) nicht unterstützt.

torch.compile lässt sich nach dem Installieren eines PyTorch Nightly Builds ausprobieren. Die neue Funktion wrappt ein Modell und gibt ein kompiliertes Modell zurück:

compiled_model = torch.compile(model)

Das PyTorch-Team zeigt den Einsatz von torch.compile anhand des Kompilierens und Ausführen eines echten Modells mit zufälligen Daten. Das erste Ausführen von compiled_model(x) kompiliert das Modell, was im Vergleich zu den nachfolgenden Ausführungen mehr Zeit in Anspruch nimmt.

import torch
import torchvision.models as models

model = models.resnet18().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
compiled_model = torch.compile(model)

x = torch.randn(16, 3, 224, 224).cuda()
optimizer.zero_grad()
out = compiled_model(x)
out.sum().backward()
optimizer.step()

Für verschiedene Anwendungsfälle stehen Presets bereit, um beispielsweise bei kleinen Modellen den Overhead zu reduzieren und somit eine höhere Geschwindigkeit zu erreichen.

# API NOT FINAL
# default: optimizes for large models, low compile-time
#          and no extra memory usage
torch.compile(model)

# reduce-overhead: optimizes to reduce the framework overhead
#                and uses some extra memory. Helps speed up small models
torch.compile(model, mode="reduce-overhead")

# max-autotune: optimizes to produce the fastest model,
#               but takes a very long time to compile
torch.compile(model, mode="max-autotune")

Die Entwicklung des 2017 von Meta (ehemals Facebook) gegründeten Open-Source-Frameworks PyTorch geschieht seit wenigen Monaten innerhalb der PyTorch Foundation, die wiederum unter dem Dach der Linux Foundation steht. Das PyTorch-Team hebt die Wichtigkeit des Feedbacks und der Mitarbeit der Entwicklercommunity hervor, um die 2.x-Serie zu verbessern. Für den Einstieg hat das Team Antworten zu 14 FAQs zu PyTorch 2.0 zusammengestellt und bietet weitere Ressourcen wie Tutorials und Modelle. Das erste stabile 2.0-Release soll im März 2023 erscheinen.

Alle weiteren Details zu PyTorch 2.0 lassen sich der PyTorch-Website entnehmen.

(mai)