PyTorch 2.0 ist da: Alles, was wir wissen

PyTorch ist ein Open-Source-Framework für Deep Learning aus der Community, das einen flexiblen und effizienten Weg bietet, Machine-Learning-Modelle zu entwickeln. Es punktet mit einer benutzerfreundlichen Oberfläche, starker Community-Unterstützung und nahtloser Integration in das Python-Ökosystem.

PyTorch 2.0 hat grundlegende Änderungen an den zentralen Compiler-Operationen eingeführt und dabei die gewohnte Einfachheit für Entwickler beibehalten. Dieses Update verspricht höhere Performance sowie erweiterten Support für Dynamic Shapes und verteiltes Rechnen.

Was ist neu in PyTorch 2.0?

PyTorch verlagert Teile von C++ zurück nach Python, wird dadurch schneller und leichter anpassbar. Mit Version 2.0 kommt torch.compile hinzu, das die Arbeitsweise von PyTorch auf Compiler-Ebene verändert. Dieses Feature ist optional und beeinflusst deinen vorhandenen Code nicht.

PyTorch 2.0 Compile

Als Basis für torch.compile wurden neue Technologien eingeführt:

• TorchDynamo. Ein Just-in-Time-(JIT)-Compiler auf Python-Ebene, der speziell entwickelt wurde, um PyTorch zu beschleunigen. Durch die Integration mit der Frame-Evaluation-API in CPython modifiziert er Python-Bytecode dynamisch zur Laufzeit und sorgt so für schnellere Ausführung.
• AOTAutograd. Ein Toolkit, das Entwickler dabei unterstützt, das Training von Modellen in PyTorch zu beschleunigen. Es zeichnet Vorwärts- und Rückwärtsgraphen im Voraus auf und bietet einfache Mechanismen, um die extrahierten Graphen nahtlos mit modernen Deep-Learning-Compilern zu übersetzen.
• PrimTorch. Durch die Reduzierung der über 2000 PyTorch-Operatoren auf einen kompakten Satz von rund 250 primitiven Operatoren vereinfacht PrimTorch die Entwicklung neuer PyTorch-Features oder -Backends erheblich.
• TorchInductor. Ein PyTorch-nativer Deep-Learning-Compiler, der PyTorch-Modelle automatisch in generierten Code für verschiedene Beschleuniger und Backends übersetzt. TorchInductor nutzt OpenAI Triton als Baustein für GPU-Beschleunigung.

Alle neuen Technologien sind in Python geschrieben und unterstützen Dynamic Shapes. Dadurch läuft der neue PyTorch-Stack schneller, bleibt flexibel und ist leicht hackbar – die Einstiegshürde sinkt.

Codebeispiele

Schauen wir uns die schnelle und einfache Implementierung des PyTorch-Compilers an.

Ohne torch.compile

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

Mit torch.compile

Um die Modellleistung zu steigern, legst du einfach den torch.compile-Wrapper um das Modell und erhältst ein kompiliertes Modell. Plug-and-play.

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

Lerne, Deep-Learning-Modelle mit der PyTorch-Bibliothek zu bauen – mit dem Kurs Deep Learning with PyTorch.

Du kannst dein Modell ohne Änderungen trainieren.

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

Oder du führst Inferenz aus.

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

Die Funktion torch.compile() bietet zusätzliche Parameter:

1. mode: Legt fest, worauf der Compiler beim Kompilieren optimieren soll.
2. dynamic: Aktiviert den Codepfad für Dynamic Shapes.
3. fullgraph: Übersetzt das Programm in einen einzigen Graphen.
4. backend: Standard ist TorchInductor, du kannst aber andere verfgübare Compiler-Backends angeben.

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

Lerne die PyTorch-API kennen und erstelle ein einfaches neuronales Netz von Grund auf mit unserem PyTorch-Tutorial.

Benchmark

Zum Start mit PyTorch 2.0 wurden 163 Open-Source-Modelle (46 HuggingFace Transformers, 61 TIMM und 56 TorchBench) herangezogen, um Performance-Benchmarks für das neue Compile-Feature zu erstellen. Der Benchmark umfasst Aufgaben wie Bildklassifikation, Bildgenerierung, Sprachmodellierung, Empfehlungssysteme und Reinforcement Learning.

Die Ergebnisse zeigen deutlich höhere Performance beim Training auf NVIDIA A100 GPUs.

Hinweis: Aktuell unterstützt das Standard-Backend nur CPUs sowie Nvidia-Volta- und -Ampere-GPU-Serien.

PyTorch-Compiler-Benchmark auf NVIDIA A100 GPU

Das ist der Anfang – in künftigen Updates kommen weitere Verbesserungen bei Performance und Skalierbarkeit.

So installierst du PyTorch 2.0

Du kannst die neue Version von PyTorch einfach mit pip installieren.

Kopiere den folgenden Befehl in dein Terminal.

Für GPUs: CUDA 11.8

Neuere GPU-Generationen zeigen deutlich bessere Performance.

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

Für GPUs: CUDA 11.7

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

Für CPUs:

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

Zur Verifikation:

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

Hugging Face mit PyTorch 2.0 beschleunigen

Probieren wir die Funktion torch.compile aus, um Hugging Face Transformers zu beschleunigen. Mit einem einzigen Dekorator in einer Zeile läuft dein Hugging-Face-Code schneller.

Hinweis: Mit torch.compile() wurde beim Training ein Performance-Boost von 30% bis 200% beobachtet – siehe TorchDynamo Performance Dashboard.

Im folgenden Beispiel wenden wir torch.compile auf das große Sprachmodell „dolly-v2-3b" für schnellere Inferenz an. Um den Code in Google Colab auszuführen, installieren wir zuerst die benötigten Python-Bibliotheken.

%%capture
%pip install transformers accelerate xformers

Anschließend laden wir Tokenizer und Sprachmodell über Hugging Face Transformers. Danach übergeben wir das nn.Module an 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

Im letzten Schritt wandeln wir Text mit dem tokenizer in Tokens um, geben sie an model.generate und dekodieren die Ausgabe mit tokenizer.batch_decode wieder zu Text.

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]

Wie zu sehen ist, vervollständigt „dolly-v2“ den Satz mit „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."

Das Ganze funktioniert auch mit der Hugging-Face-Pipeline. Gib einfach Aufgabentyp, Modell und Tokenizer an.

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

Ausgabe:

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

Die Compile-Funktion funktioniert out of the box mit den Python-Bibliotheken transformers, accelerate und TIMM.

Wenn du nicht weißt, wo du mit deiner Deep-Learning- und KI-Karriere starten sollst, ist der Machine Learning Scientist with Python-Lernpfad ein starker Einstieg. Er vermittelt dir die wichtigsten Python-Kompetenzen, um als Machine-Learning-Scientist durchzustarten.