Mamba-Implementierung in einer einzelnen PyTorch-Datei

 (github.com/johnma2006)
1 Punkte von GN⁺ 2023-12-21 | 1 Kommentare | Auf WhatsApp teilen
  • mamba-minimal ist ein Projekt, das Mamba einfach und minimal in einer einzigen PyTorch-Datei implementiert
  • Ziel ist es, dieselben numerischen Ausgaben wie die offizielle Implementierung im forward/backward pass zu erzeugen
  • Der Code ist vereinfacht und als kommentierte, gut lesbare Implementierung aufgebaut
  • Zentrale Optimierungen der offiziellen Implementierung sind nicht enthalten, daher bietet es keine Geschwindigkeit und enthält auch keine geeignete Parameterinitialisierung
  • Die Demo führt ein Beispiel zur Prompt-Vervollständigung mit state-spaces/mamba-370m und dem Tokenizer EleutherAI/gpt-neox-20b aus

Projektüberblick

  • mamba-minimal ist eine einfache Minimalimplementierung von Mamba in einer einzelnen PyTorch-Datei
  • Ziel ist es, dasselbe Verhalten wie die offizielle Implementierung mit besser lesbarem Code zu zeigen
  • Hauptmerkmale:
    • Gleichwertige numerische Ausgaben zur offiziellen Implementierung im forward pass und backward pass
    • Vereinfachter Code
    • Gut lesbare, kommentierte Implementierung

Nicht enthalten

  • Geschwindigkeit ist kein Ziel
    • Die offizielle Implementierung ist stark optimiert
    • Diese Optimierungen gehören zu den zentralen Beiträgen des Mamba-Papers
    • Diese Implementierung hält den Großteil des Codes zugunsten der Lesbarkeit bewusst einfach
  • Geeignete Parameterinitialisierung ist nicht enthalten
    • Sie wird als etwas genannt, das ergänzt werden könnte, ohne die Lesbarkeit zu beeinträchtigen

Beispiel für die Nutzung der Demo

  • In demo.ipynb ist ein Beispiel für Prompt-Vervollständigung zu sehen
  • Das Beispiel verwendet model.Mamba und den Hugging Face-transformers-AutoTokenizer
  • Verwendetes Modell und Tokenizer:
    • state-spaces/mamba-370m
    • EleutherAI/gpt-neox-20b
  • Der Beispiel-Prompt lautet Mamba is the; das generierte Ergebnis enthält einen Satz, der Mamba als Giftschlange beschreibt

Referenzen

Verwandte Beiträge

1 Kommentare

 
GN⁺ 2023-12-21
Hacker-News-Kommentare

  • Vor einiger Zeit habe ich zusammen mit einem Kollegen eine Bibliothek gebaut, die den gemeinsam genutzten Modellcode weitgehend auslagert. Damit lassen sich viele Modelle – abgesehen von Python-Imports und Kommentaren – in rund 100 Zeilen implementieren.
    BERT: https://github.com/explosion/curated-transformers/blob/main/...
    Llama 1/2: https://github.com/explosion/curated-transformers/blob/main/...
    MPT: https://github.com/explosion/curated-transformers/blob/main/...
    Auch Funktionen wie TorchScript JIT und PyTorch Flash Attention werden unterstützt.

    • Diese Bibliothek werde ich mir auf jeden Fall ansehen. Ich frage mich, ob du dir auch xformers angeschaut hast.
      xformers behandelt ein ähnliches Problem, konzentriert sich aber stärker darauf, mithilfe von Triton leistungsfähige Transformer-Module bereitzustellen. Allerdings war es nicht einfach, nur bestimmte Komponenten der Bibliothek herauszugreifen, und es traten ständig Runtime-Fehler auf, daher habe ich es vorerst beiseitegelegt. Ich baue gerade etwas auf Basis der BERT-Architektur und werde es mir als Referenz ansehen.
    • Ich bin beeindruckt von dieser Bibliothek. Die Hugging-Face-Implementierung mochte ich nicht besonders, aber das hier wirkt wie eine schöne API mit genau dem richtigen Abstraktionsniveau.
      Ich denke, ich werde sie im nächsten Projekt ausprobieren.

  • Der ursprüngliche Mamba-Code enthält viele Geschwindigkeitsoptimierungen und andere Elemente, sodass er nicht sofort leicht zu verstehen ist; diese Implementierung dürfte beim Lernen hilfreich sein.
    Wenn man Token einzeln inferiert, wird alles deutlich einfacher. Ich habe auch eine eigene Mamba-Inferenzimplementierung: https://github.com/rbitr/llm.f90/tree/master/ssm

    • Fortran also. Mich würde interessieren, warum Fortran verwendet wurde.
      Ich weiß, dass es die Grundlage für lange bewährten wissenschaftlichen Rechencode ist und oft von Bibliotheken wie PyTorch oder Numpy umhüllt wird, aber heutzutage ist es keine populäre Sprache. Mich interessiert, warum die Wahl darauf gefallen ist.

  • Es gibt Punkte zu Mamba, die ich gern so erklärt bekommen würde, dass sie auch jemand versteht, der kein Machine-Learning-Forscher ist

    1. Was ist die übergreifende Einsicht hinter State-Space-Modellen jenseits von Transformern?
    2. Welche inkrementellen Innovationen machen Mamba erfolgreicher oder interessanter als Vorgänger wie S4, H3 oder Monarch?
    3. Welche Bedeutung hat es außer der Skalierbarkeit unterhalb quadratischer Komplexität bei der Kontextlänge? Wenn man sich zum Beispiel nicht für Kontextlängen von mehr als 100k Token interessiert, frage ich mich, ob Mamba bei ähnlich großen Modellen und Datensätzen potenziell effizienter beim Trainingsaufwand ist
    • Meine Intelligenz liegt weit unter der der Paper-Autoren, aber ich habe trotzdem versucht, es zu verstehen. Ich habe Informatik studiert und habe grundlegende Kontrolltheorie sowie ein intuitives Verständnis diskreter Zeitsysteme auf Bachelor-Niveau, aber um dieses Paper richtig zu verstehen, müsste ich wohl State-Space-Modelle deutlich intensiver lernen
      Die zentrale Einsicht von Mamba besteht darin, ein altes Problem von State-Space-Modellen zu lösen. State-Space-Modelle sind gut darin, den Eingabekontext zu komprimieren, aber beim Komprimieren der Eingabe in einen Hidden State gehen Informationen verloren, die nötig sind, um den Kontext so effektiv zu nutzen wie ein Transformer
      Die Lösung ist der Aufbau dessen, was das Paper als Selektionsmechanismus bezeichnet. Dieser Mechanismus ist eingabeabhängig, sodass das Modell bei jeder veränderten Eingabe die Ausgabe jedes Schritts anpassen kann. Dafür werden einige State-Space-Variablen nicht eingabeinvariant, sondern eingabeabhängig gemacht, und es werden lineare Schichten und Ähnliches angefügt, um die Eingabe jedes Zeitpunkts auf State-Space-Variablen zu projizieren
      Wenn man State-Space-Variablen jedoch eingabeabhängig macht, entsteht Rechen-Overhead. Dies lösen sie mit einem hardwarebewussten Algorithmus, der die Speicherstruktur moderner GPUs maximal ausnutzt und Datenbewegungen in und aus HBM so weit wie möglich vermeidet
      Tri Dao ist der Entwickler von Flash Attention, und auch das war eine Methode, Hardware bei Transformern effizienter zu nutzen. Genau in diesem Bereich liegt wirklich seine Expertise
    • Attention wächst quadratisch mit der Kontextlänge, rekurrente neuronale Netze mit Gating (LSTM, GRU usw.) sind linear, und auch diese neuen Architekturen sind linear. Frühe rekurrente neuronale Netze nutzten Gating, um explodierende Gradienten zu vermeiden, aber die neuen Ansätze verwenden Theorie dynamischer Systeme, die Stabilität garantiert, sodass Gating nicht zwei Probleme gleichzeitig lösen muss und sich auf Gedächtnis konzentrieren kann
      Mamba und Based, die kurz vor NeurIPS 2023 erschienen, enthielten Multi-Query Associative Recall (MQAR) sowie Datenabhängigkeit von Gating/Selektion, inspiriert von Multi-Head Attention. Genau diese beiden Punkte fehlten Hyena und früheren State-Space-Architekturen; dadurch werden die neuen Modelle bei Aufgaben zum assoziativen Abruf so gut wie Attention und zeigen bei Aufgaben jenseits reiner Abfragen vielleicht sogar die Möglichkeit, etwas besser als Attention zu sein
      Ein großes Detail von Mamba ist natürlich die effiziente CUDA-Implementierung. Ohne sie könnte die Bedeutung dieser Architektur bei Aufgaben, für die Transformer bereits gut geeignet sind, geringer ausfallen
      Auch wenn man sich nicht stark um Kontextlänge sorgt, eröffnen sich viele neue Bereiche. DNA-Sequenzanalyse ist eine lineare Aufgabe mit langen Abhängigkeiten, und man kann auch Bilder, Videos und hochdimensionale Informationen als Token-Streams betrachten. So, als würde man Pixel abtasten wie bei alten CRT-Monitoren
      Einer der frühen Träume der KI war, dass sich eine einzelne Lerntrajektorie eines Agenten, der fortlaufend mit seiner Umgebung interagiert, kontinuierlich weiterentwickelt; solche Modelle mit unendlicher Kontextlänge könnten diesen Traum leichter erreichbar machen
      Derzeit sind nachgelagerte Anwendungen solcher Modelle für wichtige praktische Aufgaben im Vergleich zu den ausgereiften Attention-basierten Anwendungen jedoch insgesamt weniger validiert und weniger getunt. Die Analogie zu alten rekurrenten neuronalen Netzen hilft bis zu einem gewissen Grad, aber in den letzten fünf Jahren haben sich die Leute sehr stark auf Attention und Transformer spezialisiert, weshalb die Trägheit auf der Transformer-Seite groß ist
    • Ob Mamba bei ähnlich großen Modellen und Datensätzen recheneffizienter trainiert werden kann, würde ich selbst gern wissen
      Im Original-Paper wird erklärt, dass das Modell nach der Transformation der Parameter auf zwei Arten berechnet werden kann: als lineare Rekurrenz oder als globale Faltung. Beim Training, wo man im Allgemeinen die gesamte Eingabesequenz im Voraus sehen kann, nutzt man den leicht parallelisierbaren Faltungsmodus; bei autoregressiver Inferenz, wo die Eingaben schrittweise betrachtet werden, wechselt man in den effizienten rekurrenten Modus
      Das Training ist daher parallelisierbar, ähnlich wie RetNets paralleler Forward-Pass-Modus. Die Standard-Inferenz läuft im rekurrenten Modus, um möglichst langen Kontext zu erhalten; da es kein Chunking gibt, ist schwer zu beurteilen, wie viel RAM und VRAM während der Inferenz verbraucht wird
    • Dieses Video dürfte ziemlich genau das sein, wonach gesucht wurde
      Es erklärt das Paper und liefert zugleich viel Kontext dazu, wo es ins große Bild passt. Der Gedankengang ist ziemlich interessant anzuhören
      https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
    • Soweit ich weiß, ist Mamba im Grunde eine Fortsetzung der Forschung zu State-Space-Modellen, die man als lange Faltungen bezeichnen kann
      Statt quadratischer Attention, bei der berechnet wird, wie stark jedes Token auf jedes andere Token achtet, berechnet man irgendwie einen langen Faltungskernel mit derselben Länge wie die Eingabe und wendet dann conv1d an
      Nach meinem begrenzten Verständnis hat das ein wenig damit zu tun, FFT anzuwenden, Matrixmultiplikationen durchzuführen und dann per IFFT wieder zurückzugehen. Ich weiß, dass es funktioniert, aber es ist langsam. Es gibt mehrere Möglichkeiten, eine FFT zu berechnen, und eine davon ist eine Butterfly-Matrix. Vermutlich ist das nur eine Approximation, aber offenbar gut genug und auf aktueller Hardware sehr schnell und effizient
      Quadratische Komplexität klingt schlecht, aber in der Praxis sind Algorithmen unterhalb quadratischer Komplexität wegen Hardwarebeschränkungen oft langsamer. Daher ist es trotz großer Erwartungen an State-Space-Modelle nicht leicht zu sagen, dass Llama erledigt sei. Wir wissen auch noch nicht, ob Mamba beim Hochskalieren gut funktioniert, und um das herauszufinden, müsste man tatsächlich Millionen von Dollar ins Training stecken. Trotzdem bin ich optimistisch
      Ein weiteres interessantes Modell aus der Familie unterhalb quadratischer Komplexität ist RWKV. Es lohnt sich, es anzusehen, aber vermutlich wurde es im Podcast schon behandelt
      Ich habe mir das selbst beigebracht und das Paper früher auch nur grob überflogen, daher kann vieles falsch sein. Außerdem gibt es bei Attention normalerweise einen KV-Cache, der der Performance stark hilft; bei Mamba kann man das meiner Einschätzung nach nicht machen

  • Bei dem Satz „Mamba ist mit einer geschätzten Länge von über 150 m die längste Giftschlange der Welt“ musste ich lachen.
    Trotzdem wirklich hervorragend, und da auf das arXiv-Paper verwiesen wird, können auch Leute wie ich, die lieber solche Artikel konsumieren als das Paper direkt zu interpretieren, einen kleinen Blick ins Innere werfen.

    • Der Name Mamba ist gut. Aus [S]elective [S]tructured [S]tate [S]pace [S]equence models wird sSSSS, was wie ein Schlangenzischen klingt.
    • Ich dachte, die längste Giftschlange sei die Königskobra. Eine kurze Google-Suche ergab das auch.
      Es wäre amüsant, wenn sie später zu diesem Satz eine Korrektur veröffentlichen müssten.

  • Ich hatte erwartet, dass der Kern des Algorithmus ein paralleler Präfix-Scan ist. Das scheint doch der eigentliche Punkt von Mamba zu sein.
    for i in range(l):
    x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
    y = einsum(x, C[:, i, :], 'b d_in n , b n -> b d_in')
    ys.append(y)

  • Vielleicht eine dumme Frage, aber ich frage mich, wie schwierig es ist, die auf Hugging Face verfügbaren Mamba-Modelle zu trainieren.
    Das größte Modell scheint 2.8B zu haben; wie viele GPUs und wie viel Zeit bräuchte man wohl, um es auf einem Datensatz wie The Pile zu trainieren?

    • Eine ausgezeichnete Frage, die ich auch gern beantwortet hätte. Die Antwort scheint zu sein: deutlich schneller als ein Transformer gleicher Größe, und die Endergebnisse dürften in fast allen Benchmarks besser ausfallen als bei Transformern.
      Die Inferenz dürfte außerdem mit nur halb so viel RAM 3- bis 5-mal schneller laufen.

  • Ich hatte versucht, die offizielle CUDA-Version auseinanderzunehmen, bin aber nach dem ersten gescheiterten Versuch letztlich nicht weitergekommen; diese Implementierung sieht viel besser aus.

  • Noch eine Single-File-PyTorch-Implementierung – wirklich großartig. Ich hoffe, frühere Arbeiten wie hlb-CIFAR10 und verwandte Projekte sowie Vorläufer-Einflüsse wie minGPT oder DawnBench haben ein wenig dazu beigetragen, das einfache Single-File-Format voranzutreiben.
    Solche Arbeit ist wichtig für effiziente Machine-Learning-Forschung und könnte derzeit zu den wichtigsten Dingen gehören, die man für dieses Feld tun kann.
    Forschung schreitet mit der Geschwindigkeit von Innovation voran, und Innovation beschleunigt sich mit dem Kehrwert der Laufzeit von Experimenten; das hängt eindeutig mit der Kolmogorov-Komplexität des Codes für Forschungs- oder einfache Hacking-Zwecke zusammen.
    Man kann kaum genug betonen, wie wichtig solche Werkzeuge für die Forschung sind und wie sehr sie meinen persönlichen Prozess der Wissensentdeckung beschleunigt haben. Die Fähigkeit, Ideen in wenigen Minuten schnell zu skizzieren und sofort Ergebnisse mit hohem Signal-Rausch-Verhältnis zu bekommen, ist zu einem zentralen Bestandteil des Forschungsfortschritts geworden.
    Wissensdestillation und MDL (https://en.wikipedia.org/wiki/Minimum_description_length) halte ich für sehr wichtig, um unnötige Verzierungen, Ballast und den übermäßig dichten Wettbewerb um niedrigwertige „bloß nicht zurückfallen“-Themen zurückzudrängen, den der aktuelle Prozess für Paper-Einreichung und Review offenbar begünstigt.
    In letzter Zeit habe ich angefangen, Code als „Code-Skizzen“ zu veröffentlichen – kurze, in sich geschlossene Gists aus einer einzigen Datei –, weil ich dieses Problem vermeiden und zu einer etwas besseren Skalierungslösung kommen wollte. Das reduziert die Entwicklungszeit und bringt den Leuten unmittelbar groben, unpolierten, aber laufenden Code, der das Konzept enthält. Bisher scheint das ziemlich gut zu funktionieren, und ich möchte damit weitermachen.
    Ich würde gern mehr solchen Code sehen. Forschende, die Daten im großen Maßstab trainieren, sollten auch bei der Verbreitung von Informationen dateneffizient sein.

    • 2023 war schon allein deshalb ein spannendes Jahr, weil man beobachten konnte, wie sich KI-Forschung in absurdem Tempo entfaltet. Grundlagen wie ArXiV, PyTorch, GitHub, Hugging Face und knapper Open-Source-Python-Code beschleunigen die Entwicklung dieses neuen Feldes dramatisch.
      Wahrscheinlich hat die Menschheit noch nie etwas von beträchtlicher Komplexität so schnell weiterentwickelt.
      Einen ähnlichen Speed sieht man vielleicht bei SpaceX, das dieses Jahr auch zwei hochmoderne Raketen gestartet hat. Ich bin gespannt, was 2024 bringt.
    • Es gibt möglicherweise eine kleine Performance-Optimierung. Da x_proj hier keinen Bias hat, scheint es möglich, die Gewichte von x_proj und dt_proj zusammenzuführen.
      Falls es Anforderungen an eine Gewichtsanpassung gibt, ließe sich das zur Laufzeit vielleicht einfach machen, und ein einzelner Kernel samt Bias wäre am Ende vermutlich schneller. Sicher bin ich mir nicht.

  • Ich frage mich, ob es eine Diskussion zum ursprünglichen Paper gab. Die muss ich wohl verpasst haben, aber es ist ziemlich interessant.
    Den Teil „Aufgrund fehlender effizienter Implementierungen, die zu Speichermangel oder unrealistischen Rechenanforderungen führen, fehlen die vollständigen Ergebnisse für eine Kontextlänge von 8k bei den Baselines RWKV und RetNet, früheren starken rekurrenten Modellen, die ebenfalls als SSM interpretiert werden können“ habe ich nicht ganz verstanden.
    RetNet verbraucht nicht viel Speicher, und wenn man eine chunkweise Forward-Pass-Implementierung nutzt, ist der VRAM-Verbrauch durch die Chunk-Größe begrenzt. Genau das ist der Kern beim Testen der Kontextlänge.
    Ich frage mich, ob jemand das ursprüngliche Mamba-Modell getestet hat. Wie schnell ist das Training im Vergleich zu RetNet im parallelen Forward-Modus?

  • Ich mag Implementierungen immer, die Komplexes auf das Wesentliche reduzieren.