Flash Consideration ist ein leistungsoptimierender Transformator-Consideration-Mechanismus, der eine Effizienz von 15 % bietet

Foto von Sander Traa An Unsplash

Flash Consideration ist ein Transformator-Aufmerksamkeitsmechanismus zur Leistungsoptimierung, der im Hinblick auf die Echtzeitgeschwindigkeit eine Effizienz von 15 % bietet, und zwar ohne Annäherung.

Da Transformer-Modelle bei langen Sequenzen langsam und speicherhungrig sind (Zeit- und Speicherkomplexität sind quadratischer Natur), ist Flash Consideration (Papier) bietet eine Finish-to-Finish-Beschleunigung im Echtzeittakt von 15 % auf BERT-large und die dreifache Geschwindigkeit auf GPT-2.

Wenn man bedenkt, dass das Trainieren dieser großen Modelle enorme Mengen an Energie verbraucht, kann Flash Consideration mit Software program- und Hardwareoptimierung eine Effizienz von 15 % erzielen, was hinsichtlich der Verbesserung einen enormen Gewinn darstellt.

Nachfolgend werden einige der grundlegenden Konzepte von Flash Consideration und deren Implementierung erläutert.

Grundlegende Konzepte rund um Rechenleistung und Speicher

Bevor wir tiefer in die Thematik Rechenleistung und Speicher eintauchen, wollen wir sie noch einmal betrachten:

Was ist Compute?

  • Zeit, die Ihr GPU für die Berechnung tatsächlicher Gleitkommaoperationen (FLOPS) aufwendet

Was ist Speicher?

  • Zeitaufwand für die Übertragung von Tensoren innerhalb einer GPU

Im Idealfall soll unsere gCPU ständig Matrixmultiplikationen durchführen und nicht durch den Speicher eingeschränkt werden. In Wirklichkeit hat die Rechenleistung jedoch im Vergleich zum Speicher größere Fortschritte gemacht, und wir befinden uns in einer Welt, in der die gCPU im Leerlauf wartet, bis Daten geladen sind. Dies wird normalerweise als Speicher gebunden Operation. Siehe unten das erläuternde Diagramm, das dies darstellt. Matrixmultiplikation wird als Berechnung betrachtet und der Speicher speichert die Daten (betrachten Sie ihn als Lager). Die Berechnung benötigt Daten zur Verarbeitung und die Speicherbandbreite muss diese Operation unterstützen.

Foto von https://horace.io/brrr_intro.html

Was ist Speicherhierarchie?

Die A100 GPU hat 40–80 GB von Excessive Bandwidth Reminiscence mit einer Bandbreite von 1,5–2,0 TB/s Und 192 KB On-Chip-SRAM mit jeweils 108 Streaming-Multiprozessoren mit einer Bandbreite von ca. 19 TB/s.

Foto von https://arxiv.org/abs/2205.14135

Vor dem Hintergrund des oben genannten Kontexts ist Self-Consideration-Architektur erinnerungsgebunden.

Foto vom Autor

Betrachtet man die Aufmerksamkeitsmathematik, ist es eine Softmax-Operation, die die Speicherbindung verursacht.

  • Quantitativer Beweis: Wie Sie unten sehen können, nehmen Operationen wie Softmax, Dropout und Maskierung im Vergleich zur Matrixmultiplikation (Matmul) den Großteil der Zeit in Anspruch.
Foto von https://arxiv.org/abs/2205.14135

Warum wird Softmax zu einer speichergebundenen Operation?

Der Maßstab, in dem es betrieben wird, ist unser größter Engpass. Im folgenden Diagramm

  • N -> Anzahl der Token
  • d -> Anzahl der Einbettungsdimensionen
  • Wenn Abfrage und Schlüssel multipliziert werden, explodiert die Aufmerksamkeitsmatrix auf N * N, was viel Speicher beansprucht. Als Referenz (d ~128; N ~128.000 Token; Google Gemini: ~1 Million Token)
Foto von FlashAttention — Tri Dao | Stanford MLSys #67

Unten sehen Sie den Algorithmus zur Implementierung eines Selbstaufmerksamkeitsmechanismus

Foto von https://arxiv.org/abs/2205.14135

Wie im obigen Abschnitt erwähnt, ist das Übertragen von Informationen an HBM (Schreiben von S an HBM) und das anschließende Zurückladen von HBM an gCPU, um Softmax zu berechnen, und das anschließende Zurückschreiben an HBM eine Menge an Informationen, die es speichergebundener Vorgang.

Zusammen mit dem Diagramm erklären die folgenden Schritte, wie die Selbstaufmerksamkeit durch Matrixmultiplikation berechnet wird.

Schritt 1:

  • Ich habe das vereinfacht. In der Praxis wird jedes Token mit Positionscodierung hinzugefügt, um Einbettungen zu generieren, die in eine lineare Schicht eingespeist werden, um zu generieren. Zur Veranschaulichung habe ich eine Dimension von 3 verwendet (normalerweise liegt sie zwischen 64 und 128). Dies ist eine Standardeingabe für die Transformer-Architektur.

Schritt 2

  • Schlüssel -> Schlüssel‘ (Transponierung) wird berechnet und mit Abfrage multipliziert, um QK‘ zu erhalten, was N*N ist. Dies enthält die Aufmerksamkeit jedes Tokens mit den übrigen Token. Das folgende Diagramm zeigt auch die Beziehung. Da es sich um Token handelt und wir die Wichtigkeit jedes Tokens im Verhältnis zueinander berechnen müssen, wird die Softmax-Operation zeilenweise angewendet, um sie von 0 -1 zu normalisieren.
  • Dieser Schritt erfordert eine Bewegung zu HBM und ist der teuerste Vorgang wie wir besprochen haben. Das gesamte Flash-Aufmerksamkeitspapier befasst sich mit der Optimierung dieses Prozesses.

Schritt 3

  • Softmax(QK‘) * V wird als endgültige Ausgabematrix berechnet. Die Dimension ist hier dieselbe wie die Eingabeeinbettungen von Schlüssel, Abfrage und Wert.
  • Letzte Zeile in der Ausgabematrix
  • 1*5 bedeutet, dass die Einbettung von „diesem“ geändert werden sollte, um Beziehungen mit anderen Token einzubinden.
  • 2*5 bedeutet, dass die Einbettung von „ist“ geändert werden sollte, um Beziehungen mit anderen Token einzubinden.
  • Dasselbe wie oben für die restlichen Zeilen
Foto des Autors: Illustratives Diagramm zur Funktionsweise des Selbstaufmerksamkeitsmechanismus

Die Grundidee wird durch das folgende Diagramm erklärt, in dem Schlüssel-, Abfrage- und Wertblöcke vom HBM zum SRAM übertragen werden. Durch einige mathematische Methods (siehe unten) ist die hier durchgeführte Berechnung keine ungefähre Antwort, sondern eine tatsächlich richtige.

Mit dieser Implementierung kann Papier die Wall-Pace-Zeit reduzieren, indem es auf Informationen in Blöcken zugreift, ohne die Korrektheit zu beeinträchtigen.

Foto von https://arxiv.org/abs/2205.14135

Algorithmus hinter dem Papier: Wie wird Flash Consideration implementiert?

Dies ist der komplexeste Teil des Dokuments. Lassen Sie uns dieses Downside in Unteraspekte aufteilen und tiefer eintauchen.

Das folgende Diagramm unterteilt die Matrix in Blöcke und zeigt, wie jeder Block zur Berechnung des partiellen Softmax und anschließend des korrekten Softmax verwendet wird.

  • Ersteingabe: Token: Dies ist ein Flash-Consideration-Papier
  • Schlüssel: 4 (Token) x 3 (Dimensionen), Abfrage: 4 (Token) x 3 (Dimensionen) und Wert: 4 (Token) x 3 (Dimensionen)
Bild vom Autor geändert. Originalbild von https://arxiv.org/abs/2205.14135

Schritt 0

  • Angenommen, der Speicher ist 24 Byte groß
  • SRAM wird in 4 Blöcke unterteilt (Abfrage, Schlüssel, Wert und Ausgabematrix)
  • Abfrage, Schlüssel, Wert und Ausgabe erhalten jeweils = 6 Bytes zum Speichern ihrer Informationen (12 Bytes/4)
  • Jede Dimension ist 3, da jede Einbettung nicht gebrochen werden kann, additionally
  • Abfrage: 6 Bytes/ 3 (Dimension) = 2. Gleiches gilt für Wert, Schlüssel und Ausgabe
  • Daher gibt (M/4d) die Größe jedes Blocks an. In diesem Fall beträgt die Blockgröße 2. Das bedeutet, dass 2 Zeilen in den SRAM geladen werden können.
  • Im Allgemeinen beträgt die Blockgröße (M/4d) und die Anzahl der Blöcke (N*4D/M).

Schritt 1 und 2: Unten wird eine Tabelle hinzugefügt, die die Schritte 1 und 2 zur Funktionsweise von Flash Consideration veranschaulicht und die Speicher- und Rechenaspekte davon vergleicht.

Foto vom Autor: Schrittweise Aufschlüsselung der Speicher- und Rechenleistungsnutzung in Flash Achtung

Das folgende Diagramm hilft dabei, die in Flash Consideration verwendete Matrixmultiplikation (Block für Block) zu visualisieren.

Foto des Autors: Illustratives Diagramm zur Funktionsweise des Flash-Consideration-Mechanismus

Was ist der mathematische Aspekt von Softmax?

Einer der kritischsten Aspekte des Dokuments ist, wie das Zerlegen von Matrizen dennoch zur Berechnung der Softmax-Genauigkeit führt. Lassen wir das mathematische Beispiel unten aus, das zeigt, wie zwei verschiedene Matrizen zusammengefasst werden können, um erneut Softmax zu berechnen.

Instinct

  • Dies ist die schöne Eigenschaft von Exponenten, die hier ausgenutzt wird.
  • Jeder Softmax wird einzeln berechnet, aber gleichzeitig wird der Maximalwert der Zeile zusammen mit dem summierten Exponentenwert gespeichert.
  • Beim Zusammenführen mit einer anderen Matrix müssen wir prüfen, wie stark sich das Most vom globalen Most der beiden Matrizen unterscheidet. Und aufgrund des Exponenten werden sowohl Zähler als auch Nenner mit e^(aktuelles_Maximum – globales_Maximum) angepasst, um dies zu berücksichtigen.

Logik ist ziemlich komplex, deshalb wird unten ein Beispiel durchgegangen. Wenn Sie sich mit einem Beispiel vertraut gemacht haben, wird die obige Instinct sehr viel Sinn ergeben.

Foto vom Autor: Beispiel zur Demonstration, wie man eine Matrix in Unterkomponenten aufteilt und diese schließlich kombiniert, um Softmax zu berechnen

Schauen wir uns die Komplexitätsanalyse an, um ein Gefühl dafür zu bekommen, wie sich die Dinge verändert haben

Selbstaufmerksamkeit

  • Während der Berechnung von S = QK‘ wird eine N*N-Matrix erstellt, die zurück zum HRAM übertragen und anschließend aus dem HRAM abgerufen werden muss.
  • Daher ist O(N*N + N*N) = O(N*N) HBM-Zugriff

Blitzlicht

  • Äußere Schleife: Auf Schlüssel und Abfrage wird O(Nd) mal zugegriffen
  • Innere Schleife: Zum Laden von HBM sind nur O(Nd/M) erforderlich, da auf Blöcken gearbeitet wird
  • Insgesamt: O(N*N*d*d/M)
  • In der Praxis ist d viel kleiner als M. d liegt im Bereich von (64–128), während M im Bereich von 100 KB liegt. Daher ist der HBM-Zugriff optimiert
  • Wir begannen mit dem Ziel, den HBM-Zugriff zu optimieren, und mit dieser Komplexitätsanalyse sehen wir, dass das Papier die HBM-Zugriff über den Faktor (d*d/M) ohne Näherung.

So ein komplexes Papier mit enormer Effizienzsteigerung. Ich hoffe, die obige Erklärung vermittelt einen Eindruck davon, wie Flash Consideration die Leistung optimiert und verbessert. Ich habe Block Sparse Flash Consideration noch nicht behandelt. Wie sieht es im Vergleich zu anderen Optimierungstechniken, Forwards Cross Optimization usw. aus? Ich hoffe, das in einem zukünftigen Beitrag behandeln zu können.

Verweise

Von admin

Schreibe einen Kommentar

Deine E-Mail-Adresse wird nicht veröffentlicht. Erforderliche Felder sind mit * markiert