Reduzierung der Qwen-Trainingszeit um das 6-Fache auf AI Runtime

Das Trainieren oder Feintunen eines Foundation-Modells mit mehreren Milliarden Parametern erfordert erhebliche Rechenleistung, die sich oft über Hunderte bis Tausende von GPU-Trainingsstunden erstreckt. Das Optimieren der Trainingsperformance für ein LLM ist jedoch eine Herausforderung, und allgemeine Metriken wie die GPU-Auslastung zeigen möglicherweise nicht Ihre tatsächliche Auslastung. In diesem Blogbeitrag zeigen wir, wie Sie die Abstraktionsschichten durchbrechen und das Feintuning des neuesten Qwen 3.8 27B-Modells auf AI Runtime um das Sechsfache beschleunigen können. Durch systematische Performance-Optimierung reduzieren wir beim Feintuning eines festen Datensatzes die Zeit von 7 Stunden und 44 Minuten auf 1 Stunde und 17 Minuten – und das selbst bei einer bestehenden GPU-Auslastung von 98 %.

Wir bieten einen tiefen Einblick in die Techniken hinter diesen Verbesserungen: von der Messung der entscheidenden Faktoren über Implementierungsdetails bis hin zu den Design-Abwägungen bei jeder Optimierung.

Wenn Sie sofort loslegen möchten, springen Sie direkt zum Beispielcode auf AI Runtime und probieren Sie es aus!

Modell-Setup

Wir beginnen mit einer Standard-Trainingsschleife für Qwen 3.8 27B unter Verwendung der transformers-Bibliothek und des SFT-Trainingsdatensatzes HuggingFaceH4/ultrachat_200k.

• Wir wenden FSDP2 an, die neueste Modell-Sharding-Technik von torch. Dies ist erforderlich, um die Modellgewichte und den Optimizer-Status auf den GPUs unterzubringen. Ein rein verteilter, datenparalleler Ansatz würde daran scheitern, die Modellgewichte bei dieser Größe auf einer einzelnen GPU unterzubringen.
• Wir aktivieren auch Activation Checkpointing, eine Technik, die Speicherplatz gegen Rechenleistung im Backward-Pass eintauscht, um eine größere Batch-Größe pro GPU zu ermöglichen.
• Der Dataloader packt 2D-Batches mit maximal 8192 Tokens pro Batch und verwendet Padding bis zur maximalen Größe des Batch-Beispiels.

Erste Trainings-Benchmarks

Unser Baseline-Training des Modells über den gesamten Datensatz dauert 7 Stunden und 44 Minuten, und während des Trainings zeigt die GPU eine durchschnittliche Auslastung von 98 % an. Das Ergebnis ist ein Training mit 9.396 tok/s.

Diese GPU-Auslastung sieht hoch aus, aber wie bereits erwähnt, ist diese Zahl irreführend. Die Standardmetrik für die GPU-Auslastung misst nur den Prozentsatz der Zeit, in der die GPU irgendetwas getan hat. Sie sagt uns also nicht, wie gut das Modell die GPU genutzt hat, während sie aktiv war. Um das zu verstehen, müssen wir genauer hinsehen.

Eine Möglichkeit hierfür ist der Torch-Trace-Mechanismus. Wir können unsere Trainingsschleife wie folgt anpassen (und gleichzeitig die Anzahl der Schritte begrenzen, um eine schnelle Antwort zu erhalten):

Für einen Beispielschritt erhalten wir folgende Daten:

Tabelle 1: CPU – 14,4 s mit 1.714.382 Aufrufen. Die Tabelle enthält die Top-Operationen nach Gesamtzeit.

Tabelle 2: GPU – 10,3 s mit 269.448 Operationen. Die Tabelle enthält die Top-Operationen nach Gesamtzeit

Verbesserungsbereiche

Aus diesen Tabellen erkennen wir einige grundlegende Probleme beim aktuellen Modell:

• Die CPU benötigt länger zum Einrichten jedes Trainingsschritts als die GPU zu dessen Ausführung. Dies ist unerwartet, da die CPU-Seite der GPU-Ausführung im Allgemeinen asynchron ist, d. h. die CPU sollte der GPU deutlich voraus sein. Der Haupteffekt davon ist, dass die GPU ins Stocken gerät, woraus sich Optimierungspotenzial ergibt.
• Es gibt zu viele Operationen! In den obigen Tabellen haben wir Hunderte von Tausenden kleiner Operationen. Das Problem hierbei ist, dass die CPU GPU-Operationen nicht schnell genug in die Warteschlange stellen kann, um die GPU ausgelastet zu halten. Die Grundidee besteht darin, die Anzahl der Operationen zu reduzieren, die die CPU an die GPU übergibt. Es gibt einige allgemeine Techniken dafür, von denen eine zu unserer Situation passt:
  • Die CUDA-Graphen von torch zeichnen die erkannten CUDA-Kernel auf und geben sie mit einem einzigen CPU-Aufruf erneut wieder. Obwohl dies sehr leistungsfähig ist, gibt es viele Einschränkungen und es kollidiert mit der Dynamik von FSDP2. Vorerst kommt dies also nicht infrage.
  • Torch.compile verfolgt den Code und nutzt Triton, um Kernel nach Möglichkeit automatisch zu fusionieren. In diesem Beispiel verwenden wir manuell fusionierte Kernel, und im Fall von LLMs lässt dies keinen Raum für weitere Verbesserungen mit torch.compile.
  • Manuell fusionierte Kernel verwenden spezialisierte Fusionstechniken für bestimmte Operationen, die der Compiler von torch nicht von selbst anwenden kann. Glücklicherweise gibt es in der Community bereits viele etablierte, vorgefertigte Kernel. Beispiele hierfür sind „Flash Attention“, was in diesem Fall eine kleine Leistungsverbesserung bringt.

Außerdem haben wir ein Problem mit übermäßigem Padding im Batch: Der Eingabe-Tensor muss rechteckig sein, und jede Zeile wird bis zur maximalen Anzahl an Tokens für diesen Batch aufgefüllt (Padding). Wir zeigen einen Trick für LLMs, mit dem sich dies minimieren lässt.

In den folgenden Abschnitten werden wir tiefere Einblicke in die einzelnen technischen Themen geben und erklären, worum es sich handelt, welche Best Practices empfohlen werden und welche wesentlichen Abwägungen bei der Anwendung zu berücksichtigen sind.

Fusionierte Kernel

Sie werden feststellen, dass einige Kernel sehr häufig vorkommen! Insbesondere bei Qwen wird der Kernel für die Batch-Matrix-Multiplikation in einem einzigen Schritt über 30.000 Mal ausgeführt.

In diesem Fall schauen wir uns zuerst fusionierte Kernel an. Ein fusionierter Kernel fasst mehrere kleinere GPU-Kernel zu einer Einheit zusammen. Dies kann die Performance verbessern, indem kostenintensive Übertragungen zwischen dem High Bandwidth Memory einer GPU und ihrem noch schnelleren On-Die-Speicher vermieden werden. Außerdem ermöglicht es eine effizientere Skalierung und Auslastung der GPU-Kerne.

Bei so vielen zusätzlichen Kerneln funktioniert torch.compile zwar möglicherweise, aber die Ausführung nimmt viel Zeit in Anspruch. Die Reduzierung der reinen Anzahl an Ops durch manuell gefuste Kernel ist daher ein guter erster Schritt. Zudem sind bekannte Fusionen wie Flash Attention nichts, was torch.compile von selbst finden würde.

Die transformers-Bibliothek verwendet standardmäßig SDPA von PyTorch für die Standard-Attention-Layer. Es stellt sich jedoch heraus, dass viele dieser kleinen Operationen, die wir sehen, von den nicht-standardmäßigen Gated-Delta-Network-Layern stammen. Glücklicherweise gibt es für diese Layer zwei Kernel, causal-conv1d und flash-linear-attention, die entwickelt wurden, um viele dieser Operationen pro Layer zu fusionieren. Wie finden wir das heraus? Die Implementierung der transformers-Bibliothek von Qwen zeigt uns bereits, welche Kernel unterstützt werden:

• flash-linear-attention
• causal-conv1d

Es gibt zwei grundlegende Methoden, um diese Kernel zu erhalten:

• Installieren Sie die Kernel selbst. Das reicht von:
  • Einem einfachen pip install (z. B. für flash-linear-attention) von PyPI
  • Einem git clone und dem Kompilieren des entsprechenden Repositorys
  • Dem Suchen einer anderen Quelle für ein vorkompiliertes Paket, falls die Kompilierung zu lange dauert
• Verwenden Sie den Kernel Hub von Hugging Face, um vorkompilierte Kernel abzurufen.

Wir empfehlen generell die erste Option, insbesondere in Kombination mit unserer Docker-Container-Unterstützung: Erstellen Sie Ihren Container mit allen installierten Kerneln zur Build-Zeit, um die Startzeit Ihres Jobs zu verkürzen.

Kernel Hub ist zwar praktisch, aber da Kernel oft explizit von den installierten Torch- und CUDA-Versionen abhängen, ist für Ihre spezifische Konfiguration möglicherweise kein passender Kernel verfügbar. In diesem Fall greift das transformers-Modell stillschweigend auf einen unoptimierten Kernel mit geringer Leistung zurück.

Nach der Installation der Python-Pakete flash-linear-attention und causal-conv1d konnten wir unsere Ende-zu-Ende-Laufzeit um das 3,7-fache auf 2 Stunden und 5 Minuten reduzieren. In den folgenden Tabellen sehen Sie die neuen (hervorgehobenen) Kernel, die an die Stelle der elementweisen Funktionen aus den vorherigen Tabellen treten. Ebenso wurde die Gesamtzahl der ausgeführten Operationen drastisch reduziert – um das 15-fache:

Tabelle 3: CPU – 3,9 s bei 118.316 Aufrufen. Die Tabelle enthält die Top-Operationen nach Gesamtzeit.

Tabelle 4: GPU – 3,9 s bei 18.060 Aufrufen. Die Tabelle enthält die Top-Operationen nach Gesamtzeit.

Hinweis: Beim Training benötigen Sie sowohl einen Kernel für den Vorwärts- als auch für den Rückwärtsdurchlauf. Einige Kernel unterstützen jedoch nur den Vorwärtsdurchlauf, da sie lediglich zur Beschleunigung der Inferenz implementiert wurden.

Fazit

• Ziehen Sie die Verwendung gut etablierter, hochoptimierter Kernel in Betracht, die für gängige Modelle leicht verfügbar sind. Die Nutzung dieser Kernel kann bei relativ geringem Aufwand erhebliche Leistungssteigerungen bieten. In diesem spezifischen Fall haben wir eine 3,7-fache Steigerung der Trainingsgeschwindigkeit festgestellt, wodurch sich unsere Trainingszeit von 7 Stunden 44 Minuten auf 2 Stunden 5 Minuten reduzierte.

Batching, Tokenisierung und Padding

Unser Qwen-Modell ist jetzt ziemlich schnell, aber wir können noch schneller werden. Standardmäßig verwendet unser Batching-Setup rechteckige 2D-Tensoren mit 8.192 Tokens, da dies die größte Größe ist, die passt, ohne dass der Speicher ausgeht. Der Nachteil dieses Ansatzes besteht darin, dass wir durch Padding einen erheblichen Teil des Platzes im Tensor verschwenden.

Wenn Qwen nur Standard-Attention-Mechanismen verwenden würde, wäre NestedTensor von Torch gut unterstützt und würde helfen, diesen Aufwand zu reduzieren, indem das Padding vollständig entfernt wird. Allerdings unterstützt flash-linear-attention nur einen traditionelleren Ansatz mit gepackten 1D-Tensoren:

Aber wir können das nicht einfach willkürlich tun: Das einfache Übergeben des 1D-Tensors an das Modell mag für das Pretraining akzeptabel sein, aber für das Finetuning müssen wir sicherstellen, dass sich die Tokens in einem Beispiel nur auf Tokens desselben Beispiels beziehen. Um das zu beheben, müssen wir die korrekten position_ids des Beispiels an das Modell übergeben. Das Modell wiederum übergibt diese an die fusionierten Attention-Kernel, die dann die richtige Attention-Maske anwenden. Wenn Sie sich den Beispielcode hier ansehen, können Sie sehen, wie wir die erforderlichen Eingaben berechnen, damit dies funktioniert.

Mit dieser einfach aussehenden Änderung verkürzt sich unsere Trainingszeit von 2 Stunden 5 Minuten auf 1 Stunde und 17 Minuten.

Dafür gibt es zwei Gründe:

1. Der Hauptgrund ist, dass dieses Packing die Anzahl der Sequenzen pro Batch erhöht und die Gesamtzahl der Batches um etwa 36 % reduziert. Beachten Sie, dass dies von der Verteilung und Länge der Beispiele in Ihrem Datensatz abhängt.
2. Ein zweiter Grund ist, dass die Flash-Linear-Attention-Kernel für jede Eingabeform Just-in-Time kompiliert werden müssen und die Form zuvor dynamisch war, was zu einigen zusätzlichen Kompilierungsminuten führte.

Kompromiss: Diese Art der Änderung macht die Eingabe-Pipeline komplexer und birgt ein hohes Risiko für Fehler, die sich auf die Modellqualität auswirken. In diesem Fall sind die Einsparungen von 36 % es jedoch wert. Es gibt weitere Techniken zur Reduzierung des Paddings, wie z. B. die Beibehaltung des ursprünglichen Batching-Mechanismus bei gleichzeitigem Sortieren aller Eingaben nach Länge, sodass das Padding minimiert wird. Diese Reihenfolge kann jedoch unerwartete Korrelationen in benachbarten Samples erzeugen und das Modell verfälschen.

Fazit: Beim Training von LLMs über einen Datensatz mit Eingaben, die kürzer als Ihre maximal trainierbaren Tokens sind (z. B. in diesem Fall 8.192), kann das effiziente Packen von Daten in Batches eine erhebliche Zeitersparnis bringen. Für diesen Datensatz konnten wir eine 1,6-fache Steigerung der Trainingsgeschwindigkeit erzielen, wodurch sich unsere Trainingszeit von 2 Stunden 5 Minuten auf 1 Stunde und 17 Minuten verkürzte.

Zusammenfassung

In diesem Blogbeitrag haben wir Ihnen gezeigt, wie Sie Ihre Trainingseffizienz um das Sechsfache steigern können. Entscheidend waren (1) zuerst das Messen der relevanten Metriken – nicht nur der GPU-Auslastung, sondern auch der Bereiche, in denen Ihr Modell während des Trainings Zeit verbringt –, (2) das Optimieren und Fusionieren von Kerneln mithilfe bekannter Bibliotheken wie Flash Linear Attention und (3) die gezielte Anpassung der Zuführung von Datenbatches in das Modell.

Um dies im großen Maßstab umzusetzen, benötigen Sie eine Plattform, die einen performanten Weg zum Standard macht. In AIR sind viele dieser Leistungsverbesserungen bereits integriert. Wir liefern die neueste CUDA-Version und Flash Attention direkt in unseren Basisumgebungen mit, erfassen automatisch Kernmetriken wie GPU-Auslastung, -Temperatur und -Zustand für Sie und stellen sicher, dass Torch-Versionen (und die Profiler) auf dem neuesten Stand sind. Unser interner Agent, Genie Code, ist außerdem darauf abgestimmt, nach solchen Möglichkeiten zu suchen, um Ihr Training weiter zu optimieren. Darüber hinaus ist das Lesen von Daten aus Databricks durch unseren UCVolumeDataset optimiert.

Probieren Sie AIR doch einfach für Ihren nächsten Trainings-Workload aus! Bereit zum Ausprobieren? Im AI Runtime Cookbook finden Sie Details dazu, wie Sie dieses Modell selbst ausführen können.

(Dieser Blogbeitrag wurde mit KI-gestützten Tools übersetzt.) Originalbeitrag