PyTorch 2.0 Zum Anfassen: Beschleunigen Sie HuggingFace- Und TIMM-Modelle!

PyTorch 2.0 kann die Modelltrainingsgeschwindigkeit mit einer einfachen Zeile torch.compile() um 30%-200% erhöhen. Dieses Tutorial zeigt, wie diese Beschleunigung tatsächlich reproduziert werden kann.

torch.compile() **Kann****verschiedene Compiler-Backends einfach ausprobieren,Dies beschleunigt die Ausführung des PyTorch-Codes. Es alstorch.jit.script()**Ein Drop-In-Ersatz für , der direkt auf nn.Module ausgeführt werden kann, ohne den Quellcode zu ändern.

Im vorherigen Artikel haben wir vorgestellt, dass torch.compile beliebigen PyTorch-Code, Kontrollfluss und Mutation unterstützt und bis zu einem gewissen Grad dynamische Formen unterstützt.

Durch das Testen von 163 Open-Source-Modellen haben wir festgestellt, dass torch.compile() eine Beschleunigung von 30%-200% bringen kann.

opt_module = torch.compile(module)

Die Testergebnisse sind detailliert aufgeführt in:

Dieses Tutorial zeigt Ihnen, wie Sie**torch.compile()**Beschleunigen Sie das Modelltraining.

Voraussetzungen und Einstellungen

Für GPUs (neuere GPUs weisen deutlichere Leistungsverbesserungen auf):

pip3 install numpy --pre torch[dynamo] --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

Für CPU:

pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu

Optional: Überprüfen der Installation

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

Optional: Docker-Installation

Alle erforderlichen Abhängigkeiten werden in der PyTorch Nightly Binaries-Datei bereitgestellt, die über folgenden Link heruntergeladen werden kann:

docker pull ghcr.io/pytorch/pytorch-nightly

Für Ad-hoc-ExperimenteStellen Sie einfach sicher, dass der Container auf alle GPUs zugreifen kann:

docker run --gpus all -it ghcr.io/pytorch/pytorch-nightly:latest /bin/bash

Start

Einfaches Beispiel

Schauen wir uns zunächst ein einfaches Beispiel an. Dabei fällt auf, dass die Beschleunigung bei neueren GPUs stärker ausgeprägt ist.

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()

Dieses Beispiel erhöht die Geschwindigkeit zwar nicht wirklich, kann aber für den Einstieg verwendet werden.

In diesem Beispieltorch.cos()Undtorch.sin()sind Beispiele für punktweise Operationen, sie operieren Vektoren Element für Element. Eine bekanntere punktweise Operation isttorch.relu().

Punkt-für-Punkt-Operationen im Eager-Modus sind nicht optimal, da jeder Operator einen Tensor aus dem Speicher lesen, einige Änderungen vornehmen und diese Änderungen dann zurückschreiben muss.

Eine der wichtigsten Optimierungen in PyTorch 2.0 ist die Fusion.

In diesem Fall können wir also 2 Lese- und 2 Schreibvorgänge in 1 Lese- und 1 Schreibvorgang umwandeln, was bei neueren GPUs von entscheidender Bedeutung ist, bei denen der Engpass eher die Speicherbandbreite (wie schnell Daten an die GPU gesendet werden können) als die Rechenleistung (wie schnell die GPU Gleitkommaoperationen ausführen kann) ist.

Die zweite wichtige Optimierung von PyTorch 2.0 sind CUDA-Diagramme.

CUDA-Diagramme tragen dazu bei, den Aufwand beim Starten einzelner Kernel aus Python-Programmen zu eliminieren.

torch.compile() unterstützt viele verschiedene Backends, das bekannteste davon ist Inductor, das Triton-Kernel generieren kann.

Diese Kernel sind in Python geschrieben.**Aber es ist besser als die meisten handgeschriebenen CUDA-Kernel.**Angenommen, das obige Beispiel heißt trig.py, können Sie den Code, der den Triton-Kernel generiert, tatsächlich überprüfen, indem Sie Folgendes ausführen:

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)

Aus dem obigen Code können wir Folgendes ersehen:SündenDie Fusion fand statt, weil die beidenSündeDie Operationen finden in einem Triton-Kernel statt und temporäre Variablen werden in Registern gespeichert, auf die sehr schnell zugegriffen werden kann.

Reales Modellbeispiel

Nehmen Sie resnet50 im PyTorch Hub als Beispiel:

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))

Im tatsächlichen Betrieb werden Sie feststellen, dass der erste Durchlauf sehr langsam ist, da das Modell kompiliert wird. Die nachfolgende Betriebsgeschwindigkeit wird höher sein.Daher ist es üblich, das Modell vor dem Beginn des Benchmarkings aufzuwärmen.

Wie Sie sehen, verwenden wir hier „Induktor“, um den Compilernamen darzustellen, aber es ist nicht das einzige verfügbare Backend. Sie können es im REPL ausführen**torch._dynamo.list_backends()**um die vollständige Liste der verfügbaren Backends anzuzeigen.

Sie können auch versuchenaot_cudagraphsodernvfuser.

Beispiel für ein Hugging Face-Modell

Die PyTorch-Community verwendet häufig vortrainierte Modelle von Transformatoren oder TIMM:

Eines der Designziele von PyTorch 2.0 besteht darin, dass jeder Kompilierungsstapel in der überwiegenden Mehrheit der tatsächlich ausgeführten Modelle sofort einsatzbereit sein muss.

Hier laden wir ein vortrainiertes Modell direkt vom HuggingFace-Hub herunter und optimieren es:

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)

Wenn Sie es aus dem Modell entfernenzu (Gerät = "cuda:0")Undcodierte_Eingabe, PyTorch 2.0 generiert C++-Kernel, die für die Ausführung auf der CPU optimiert sind.

Sie können sich die Triton- oder C++-Kernel von BERT ansehen, die offensichtlich komplexer sind als die trigonometrischen Beispiele oben. Aber wenn Sie PyTorch kennen, können Sie es überspringen.

Um bessere Ergebnisse zu erzielen, kann derselbe Code mit Folgendem verwendet werden:

* https://github.com/huggingface/accelerate

* DDP

Versuchen Sie es noch einmal mit dem TIMM-Beispiel:

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))

**Das Ziel von PyTorch besteht darin, einen Compiler zu erstellen, der sich an mehr Modelle anpassen und die Ausführung der meisten Open-Source-Modelle beschleunigen kann.Besuchen Sie jetzt den HuggingFace Hub,**Beschleunigen Sie das TIMM-Modell mit PyTorch 2.0!