Das PyTorch-Team implementierte das „Split Everything“-Modell achtmal schneller als die ursprüngliche Implementierung

Von Anfang des Jahres bis heute hat sich die generative KI rasant entwickelt. Aber oft stehen wir vor einem schwierigen Problem: Wie können wir das Training, die Argumentation usw. der generativen KI beschleunigen, insbesondere bei Verwendung von PyTorch?

In diesem Artikel liefern uns Forscher des PyTorch-Teams eine Lösung. Der Artikel konzentriert sich auf die Verwendung von reinem nativem PyTorch zur Beschleunigung generativer KI-Modelle. Außerdem werden neue PyTorch-Funktionen und praktische Beispiele für deren Kombination vorgestellt.

Was war das Ergebnis? Das PyTorch-Team gab an, das „Split Everything“ (SAM)-Modell von Meta neu geschrieben zu haben, was zu einem Code führte, der achtmal schneller als die ursprüngliche Implementierung ist, ohne an Genauigkeit zu verlieren, alles mit nativem PyTorch optimiert.

Blog-Adresse: https://pytorch.org/blog/accelerating-generative-ai/

Nachdem Sie diesen Artikel gelesen haben, werden Sie das folgende Verständnis erlangen:

• Torch kompilieren : PyTorch-Modell-Compiler, PyTorch 2.0 fügt eine neue Funktion namens Torch.compile () hinzu, die vorhandene Modelle mit einer Codezeile beschleunigen kann;
• GPU-Quantisierung: Beschleunigen Sie das Modell durch Reduzierung der Berechnungsgenauigkeit; (Scaled Dot Product Attention): eine speichereffiziente Aufmerksamkeitsimplementierung;
• Halbstrukturiert (2:4) Sparseness: ein für GPU optimiertes spärliches Speicherformat; } zusammen, um Daten uneinheitlicher Größe in einem einzigen Tensor zusammenzufassen, z. B. Bilder unterschiedlicher Größe;
• Benutzerdefinierte Triton-Operationen: Verwenden Sie Triton Python DSL, um GPU-Operationen zu schreiben, und integrieren Sie sie einfach über einen benutzerdefinierten Operator in verschiedene Komponenten von PyTorch Anmeldung.
• Die nativen Funktionen von PyTorch sorgen für einen höheren Durchsatz und einen geringeren Speicheraufwand.

Weitere Informationen zu dieser Forschung finden Sie im von Meta vorgeschlagenen SAM. Ausführliche Artikel finden Sie in „CV existiert nicht mehr? Meta veröffentlicht „Split Everything“-KI-Modell, CV könnte den GPT-3-Moment einläuten“

Analyse, Identifizierung von Engpässen und wie diese neuen Funktionen in PyTorch integriert werden können, um die Probleme von SAM zu lösen. Darüber hinaus werden wir auch einige neue Funktionen von PyTorch vorstellen, darunter Torch.compile, SDPA, Triton-Kernel, Nested Tensor und semi-strukturierte Sparsity (halbstrukturierte Sparsity)

Der Inhalt wird Schicht für Schicht vertieft Am Ende dieses Artikels stellen wir die schnelle Version SAM vor. Interessierte Leser können es von GitHub herunterladen. Darüber hinaus wurden diese Daten mithilfe der Perfetto-Benutzeroberfläche visualisiert, um den Anwendungswert verschiedener Funktionen von PyTorch zu demonstrieren. Dieses Projekt finden Sie unter der GitHub-Adresse: https://github.com/pytorch-labs/segment-anything-fast Der Quellcode von

schreibt das Split-Everything-Modell SAM neu

Die Studie weist darauf hin, dass der in diesem Artikel verwendete SAM-Basisdatentyp float32 dtype ist, die Stapelgröße 1 ist und PyTorch Profiler zum Anzeigen verwendet wird Ergebnisse der Kernverfolgung Wie folgt:

In diesem Artikel wurde festgestellt, dass SAM zwei Stellen hat, die optimiert werden können:

Der erste ist der lange Aufruf von aten::index, der durch verursacht wird Tensor-Indexoperation (z. B. []) Wird durch die zugrunde liegenden Aufrufe generiert. Die tatsächliche Zeit, die die GPU für aten::index aufwendet, ist jedoch relativ gering. Der Grund dafür ist, dass aten::index beim Starten von zwei Kernen cudaStreamSynchronize zwischen den beiden blockiert. Das bedeutet, dass die CPU darauf wartet, dass die GPU die Verarbeitung beendet, bis der zweite Kern gestartet wird. Um SAM zu optimieren, ist dieses Papier daher der Ansicht, dass man sich bemühen sollte, die blockierende GPU-Synchronisierung zu beseitigen, die Leerlaufzeiten verursacht.

Das zweite Problem besteht darin, dass SAM viel GPU-Zeit mit der Matrixmultiplikation verbringt (dunkelgrüner Teil, wie im Bild gezeigt), was im Transformers-Modell sehr häufig vorkommt. Wenn wir die GPU-Zeit des SAM-Modells bei der Matrixmultiplikation reduzieren können, können wir die Geschwindigkeit von SAM erheblich verbessern.

Als Nächstes vergleichen wir den Durchsatz (img/s) und den Speicheraufwand (GiB) von SAM Grundlinie. Dann gibt es noch den Optimierungsprozess

Der Satz, der neu geschrieben werden muss, lautet: Bfloat16 mit halber Präzision (plus GPU-Synchronisation und Stapelverarbeitung)

Um das oben genannte Problem zu lösen, ist das Das heißt, die für die Matrixmultiplikation erforderliche Zeit zu reduzieren. In diesem Artikel geht es um bfloat16. bfloat16 ist ein häufig verwendeter Typ mit halber Genauigkeit, der viel Rechenzeit und Speicher sparen kann, indem die Präzision jedes Parameters und jeder Aktivierung verringert wird Darüber hinaus wurde in diesem Artikel festgestellt, dass es zwei Orte gibt, die optimiert werden können, um die GPU-Synchronisierung zu entfernen der Bildencoder von SAM. Es gibt zwei Variablen q_coords und k_coords, die als Koordinatenskalierer fungieren, und diese Variablen werden auf der CPU zugewiesen und verarbeitet. Sobald diese Variablen jedoch zur Indizierung in rel_pos_resized verwendet werden, verschiebt der Indizierungsvorgang diese Variablen automatisch auf die GPU, was zu GPU-Synchronisierungsproblemen führt. Um dieses Problem zu lösen, haben Untersuchungen ergeben, dass dieser Teil durch Umschreiben mithilfe der Funktion „torch.where“ wie oben gezeigt gelöst werden kann Zeitabstand zwischen einzelnen Kernel-Aufrufen, insbesondere bei kleinen Batches (hier 1). Um ein tieferes Verständnis dieses Phänomens zu erlangen, haben wir mit der Leistungsanalyse der SAM-Inferenz mit einer Stapelgröße von 8 begonnen Ausgaben für elementweise Kernel und Softmax-Operationen

Jetzt können Sie sehen, dass der relative Overhead der Matrixmultiplikation viel geringer ist.

Durch die Kombination von GPU-Synchronisierung und bfloat16-Optimierung wird die SAM-Leistung um das Dreifache verbessert.

Torch.compile (+Graph Breaks und CUDA Graphs)

Während des Studiums von SAM habe ich viele kleine Operationen entdeckt. Forscher glauben, dass die Verwendung eines Compilers zur Konsolidierung dieser Vorgänge sehr vorteilhaft ist. Daher hat PyTorch die folgenden Optimierungen vorgenommen, um Torch.compile

Sequenzen von Vorgängen wie nn.LayerNorm oder nn.GELU in einem einzigen GPU-Kernel zu fusionieren

Fuse-Operationen unmittelbar nach dem Matrixmultiplikationskernel, um die Anzahl der GPU-Kernelaufrufe zu reduzieren.

Durch diese Optimierungen reduziert die Forschung die Anzahl der globalen GPU-Speicher-Roundtrips und beschleunigt so die Inferenz. Wir können Torch.compile jetzt auf dem Bildencoder von SAM ausprobieren. Um die Leistung zu maximieren, werden in diesem Artikel einige erweiterte Kompilierungstechniken verwendet:

Kernverfolgung

Den Ergebnissen zufolge schneidet Torch.compile sehr gut ab. Es ist zu beobachten, dass Softmax a großen Teil der Zeit , und dann jede GEMM-Variante. Die folgenden Maße gelten für Losgrößen ab 8 Stück.

SDPA:scaled_dot_product_attention

Als nächstes führte dieser Artikel Experimente zu SDPA (scaled_dot_product_attention) durch, wobei der Schwerpunkt auf dem Aufmerksamkeitsmechanismus lag. Im Allgemeinen skalieren native Aufmerksamkeitsmechanismen quadratisch mit der Sequenzlänge in Zeit und Gedächtnis. Die SDPA-Operationen von PyTorch basieren auf den speichereffizienten Aufmerksamkeitsprinzipien von Flash Attention, FlashAttentionV2 und xFormer, die die GPU-Aufmerksamkeit erheblich beschleunigen können. In Kombination mit Torch.compile ermöglicht dieser Vorgang den Ausdruck und die Fusion eines gemeinsamen Musters in Varianten von MultiheadAttention. Nach einer kleinen Änderung kann das Modell nun Scaled_dot_product_attention verwenden.

Kernverfolgung

Sie können jetzt sehen, dass der speichereffiziente Aufmerksamkeitskernel viel Rechenzeit auf der GPU beansprucht:

Mit PyTorch‘ s natives Scaled_dot _product_attention , ja Erhöht die Chargengröße erheblich. Die folgende Grafik zeigt die Änderungen für Losgrößen ab 32.

Als nächstes führte die Studie Experimente zu Triton, NestedTensor, Batch Predict_torch, int8-Quantisierung, halbstrukturierter (2:4) Sparsity und anderen Operationen durch

In diesem Artikel wird beispielsweise ein benutzerdefiniertes positionelles For verwendet Beim Triton-Kernel wurden Messungen mit einer Chargengröße von 32 beobachtet.

Verwendung der Nested-Tensor-Technologie und Anpassen der Chargengröße auf 32 und mehr

Nach dem Hinzufügen der Quantisierung ändern sich die Messergebnisse der Chargengröße von 32 und mehr.

Das Ende des Artikels ist halbstrukturierte Sparsity. Die Studie zeigt, dass die Matrixmultiplikation immer noch einen Engpass darstellt, der angegangen werden muss. Die Lösung besteht darin, die Sparsifizierung zur Annäherung an die Matrixmultiplikation zu verwenden. Durch dünn besetzte Matrizen (d. h. das Nullen der Werte) können weniger Bits zum Speichern von Gewichten und Aktivierungstensoren verwendet werden. Der Vorgang, bei dem festgelegt wird, welche Gewichte in einem Tensor auf Null gesetzt werden, wird als Pruning bezeichnet. Durch das Weglassen kleinerer Gewichte kann die Modellgröße möglicherweise ohne nennenswerten Genauigkeitsverlust verringert werden.

Es gibt viele Möglichkeiten zum Beschneiden, von völlig unstrukturiert bis stark strukturiert. Während unstrukturiertes Bereinigen theoretisch nur minimale Auswirkungen auf die Genauigkeit hat, kann es im spärlichen Fall der GPU zu erheblichen Leistungseinbußen kommen, obwohl sie bei der Durchführung großer Multiplikationen mit dichter Matrix sehr effizient ist. Eine kürzlich von PyTorch unterstützte Beschneidungsmethode ist die halbstrukturierte (oder 2:4) Sparsity, die darauf abzielt, ein Gleichgewicht zu finden. Diese spärliche Speichermethode reduziert den ursprünglichen Tensor um 50 % und erzeugt gleichzeitig eine dichte Tensorausgabe. Eine Erklärung finden Sie in der Abbildung unten.

Um dieses spärliche Speicherformat und den damit verbundenen schnellen Kernel zu verwenden, müssen Sie als Nächstes die Gewichte bereinigen. In diesem Artikel werden die kleinsten zwei Gewichte zum Bereinigen bei einer Sparse von 2:4 ausgewählt. Das Ändern der Gewichte vom standardmäßigen PyTorch-Layout („strided“) in dieses neue halbstrukturierte Sparse-Layout ist einfach. Um apply_sparse (Modell) zu implementieren, sind nur 32 Zeilen Python-Code erforderlich:

Bei einer Sparsity von 2:4 beobachten wir SAM-Spitzenleistung mit vit_b und einer Batchgröße von 32

Endlich Die Zusammenfassung dieses Artikels lautet wie folgt: Dieser Artikel stellt den bisher schnellsten Weg vor, Segment Anything auf PyTorch zu implementieren. Mithilfe einer Reihe offiziell veröffentlichter neuer Funktionen wird in diesem Artikel das ursprüngliche SAM in reinem PyTorch neu geschrieben kein Verlust an Genauigkeit

Für interessierte Leser können Sie den Originalblog für weitere Informationen überprüfen

Das obige ist der detaillierte Inhalt vonDas PyTorch-Team implementierte das „Split Everything“-Modell achtmal schneller als die ursprüngliche Implementierung. Für weitere Informationen folgen Sie bitte anderen verwandten Artikeln auf der PHP chinesischen Website!