Extrem lange Sequenzen, extrem schnell: LASP-Sequenzparallelität für eine neue Generation effizienter großer Sprachmodelle

Die AIxiv-Kolumne ist eine Kolumne, in der diese Website akademische und technische Inhalte veröffentlicht. In den letzten Jahren sind in der AIxiv-Kolumne dieser Website mehr als 2.000 Berichte eingegangen, die Spitzenlabore großer Universitäten und Unternehmen auf der ganzen Welt abdecken und so den akademischen Austausch und die Verbreitung wirksam fördern. Wenn Sie hervorragende Arbeiten haben, die Sie teilen möchten, können Sie gerne einen Beitrag leisten oder uns für die Berichterstattung kontaktieren. E-Mail-Adresse: liyazhou@jiqizhixin.com; zhaoyunfeng@jiqizhixin.com.

Von den internationalen Spitzenmodellen GPT-4 128K und Claude 200K bis zum heimischen „Red Fried Chicken“ Kimi Chat, der mehr als 2 Millionen Textwörter unterstützt, sind große Sprachmodelle (LLM) ausnahmslos in langen Kontexten aufgerollt Technologie. Wenn die klügsten Köpfe der Welt an etwas arbeiten, liegt die Bedeutung und Schwierigkeit der Sache auf der Hand.

Extrem lange Kontexte können den Produktivitätswert großer Modelle erheblich steigern. Mit der Beliebtheit von KI geben sich Benutzer nicht mehr damit zufrieden, mit großen Modellen zu spielen und ein paar Denksportaufgaben durchzuführen. Benutzer beginnen, den Wunsch zu verspüren, große Modelle zu verwenden, um die Produktivität wirklich zu verbessern. Schließlich kann die PPT-Datei, deren Erstellung früher eine Woche gedauert hat, jetzt in wenigen Minuten erstellt werden, indem man dem großen Modell einfach eine Reihe von Eingabeaufforderungen und ein paar Referenzdokumenten zufügt. Wer würde das als Berufstätige nicht lieben?

In letzter Zeit sind einige neue effiziente Sequenzmodellierungsmethoden entstanden, wie Lightning Attention (TransNormerLLM), State Space Modeling (Mamba), Linear RNN (RWKV, HGRN, Griffin) usw., die zu einer heißen Forschungsrichtung geworden sind. Forscher sind bestrebt, die bereits ausgereifte, sieben Jahre alte Transformer-Architektur zu transformieren, um eine neue Architektur mit vergleichbarer Leistung, aber nur linearer Komplexität zu erhalten. Diese Art von Ansatz konzentriert sich auf das Design der Modellarchitektur und bietet eine hardwarefreundliche Implementierung auf Basis von CUDA oder Triton, die eine effiziente Berechnung innerhalb einer Einzelkarten-GPU wie FlashAttention ermöglicht.

Gleichzeitig hat auch ein anderer Controller des Langsequenztrainings eine andere Strategie verfolgt: Sequenzparallelität gewinnt immer mehr an Bedeutung. Durch die Aufteilung der langen Sequenz in mehrere gleichmäßig aufgeteilte kurze Sequenzen in der Sequenzdimension und die Verteilung der kurzen Sequenzen auf verschiedene GPU-Karten zum parallelen Training sowie durch die Kommunikation zwischen Karten wird der Effekt des sequenzparallelen Trainings erreicht. Von der frühesten Colossal-AI-Sequenzparallelität über die Megatron-Sequenzparallelität bis hin zu DeepSpeed ​​​​Ulysses und in jüngerer Zeit Ring Attention haben Forscher weiterhin elegantere und effizientere Kommunikationsmechanismen entwickelt, um die Trainingseffizienz der Sequenzparallelität zu verbessern. Natürlich sind diese bekannten Methoden alle für traditionelle Aufmerksamkeitsmechanismen konzipiert, die wir in diesem Artikel Softmax Attention nennen. Diese Methoden wurden bereits von verschiedenen Experten analysiert, sodass in diesem Artikel nicht näher darauf eingegangen wird.

Dann kommt die Frage. Wie man mit der neuen effizienten Sequenzmodellierungsmethode Sequenzparallelität erreichen kann, um ein effizientes Training großer Sprachmodelle mit wirklich unendlicher Sequenzlänge (natürlich muss man über eine unbegrenzte GPU verfügen) über die Speicherbeschränkungen einer Einzelkarten-GPU hinweg zu erreichen, ist zu einem Problem geworden offene Frage. Ausgereifte sequenzparallele Methoden wie DeepSpeed ​​​​Ulysses und Megatron-SP können zwar auf lineare Sequenzmodellierungsmethoden angewendet werden, sind jedoch von Natur aus keine optimalen Lösungen, die auf Softmax Attention basieren.

• Papiertitel: Linear Attention Sequence Parallelism
• Papieradresse: https://arxiv.org/abs/2404.02882
• LASP-Codeadresse: https://github.com/OpenNLPLab/ LASP

Das LASP, das in diesem Artikel vorgestellt wird, ist entstanden. Forscher des Shanghai Artificial Intelligence Laboratory schlugen die LASP-Methode (Linear Attention Sequence Parallelism) vor, um die lineare Rechtsmultiplikationseigenschaft von Linear Attention vollständig zu nutzen und eine effiziente sequenzparallele Berechnung zu erreichen. Unter der Konfiguration einer A100 80G-GPU mit 128 Karten, des TransNormerLLM 1B-Modells und eines FSDP-Backends kann LASP die Sequenzlänge auf bis zu 4096 KB erweitern, was 4 MB entspricht. Im Vergleich zu ausgereiften sequenzparallelen Methoden beträgt die längste trainierbare Sequenzlänge von LASP das Achtfache der von Megatron-SP und das Vierfache der von DeepSpeed ​​​​Ulysses, und die Geschwindigkeit ist 136 % bzw. 38 % schneller.

Es ist zu beachten, dass der Name der Methode zur Verarbeitung natürlicher Sprache Linear Attention umfasst, sie ist jedoch nicht auf die Methode Linear Attention beschränkt, sondern kann weit verbreitet verwendet werden, einschließlich Lightning Attention (TransNormerLLM), State Space Modeling (Mamba), Lineares RNN (RWKV, HGRN, Griffin) und andere lineare Sequenzmodellierungsmethoden.

Einführung in die LASP-Methode

Um die Idee von LASP besser zu verstehen, schauen wir uns zunächst die traditionelle Berechnungsformel von Softmax an. Achtung: O=softmax((QK^T)⊙M)V, wobei Q, K, V, M und O sind Abfrage-, Schlüssel-, Wert-, Masken- und Ausgabematrizen. M ist hier eine untere dreieckige All-1-Matrix in einseitigen Aufgaben (z. B. GPT) und kann in zweiseitigen Aufgaben (z. B. BERT) ignoriert werden , es gibt keine Maskenmatrix für Zwei-Wege-Aufgaben. Zur Erläuterung werden wir LASP im Folgenden in vier Punkte aufteilen:

Prinzip der linearen Aufmerksamkeit

Lineare Aufmerksamkeit kann als eine Variante von Softmax Attention angesehen werden. Lineare Aufmerksamkeit entfernt den rechenintensiven Softmax-Operator, und die Berechnungsformel von Aufmerksamkeit kann als prägnante Form von O=((QK^T)⊙M) V geschrieben werden. Aufgrund der Existenz der Maskenmatrix M in der Einwegaufgabe kann diese Form jedoch immer noch nur eine linke Multiplikationsberechnung durchführen (d. h. zuerst QK^T berechnen), sodass die lineare Komplexität von O (N) nicht erhalten werden kann . Da es jedoch für bidirektionale Aufgaben keine Maskenmatrix gibt, kann die Berechnungsformel weiter auf O=(QK^T) V vereinfacht werden. Das Clevere an der linearen Aufmerksamkeit ist, dass die Berechnungsformel durch einfache Verwendung des assoziativen Gesetzes der Matrixmultiplikation weiter in Folgendes umgewandelt werden kann: O=Q (K^T V). Diese Berechnungsform wird als rechte Multiplikation bezeichnet Aufmerksamkeit ist verlockend O(N)-Komplexität kann in dieser bidirektionalen Aufgabe erreicht werden!

LASP-Datenverteilung

LASP teilt die langen Sequenzdaten zunächst in mehrere gleichmäßig unterteilte Teilsequenzen aus der Sequenzdimension auf und verteilt die Teilsequenzen dann an alle GPUs in der Sequenz-Parallelkommunikationsgruppe, sodass jede einzelne GPU verfügt eine Teilfolge zur parallelen Berechnung nachfolgender Folgen.

LASP-Kernmechanismus

Da GPT-ähnliche Modelle, die nur auf Decoder basieren, nach und nach zum De-facto-Standard für LLM werden, wurde LASP unter vollständiger Berücksichtigung des einseitigen Casual-Task-Szenarios entwickelt. Aus der segmentierten Teilsequenz Xi werden Qi, Ki und Vi berechnet, die entsprechend den Sequenzdimensionen segmentiert sind. Jeder Index i entspricht einem Chunk und einem Gerät (d. h. einer GPU). Aufgrund der Existenz der Maskenmatrix unterscheidet der LASP-Autor geschickt die Qi, Ki und Vi, die jedem Chunk entsprechen, in zwei Typen, nämlich: Intra-Chunk und Inter-Chunk. Unter diesen ist Intra-Chunk der Chunk auf der Diagonale, nachdem die Maskenmatrix in Blöcke unterteilt wurde. Es kann davon ausgegangen werden, dass die Maskenmatrix noch vorhanden ist und Inter-Chunk weiterhin verwendet werden muss Die Off-Diagonal-Linie der Maskenmatrix, die nicht berücksichtigt werden kann. Wenn die Maskenmatrix vorhanden ist, kann natürlich die rechte Multiplikation verwendet werden. Je mehr Chunks geteilt werden, desto kleiner wird der Anteil der Chunks auf der Diagonale Je mehr Blöcke die Aufmerksamkeit berechnet, desto größer ist der Anteil der Blöcke außerhalb der Diagonale. Unter anderem muss für die Berechnung des richtigen multiplizierten Inter-Chunks während der Vorwärtsberechnung jedes Gerät eine Punkt-zu-Punkt-Kommunikation verwenden, um den KV des vorherigen Geräts zu empfangen und seinen eigenen aktualisierten KV an das nächste Gerät zu senden. Bei der umgekehrten Berechnung ist es genau umgekehrt, mit der Ausnahme, dass die Objekte von Send und Recive zum Gradienten dKV von KV werden. Der Vorwärtsberechnungsprozess ist in der folgenden Abbildung dargestellt:

LASP-Code-Implementierung

Um die Recheneffizienz von LASP auf der GPU zu verbessern, führte der Autor Kernel Fusion für die Berechnungen von Intra-Chunk und Inter durch -Chunk bzw. Die Aktualisierungsberechnungen von KV und dKV sind auch in die Intra-Chunk- und Inter-Chunk-Berechnungen integriert. Um eine Neuberechnung des Aktivierungs-KV während der Backpropagation zu vermeiden, entschieden sich die Autoren außerdem dafür, ihn unmittelbar nach der Forward-Propagation-Berechnung im HBM der GPU zu speichern. Bei der anschließenden Backpropagation greift LASP zur Nutzung direkt auf den KV zu. Es ist zu beachten, dass die in HBM gespeicherte KV-Größe d x d beträgt und von der Sequenzlänge N völlig unabhängig ist. Wenn die Eingabesequenzlänge N groß ist, wird der Speicherbedarf von KV unbedeutend. Innerhalb einer einzelnen GPU implementierte der Autor Lightning Attention, das von Triton implementiert wurde, um den E/A-Overhead zwischen HBM und SRAM zu reduzieren und dadurch lineare Aufmerksamkeitsberechnungen für eine einzelne Karte zu beschleunigen.

Leser, die mehr Details erfahren möchten, können Algorithmus 2 (LASP-Vorwärtsprozess) und Algorithmus 3 (LASP-Rückwärtsprozess) im Artikel sowie den detaillierten Ableitungsprozess im Artikel lesen.

Verkehrsanalyse

Im LASP-Algorithmus ist zu beachten, dass die Vorwärtsausbreitung eine KV-Aktivierungskommunikation auf jeder Ebene des linearen Aufmerksamkeitsmoduls erfordert. Der Datenverkehr beträgt Bd^2/h, wobei B die Batch-Größe und h die Anzahl der Köpfe ist. Im Gegensatz dazu verwendet Megatron-SP eine All-Gather-Operation nach den beiden Layer-Norm-Schichten in jeder Transformer-Schicht und eine Reduce-Scatter-Operation nach den Attention- und FFN-Schichten, wodurch die Kommunikationsmenge 2BNd + 4BNd/T beträgt T ist die sequenzparallele Dimension. DeepSpeed-Ulysses verwendet eine All-to-All-Set-Kommunikationsoperation, um die Eingaben Q, K, V und Ausgabe O jeder Aufmerksamkeitsmodulschicht zu verarbeiten, was zu einem Kommunikationsvolumen von 4BNd/T führt. Der Vergleich des Kommunikationsvolumens zwischen den drei ist in der folgenden Tabelle dargestellt. Dabei ist d/h das Kopfmaß, normalerweise auf 128 eingestellt. In praktischen Anwendungen kann LASP das niedrigste theoretische Kommunikationsvolumen erreichen, wenn N/T>=32. Darüber hinaus wird das Kommunikationsvolumen von LASP nicht durch die Sequenzlänge N oder die Teilsequenzlänge C beeinflusst, was ein großer Vorteil für die parallele Berechnung extrem langer Sequenzen über große GPU-Cluster hinweg ist.

Datensequenz-Hybrid-Parallelität

Datenparallelität (d. h. Datensegmentierung auf Stapelebene) ist eine Routineoperation für verteiltes Training. Sie wurde auf der Grundlage der ursprünglichen Datenparallelität (PyTorch DDP) entwickelt Erreichen Sie eine speichersparendere Slice-Datenparallelität. Von der ursprünglichen DeepSpeed ​​​​ZeRO-Serie bis zum offiziell von PyTorch unterstützten FSDP ist die Slice-Datenparallelität ausgereift genug und wird von immer mehr Benutzern verwendet. Als Datensegmentierungsmethode auf Sequenzebene ist LASP mit verschiedenen datenparallelen Methoden kompatibel, darunter PyTorch DDP, Zero-1/2/3 und FSDP. Dies sind zweifellos gute Nachrichten für LASP-Benutzer.

Genauigkeitsexperiment

Experimentelle Ergebnisse zu TransNormerLLM (TNL) und Linear Transformer zeigen, dass LASP als Systemoptimierungsmethode mit verschiedenen DDP-Backends kombiniert werden kann und eine Leistung auf dem Niveau von Baseline erzielt.

Skalierbarkeitsexperiment

Dank des effizienten Kommunikationsmechanismusdesigns kann LASP problemlos auf Hunderte von GPU-Karten erweitert werden und behält eine gute Skalierbarkeit bei. „Geschwindigkeitsvergleichsexperiment“ Ulysses ist um 136 % bzw. 38 % schneller.

? Extrem schnelle Sequenzparallelität in wenigen Minuten.

Code-Portal: https://github.com/OpenNLPLab/LASP

Das obige ist der detaillierte Inhalt vonExtrem lange Sequenzen, extrem schnell: LASP-Sequenzparallelität für eine neue Generation effizienter großer Sprachmodelle. Für weitere Informationen folgen Sie bitte anderen verwandten Artikeln auf der PHP chinesischen Website!