Wissensdestillation mit deutlich weniger GPU-Speicher
💡 Quellen-Hinweis: Dieser Beitrag ist eine redaktionelle Zusammenfassung basierend auf einem Originalbericht von Hugging Face Blog.
Wissensdestillation überträgt Fähigkeiten eines großen Teacher-Modells auf einen kleineren Student. Das spart später Rechenleistung, ist während des Trainings jedoch ausgesprochen teuer: Üblicherweise müssen beide Modelle gleichzeitig geladen werden und für jedes Token entstehen Wahrscheinlichkeitswerte über das gesamte Vokabular. Multiverse Computing beschreibt nun zwei Systemänderungen, die diesen Engpass deutlich verkleinern und auch lange Kontexte auf wesentlich weniger GPUs trainierbar machen sollen.
Teacher-Ausgaben einmalig zwischenspeichern
Der erste Baustein ist eine Offline-Destillation. Statt den Teacher bei jedem Trainingsschritt erneut auszuführen, werden dessen 100 wahrscheinlichste Logits pro Token einmal berechnet und gespeichert. Danach arbeitet das Student-Training nur noch mit diesem Cache. Der Teacher belegt dabei keinen Speicher mehr und dieselben Daten lassen sich für mehrere Experimente wiederverwenden. Das ist besonders relevant bei großen Vokabularen: Für gpt-oss-120b nennt der Beitrag bei 32.000 Token und Batchgröße vier allein rund 50 GB für einen einzigen Teacher-Wahrscheinlichkeitstensor.
KL-Verlust wird in Blöcken berechnet
Der zweite Baustein betrifft die KL-Divergenz. Eine dichte Standardberechnung erzeugt eine riesige Matrix aus Sequenzlänge und Vokabular. Die vorgestellte fused chunked KL loss verbindet deshalb die Output-Projektion des Modells direkt mit der Verlustberechnung. Sie verarbeitet jeweils nur einen Sequenzblock, verwirft dessen Logits anschließend und berechnet ihn im Rückwärtslauf erneut. Das kostet zusätzliche Rechenarbeit, verhindert aber die große Speicherspitze. Eine weniger aggressive Forward-Chunked-Variante ist bei kürzeren Kontexten schneller, hält die vollständigen Student-Logits jedoch weiterhin für den Backward-Pass vor.
Messwerte und praktische Bedeutung
Auf einer einzelnen H200 sank der Spitzenverbrauch bei einem 8K-Test von 102,8 GB für Online-Destillation auf 58,3 GB mit der vollständig blockweisen Methode. In einem isolierten 32K-Test fiel der Bedarf laut Beitrag sogar von 85,2 auf 5,45 GiB. Bei der Destillation eines GPT-OSS-20B-Modells mit 32.768 Token konnte das Setup von vier GPU-Knoten auf einen reduziert werden; die Schrittzeit sank von 57 auf 12,23 Sekunden. Die Verlustkurven der getesteten Verfahren lagen nahezu übereinander. Der resultierende Student mit etwa 3,2 Milliarden Parametern bewahrte einen großen Teil der Leistung eines Llama-3.1-8B-Teachers. Für KI-Teams bedeutet das vor allem günstigere Iterationen bei Modellkompression und Long-Context-Anpassung. Die Implementierung ist offen verfügbar, sollte aber mit eigenen Modellen und Qualitätsmetriken validiert werden.
Originalquelle: Hugging Face Blog – Making Knowledge Distillation Cheap Enough to Run at Scale
