SwiftQK: Fast and Communication-Efficient Tensor Parallelism for Query-Key Normalization

Stellen Sie sich vor, Sie versuchen, ein riesiges, unglaublich kluges Robotergehirn zu betreiben, das Geschichten schreiben, mathematische Probleme lösen und wie ein Mensch chatten kann. Dieses „Gehirn“ wird als Large Language Model (LLM) bezeichnet. Um es zum Laufen zu bringen, verwenden Wissenschaftler Tausende von leistungsstarken Computerchips, die als GPUs bezeichnet werden. Aber diese Chips haben ein Problem: Sie sind wie brillante, aber winzige Arbeiter, die gleichzeitig nur eine kleine Menge an Informationen in ihren Taschen tragen können. Um große Probleme zu lösen, müssen sie zusammenarbeiten und Informationen hin und her weitergeben. Diese Teamarbeit wird „Tensor Parallelism“ genannt.

Es gibt jedoch einen kniffligen Teil des Robotergehirns, die sogenannte „Query-Key Normalization“. Stellen Sie sich dies als eine Qualitätskontrolle vor, bei der der Roboter sicherstellt, dass seine Gedanken ausgewogen und stabil sind, bevor er mit dem Schreiben beginnt. In der Vergangenheit musste jeder Arbeiter, um diese Kontrolle durchzuführen, sein gesamtes Notizbuch jedem anderen Arbeiter zeigen, damit sie einen einzigen „Balance-Score“ berechnen konnten. Dies bedeutete einen riesigen Verkehrsstau von Notizen, die hin und her gereicht wurden, was alles verlangsamte. Die Forscher in dieser Arbeit bemerkten, dass dieser Verkehrsstau der Hauptgrund war, warum ihr superschnelles Robotergehirn ins Stocken geriet. Sie wollten einen Weg finden, die Qualitätskontrolle durchzuführen, ohne dass jeder seine gesamten Notizbücher austauschen muss.

Hier kommt SwiftQK ins Spiel, ein cleverer neuer Trick, der von einem Team von Informatikern erfunden wurde, um diesen Verkehrsstau zu beheben. Anstatt jede GPU zu zwingen, ihr gesamtes Notizbuch auszutauschen (was riesig und langsam ist), ändert SwiftQK die Regeln des Spiels. Es erkennt, dass man, um den „Balance-Score“ zu berechnen, gar nicht das ganze Notizbuch braucht; man benötigt nur eine einzige Zahl, die die gesamte „Lautstärke“ oder „Energie“ der Notizen repräsentiert.

So funktioniert SwiftQK, unter Verwendung einer spielerischen Analogie: Stellen Sie sich eine Gruppe von Freunden vor, die versuchen, die Gesamtlautstärke eines Liedes zu ermitteln, das auf ihren Telefonen spielt. Auf die alte Art müsste jeder seine kompletten Songtexte in die Gruppe rufen, damit sie alle zusammengezählt werden können. Das dauert ewig. Mit SwiftQK flüstert jeder Freund nur eine einzige Zahl – die Gesamtlautstärke seines Liedes – der Gruppe zu. Die Gruppe addiert diese wenigen Zahlen, um die Gesamtlautstärke sofort zu erhalten.

Aber SwiftQK flüstert nicht nur; es macht etwas noch Klügeres. Während die Freunde ihre Zahlen untereinander flüstern, stehen die anderen nicht einfach nur herum und warten. Sie beginnen sofort mit ihrer eigenen Hausaufgabe (das Multiplizieren ihrer Notizen mit einem speziellen Gewicht). Das „Flüstern“ (Kommunikation) und die „Hausaufgabe“ (Berechnung) geschehen zur exakt gleichen Zeit, sodass sie perfekt überlappen und keine Zeit verschwendet wird. Die Forscher nennen dies einen „deadlock-safe persistent kernel“, was eine schicke Art und Weise ist zu sagen, dass sie ein System aufgebaut haben, in dem jeder genau weiß, wann er sprechen und wann er arbeiten muss, damit niemand auf einen Freund warten muss, der gerade beschäftigt ist.

Die Ergebnisse dieser neuen Methode sind beeindruckend. Als das Team SwiftQK bei aktuellen, leistungsstarken Sprachmodellen testete, stellte es fest, dass es den Schritt der „Qualitätskontrolle“ im Vergleich zur alten Methode des Austauschs ganzer Notizbücher um 81,4 % bis 93,9 % schneller machte. In einem realen Test, bei dem der Roboter gleichzeitig Fragen für viele Menschen beantwortete, reduzierte SwiftQK die Zeit, die es dauerte, eine Antwort zu erhalten (die sogenannte TPOT), um 29,5 % im Vergleich zur Standardmethode. Selbst im Vergleich zu einer leicht verbesserten Version der alten Methode, die ebenfalls versuchte, Zahlen zu flüstern, war SwiftQK immer noch 14,3 % schneller.

Die Arbeit zeigt, dass wir riesige KI-Gehirne viel reibungsloser und schneller laufen lassen können, indem wir klug entscheiden, welche Daten geteilt werden müssen, und indem wir sicherstellen, dass die Computer arbeiten, während sie kommunizieren. Es beweist, dass man nicht eine ganze Bibliothek versenden muss, um einen einzigen Tippfehler zu korrigieren; manchmal reicht es aus, nur eine einzige Zahl zu senden, solange man sie zum richtigen Zeitpunkt sendet.