So skalierst du dein Modell: Eine systemische Perspektive auf LLMs auf TPUs

 (jax-ml.github.io)
7 Punkte von GN⁺ 2025-02-07 | 1 Kommentare | Auf WhatsApp teilen
  • Die Optimierung der Deep-Learning-Performance im großen Maßstab wirkt oft wie „Alchemie“, doch tatsächlich lässt sich die Modelleffizienz mit verständlichen, einfachen Prinzipien steigern
  • Von einem einzelnen Beschleuniger bis zu Zehntausenden von Beschleunigern gelten relativ einfache Prinzipien überall. Wer sie versteht, kann damit unter anderem Folgendes tun:
    • grob einschätzen, wie nah einzelne Teile eines Modells am theoretischen Optimum liegen
    • eine Grundlage schaffen, um bei unterschiedlichen Größenordnungen geeignete Parallelisierungstechniken auszuwählen
    • Kosten und Zeitaufwand für Training und Ausführung großer Transformer-Modelle abschätzen
    • Algorithmen entwerfen, die die Eigenschaften bestimmter Hardware ausnutzen
    • Hardware entwerfen, indem die Grenzen der aktuellen Algorithmusleistung klar verstanden werden
  • Erforderliches Vorwissen
    • Grundverständnis von LLMs und der Transformer-Architektur erforderlich
    • Verständnis großskaliger Betriebsweisen ist nicht zwingend nötig
    • Grundwissen zum LLM-Training und Erfahrung mit JAX sind von Vorteil
    • Empfohlen werden ein Blogpost zur Transformer-Architektur und Folien zur LLM-Skalierung in JAX
  • Ziele
    • die Fähigkeit entwickeln, abzuschätzen, wie ein Modell auf gegebener Hardware sinnvoll parallelisiert werden kann
    • die Fähigkeit entwickeln, Zeit- und Kostenaufwand für Training und Inferenz grob zu berechnen

Warum man sich dafür interessieren sollte

  • Noch vor 3–4 Jahren mussten die meisten ML-Forschenden über solche großskaligen Optimierungen nicht viel wissen
    • Heute arbeiten selbst „kleine“ Modelle nahe an den Hardwaregrenzen, weshalb ein Verständnis effizienter Verfahren im großen Maßstab unverzichtbar geworden ist
    • Die Geschichte des ML lässt sich als Wechselspiel zwischen Systeminnovationen und Softwareverbesserungen verstehen
    • Da aktuelle Transformer-Modelle die Hardwaregrenzen ausreizen, scheitern neue Architekturen oder Forschungsansätze in der Praxis leicht, wenn die Modelleffizienz nicht verstanden wird
    • Selbst wenn ein Benchmark 20 % Leistungsgewinn zeigt, ist der praktische Nutzen gering, wenn die Hardwareeffizienz zugleich um 20 % sinkt
  • Das zentrale Ziel der Modellskalierung ist es, den Durchsatz beim Hinzufügen weiterer Chips (Beschleuniger) linear zu steigern
    • Das wird als „Strong Scaling“ bezeichnet
    • Zusätzliche Chips verkürzen die Rechenzeit, verursachen aber Kommunikationskosten zwischen den Chips
    • Dauert die Kommunikation länger als die Berechnung, gerät man in einen „Communication-Bound“-Zustand, in dem Strong Scaling nicht mehr möglich ist
    • Wer die Hardware gut genug versteht, um vorherzusagen, wo diese Engpässe auftreten, kann Modelle so entwerfen oder umstrukturieren, dass sie vermieden werden
  • Ziel dieses Buchs ist es, zu erklären, wie TPU-(und GPU-)Hardware funktioniert und wie sich die Transformer-Architektur so entwickelt hat, dass sie auf heutiger Hardware gut läuft
    • Es soll sowohl Forschenden helfen, die neue Architekturen entwerfen, als auch Ingenieurinnen und Ingenieuren, die LLMs der aktuellen Generation möglichst schnell ausführen wollen

Gesamtüberblick

  • Dieser Text ist wie folgt aufgebaut
  • In Abschnitt 1 wird mit der Roofline-Analyse erklärt, welche Faktoren die Leistungsgrenzen eines Modells bestimmen (Kommunikation, Rechenleistung, Speicher)
  • In Abschnitt 2 und Abschnitt 3 geht es um den inneren Aufbau von TPUs und GPUs sowie um die Verbindung zwischen Chips
    • Damit werden unter anderem folgende Fragen beantwortet
      • Wie schnell kann eine Matrixmultiplikation einer bestimmten Größe theoretisch ausgeführt werden
      • Ab welchem Punkt wird die Berechnung durch Speicherbandbreite oder Kommunikationsbandbreite begrenzt
      • Wie ist ein TPU-Cluster aufgebaut und wie lange dauert es ungefähr, Daten von einem Chip zu einem anderen zu verschieben
      • Wie lassen sich verteilte Matrizen effizient multiplizieren
  • In Abschnitt 4 werden die Formeln der Transformer-Architektur (Matrixgrößen, Parameterzahl, FLOPs) im Detail behandelt
  • Abschnitt 5 und Abschnitt 7 bilden den Kern und stellen verschiedene Methoden vor, um Modelle über mehrere Chips zu parallelisieren
    • Data parallel, Tensor parallel, Pipeline parallel, Expert parallel
    • Außerdem werden Speicherspartechniken wie ZeRO, Rematerialisation, Host offload und Gradient accumulation behandelt
  • Abschnitt 6 und Abschnitt 8 zeigen am Beispiel des LLaMA-3-Modells auf TPUs, wie Training und Inferenz ablaufen und welche realen Kosten, Zeiten und Konfigurationen dabei anfallen
  • Abschließend behandeln Abschnitt 9 und Abschnitt 10 praktische Methoden zum Profiling, Debugging und zur Anwendung von Parallelverarbeitung in JAX

Im Detail: Zusammenfassung der wichtigsten Abschnitte des Buchs

Verwandte Beiträge

1 Kommentare

 
GN⁺ 2025-02-07
Hacker-News-Kommentare
  • Es gibt die Erwartung, dass JAX in den kommenden Jahren pytorch/cuda ersetzen wird. Das PTX-Problem mit dem Deepseek-Team zeigt den Wert, in Low-Level-Ansätze zu investieren, um die Hardwareleistung maximal auszuschöpfen
    • Wurde intern bei Google als Leitfaden für Performance-Arbeit verwendet. Es ist überraschend, dass das veröffentlicht wurde, aber Details zu Gemini scheinen entfernt worden zu sein
    • Gut an diesem Leitfaden ist, dass dank JAX/XLA ein direkter Wechsel zu GPUs möglich ist
    • Es gibt die Meinung, dass man sich fragt, warum JAX Tracing statt AST verwendet
    • Ein Link zum Tweet-Thread des Autors wurde geteilt
    • Jemand sucht nach einer Möglichkeit, eine Jekyll-Seite in PDF umzuwandeln
    • Es gibt Lob und Dank dafür, dass es ein großartiger Artikel ist
    • Es gibt die Meinung, dass man sich fragt, wie diese großartigen Animationen gemacht werden