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() Kannverschiedene Compiler-Backends einfach ausprobieren,Dies beschleunigt die Ausführung des PyTorch-Codes. Es als torch.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/cu117Für CPU:
pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpuOptional: Überprüfen der Installation
git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.pyOptional: 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-nightlyFü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/bashStart
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() Und torch.sin() sind Beispiele für punktweise Operationen, sie operieren Vektoren Element für Element. Eine bekanntere punktweise Operation ist torch.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ünden Die Fusion fand statt, weil die beiden Sünde Die 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 versuchen aot_cudagraphs oder nvfuser .
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 entfernen zu (Gerät = "cuda:0") Und codierte_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!