Ärgern Sie sich noch über die „Schneckentempo“-Verarbeitung langer Texte durch große Sprachmodelle? Keine Sorge! Die Tsinghua-Universität präsentiert eine bahnbrechende Technologie – den APB-Sequenz-Parallel-Inferenz-Framework, der großen Sprachmodellen einen „Turbolader“ verleiht! Tests zeigen, dass diese Technologie bei der Verarbeitung sehr langer Texte bis zu 10-mal schneller ist als Flash Attention! Ja, Sie haben richtig gehört, das Zehnfache!
Mit dem Erfolg von ChatGPT und anderen großen Sprachmodellen hat sich auch die „Lesefähigkeit“ der KI deutlich verbessert. Die Verarbeitung von Texten mit mehreren zehntausend Wörtern ist keine Herausforderung mehr. Bei der Verarbeitung riesiger Datenmengen stoßen herkömmliche große Sprachmodelle jedoch an ihre Grenzen – die Transformer-Architektur ist zwar leistungsstark, aber ihr Kern, der Aufmerksamkeitsmechanismus, funktioniert wie ein „Superscanner“. Je länger der Text, desto exponentieller wächst der Scanbereich und desto langsamer wird die Verarbeitung.
Um dieses Problem zu lösen, haben Wissenschaftler der Tsinghua-Universität zusammen mit mehreren Forschungsinstituten und Technologieunternehmen einen neuen Ansatz entwickelt: den APB-Framework. Das Herzstück dieses Frameworks ist die geschickte Kombination aus „Sequenzparallelität“ und „spärlicher Aufmerksamkeit“.
Vereinfacht gesagt, funktioniert der APB-Framework wie ein hocheffizientes „Team“. Er zerlegt lange Texte in kleinere Abschnitte und verteilt sie an mehrere GPU-„Teammitglieder“ zur parallelen Verarbeitung. Besonders clever ist, dass APB jedem „Teammitglied“ die Fähigkeiten „lokale KV-Cache-Komprimierung“ und „effiziente Kommunikation“ verleiht. So können sie ihre Aufgaben parallel bearbeiten und gleichzeitig wichtige Informationen effizient austauschen, um komplexe semantische Abhängigkeiten in langen Texten gemeinsam zu lösen.
Noch überraschender ist, dass der APB-Framework die Geschwindigkeit nicht auf Kosten der Leistung erzielt. Im Gegenteil: Bei Tests mit extrem langen Texten (128K) ist APB nicht nur deutlich schneller, sondern übertrifft auch die Leistung von herkömmlichem Flash Attention! Sogar Star Attention, von Nvidia stark beworben, wird von APB übertroffen – mit einer Geschwindigkeitssteigerung um das 1,6-fache. Ein echter „Allrounder“.
Die wichtigste Anwendung dieser bahnbrechenden Technologie ist die drastische Verkürzung der Antwortzeit des ersten Tokens bei der Verarbeitung langer Texte durch große Sprachmodelle. Das bedeutet, dass große Sprachmodelle, die mit dem APB-Framework ausgestattet sind, lange Anweisungen von Benutzern sofort verstehen und blitzschnell reagieren können. Das lange Warten auf „Wird geladen…“ gehört der Vergangenheit an.
Wie erreicht der APB-Framework diese erstaunliche Beschleunigung?
Der APB-Framework kennt die Schwachstelle bei der Verarbeitung langer Texte: den Rechenaufwand. Der Rechenaufwand herkömmlicher Aufmerksamkeitsmechanismen ist proportional zum Quadrat der Textlänge. Lange Texte sind daher ein „Rechenloch“. Um diesen Engpass zu überwinden, setzt der APB-Framework zwei „Geheimwaffen“ ein:
Erster Trick: Steigerung des Parallelitätsgrades – „Viele Hände machen schnell fertig“
Der APB-Framework nutzt die Vorteile des verteilten Rechnens voll aus und verteilt die Rechenaufgaben auf mehrere GPUs. Ähnlich wie bei der „Teamwork“-Methode wird die Effizienz natürlich deutlich gesteigert. Besonders im Bereich der Sequenzparallelität zeigt der APB-Framework eine sehr hohe Skalierbarkeit und ist unabhängig von der Modellstruktur. Auch sehr lange Texte können problemlos verarbeitet werden.
Zweiter Trick: Reduzierung unnötiger Berechnungen – „Gut Ding will Weile haben“
Der APB-Framework führt einen spärlichen Aufmerksamkeitsmechanismus ein, der nicht „alles auf einmal“ berechnet, sondern „selektiv“. Er funktioniert wie ein Experte mit „scharfem Blick“ und konzentriert sich nur auf die wichtigen Informationen im Text, ignoriert unwichtige Teile und reduziert so den Rechenaufwand erheblich.
Die beiden Tricks „Parallelität“ und „Spärlichkeit“ erscheinen einfach, bergen aber „Geheimnisse“. Die eigentliche „Kernkompetenz“ des APB-Frameworks liegt darin, wie man unter dem Sequenz-Parallel-Framework eine effiziente spärliche Aufmerksamkeitsberechnung realisiert.
In einer Sequenz-Parallel-Umgebung verfügt jede GPU nur über einen Teil der Textinformationen. Eine spärliche Aufmerksamkeit mit „globalem Verständnis“ zu erreichen, ist wie „Elefanten fühlen“, eine große Herausforderung. Bisherige Methoden wie Star Attention und APE haben entweder die Leistung beeinträchtigt oder hatten eingeschränkte Anwendungsbereiche und konnten dieses Problem nicht perfekt lösen.
Der APB-Framework umgeht geschickt die „Problematik der umfangreichen Kommunikation“ und entwickelt einen neuen Ansatz: einen für sequenzparallele Szenarien geeigneten, kommunikationsarmen spärlichen Aufmerksamkeitsmechanismus. Die Kernkomponenten dieses Mechanismus umfassen:
Kleinere Anchor-Blöcke: Anchor-Blöcke funktionieren wie ein „Navigationssystem“, das den Aufmerksamkeitsmechanismus auf wichtige Informationen lenkt. Der APB-Framework reduziert innovativ die Größe der Anchor-Blöcke, wodurch sie leichter und flexibler werden und der Rechenaufwand reduziert wird.
Einzigartige Passing-Blöcke: Passing-Blöcke sind die „Seele“ des APB-Frameworks. Sie lösen auf clevere Weise das Problem der Langstrecken-semantischen Abhängigkeiten. Durch die „Komprimierung und Bündelung“ wichtiger Informationen der vorherigen GPU-Verarbeitung und die Weitergabe an nachfolgende GPUs kann jedes „Teammitglied“ den „Gesamtüberblick“ behalten und den „Kontext“ des langen Textes verstehen.
Abfrage-sensitives Kontext-Komprimierung: Der APB-Framework führt auch einen „Abfrage-sensitiven“ Mechanismus ein, damit der Kontext-Kompressor das „Problem verstehen“ und die wichtigen Informationen, die für die Abfrage relevant sind, genauer filtern und speichern kann, um die Effizienz und Genauigkeit weiter zu verbessern.
Basierend auf diesen „Geheimtechniken“ erstellt der APB-Framework einen reibungslosen Inferenzprozess:
Kontext-Segmentierung: Der lange Text wird gleichmäßig auf die einzelnen GPUs verteilt, und am Anfang wird ein Anchor-Block eingefügt, um die Abfragefrage „einzubetten“.
Kontext-Komprimierung: Mit Locret werden die KV-Caches „intelligent komprimiert“.
Effiziente Kommunikation: Mit dem AllGather-Operator werden die komprimierten KV-Caches an nachfolgende GPUs „weitergegeben“, um Passing-Blöcke zu erstellen.
Schnelle Berechnung: Es wird ein spezieller Flash Attention Kernel verwendet, zusammen mit einer optimierten Aufmerksamkeitsmaske für eine effiziente Berechnung. Der Passing-Block „scheidet nach Abschluss seiner Arbeit aus“ und nimmt nicht an weiteren Berechnungen teil.
Die Versuchsergebnisse belegen die hervorragende Leistung des APB-Frameworks. Bei Tests mit verschiedenen Modellen wie Llama-3.1-8B-instruct, Qwen-2.5-14B-instruct und Yi-34B-200K sowie verschiedenen Benchmarks wie InfiniteBench und RULER übertrifft der APB-Framework die Konkurrenz und erreicht das beste Gleichgewicht zwischen Leistung und Geschwindigkeit.
Besonders hervorzuheben ist, dass der Geschwindigkeitsvorteil des APB-Frameworks mit zunehmender Textlänge immer deutlicher wird und eine „je länger, desto schneller“-Wirkung erzielt. Der Grund dafür liegt darin, dass der Rechenaufwand des APB-Frameworks deutlich geringer ist als bei anderen Methoden, und dieser Unterschied vergrößert sich mit zunehmender Textlänge.
Eine detailliertere Analyse der Vorfüllzeit zeigt, dass die Sequenz-Parallel-Technologie die Rechenzeit für Aufmerksamkeit und FFN (Feedforward Neural Network) deutlich reduzieren kann. Der spärliche Aufmerksamkeitsmechanismus des APB-Frameworks komprimiert die Aufmerksamkeitsberechnungszeit weiter auf ein Minimum. Im Vergleich zu Star Attention nutzt der APB-Framework geschickt Passing-Blöcke, um Langstrecken-semantische Abhängigkeiten zu übertragen, reduziert die Größe der Anchor-Blöcke deutlich und reduziert so die zusätzlichen Kosten von FFN. Ein perfektes Beispiel für „beides gleichzeitig haben“.
Noch erfreulicher ist, dass der APB-Framework eine hervorragende Kompatibilität aufweist, sich flexibel an verschiedene verteilte Umgebungen und Modellgrößen anpassen lässt und unter verschiedenen „strengen“ Bedingungen eine stabile Hochleistung und -effizienz aufrechterhält.
Mit der Einführung des APB-Frameworks wird der „Engpass“ bei der Inferenz langer Texte durch große Sprachmodelle beseitigt, und der Spielraum für KI-Anwendungen wird unbegrenzt erweitert. Ob intelligenter Kundenservice, Finanzanalyse, wissenschaftliche Forschung oder Content-Erstellung – wir werden eine „schnellere, stärkere und intelligentere“ KI-Ära erleben!
Projekt-Adresse: https://github.com/thunlp/APB
Paper-Adresse: https://arxiv.org/pdf/2502.12085