Databricks verkürzt Fine-Tuning von Qwen 3.8 27B auf AI Runtime um das Sechsfache

Eine hohe Auslastungszahl, die das Problem verdeckte

Databricks geht von einer Standard-Trainingsschleife für Qwen 3.8 27B aus, die mit der Bibliothek transformers und dem SFT-Datensatz HuggingFaceH4/ultrachat_200k aufgebaut ist. Das Modell wird mit FSDP2, der neuesten Sharding-Technik von torch, auf mehrere GPUs verteilt, damit Gewichte und Optimizer-Zustand in den Speicher passen. Activation Checkpointing ist aktiviert, damit eine größere Batchgröße pro GPU möglich ist. Der Dataloader packt 2D-Batches mit höchstens 8192 Tokens pro Batch und füllt per Padding auf das größte Beispiel im jeweiligen Batch auf.

Dieser Ausgangslauf benötigte 7 Stunden und 44 Minuten für den gesamten Datensatz und erreichte 9.396 tok/s, während die GPU eine durchschnittliche Auslastung von 98 Prozent anzeigte. Databricks hält diese Zahl für irreführend. Die Standardmetrik, so das Unternehmen, "zählt nur den Prozentsatz der Zeit, in dem die GPU irgendetwas getan hat", sagt also wenig darüber aus, wie gut die GPU in dieser Zeit tatsächlich genutzt wird.

Was der torch-Trace zeigte

Für einen genaueren Blick nutzte Databricks den torch-Trace-Mechanismus über eine begrenzte Zahl von Schritten. In einem Beispielschritt benötigte die CPU-Seite 14,4 Sekunden bei 1.714.382 Aufrufen, die GPU-Seite 10,3 Sekunden bei 269.448 Operationen. Das ist das Gegenteil dessen, was zu erwarten wäre, denn normalerweise läuft die CPU der GPU deutlich voraus. Die Folge: Die GPU wartet auf Arbeit.

Als Ursache identifizierte das Unternehmen zu viele kleine Operationen, die die CPU nicht schnell genug einreihen kann, um die GPU auszulasten. Ziel wurde daher, die Zahl der Operationen zu senken. Databricks erwog mehrere allgemeine Techniken:

• CUDA Graphs: Sie zeichnen Kernel auf und spielen sie mit einem einzigen CPU-Aufruf wieder ab, haben aber viele Einschränkungen und stehen im Konflikt mit der Dynamik von FSDP2, weshalb sie ausschieden.
• torch.compile: Es fusioniert Kernel automatisch mit Triton, doch Databricks entschied sich für manuell fusionierte Kernel, die bei LLMs nach eigener Aussage keinen Spielraum für Verbesserungen durch torch.compile lassen.
• Manuell fusionierte Kernel: spezialisierte Fusionen, die der Compiler von torch nicht von selbst vornehmen kann und von denen viele in der Community bereits existieren.

Ein eigenes Problem war übermäßiges Padding, weil jede Zeile des rechteckigen Eingabetensors auf die maximale Tokenzahl ihres Batches aufgefüllt wird. Databricks sagt, es nutze für LLMs einen Trick, um dies zu minimieren.

Fusionierte Kernel bringen den ersten großen Gewinn

Im Trace lief der Kernel für die Batch-Matrixmultiplikation in einem Schritt mehr als 30.000 Mal. Databricks erläutert, dass ein fusionierter Kernel mehrere kleinere GPU-Kernel zu einer Einheit zusammenfasst. Das vermeidet teure Transfers zwischen dem High Bandwidth Memory und dem schnelleren On-Die-Speicher und kann eine effizientere Zuteilung der GPU-Kerne ermöglichen.

Standardmäßig nutzt transformers für gewöhnliche Attention-Layer das SDPA von torch. Laut dem Beitrag stammten viele der kleinen Operationen von den nicht standardmäßigen Gated-Delta-Network-Layern in Qwen. Für diese Layer gibt es zwei Kernel, die viele Operationen fusionieren: causal-conv1d und flash-linear-attention. Die transformers-Implementierung von Qwen führt selbst auf, welche Kernel sie unterstützt.

Es gibt zwei Wege, sie zu beziehen: selbst installieren, etwa per pip install, per git clone samt Kompilierung oder über ein vorkompiliertes Paket, oder vorkompilierte Kernel aus dem Kernel Hub von Hugging Face holen. Databricks empfiehlt grundsätzlich den ersten Weg, kombiniert mit seiner Docker-Container-Unterstützung, damit die Kernel schon beim Build installiert werden und Jobs schneller starten. Das Unternehmen warnt, dass im Kernel Hub womöglich kein Kernel zu den eigenen torch- und CUDA-Versionen passt. In diesem Fall greift das transformers-Modell stillschweigend auf einen nicht optimierten Kernel mit niedriger Leistung zurück.

Nach der Installation von flash-linear-attention und causal-conv1d verbesserte sich die Gesamtzeit um den Faktor 3,7 auf 2 Stunden 5 Minuten. Die Gesamtzahl der Operationen sank um den Faktor 15. Im anschließenden Trace benötigte die CPU-Seite 3,9 Sekunden bei 118.316 Aufrufen, die GPU-Seite 3,9 Sekunden bei 18.060 Aufrufen, beide Seiten waren nun also im Gleichgewicht. Unter den häufigsten GPU-Operationen tauchen neue Einträge wie flash_attention_3 und chunk_gated_delta_rule_fwd auf.

Was das für Teams bedeutet, die große Modelle feinabstimmen

Die zentrale Botschaft des Beitrags: Eine hohe GPU-Auslastung belegt kein effizientes Training. Databricks sagt, AI Runtime erlaube es Nutzern, Leistungsprobleme mit einem normalen torch-Trace eines LLM zu diagnostizieren, und verweist auf Beispielcode, den Leser auf der Plattform ausprobieren können. Der gesamte Weg von 2 Stunden 5 Minuten bis zu den finalen 1 Stunde 17 Minuten stammt laut der Zusammenfassung des Unternehmens aus weiteren Verbesserungen an der Datenpipeline, die der hier vorliegende Auszug des Beitrags nicht im Detail beschreibt.

Für Teams, die hunderte bis tausende GPU-Stunden für Fine-Tuning aufwenden, lautet die praktische Lehre: prüfen, welche Kernel die Modellbibliothek unterstützt, sicherstellen, dass sie für die exakt verwendeten torch- und CUDA-Versionen tatsächlich installiert sind, und auf stille Fallbacks achten. Die Zahlen sind eigene Messungen von Databricks auf einem festen Datensatz mit einem bestimmten Setup, bei anderen Modellen und Konfigurationen werden die Ergebnisse abweichen.