YaFSDP – ein LLM-Trainingstool, das die GPU-Nutzung um 20 % senkt – ist jetzt erhältlich

Die Entwicklung großer Sprachmodelle erfordert erhebliche Investitionen in Zeit und GPU-Ressourcen, was sich direkt in hohen Kosten niederschlägt. Je größer das Modell, desto ausgeprägter werden diese Herausforderungen.

Yandex hat kürzlich eine neue Lösung vorgestellt: YaFSDP , ein Open-Source-Tool, das verspricht, das LLM-Training zu revolutionieren, indem es den GPU-Ressourcenverbrauch und die Trainingszeit deutlich reduziert. In einem Vortrainingsszenario mit einem Modell mit 70 Milliarden Parametern können durch die Verwendung von YaFSDP die Ressourcen von etwa 150 GPUs eingespart werden. Dies entspricht potenziellen monatlichen Einsparungen von etwa 0,5 bis 1,5 Millionen US-Dollar, je nach Anbieter oder Plattform des virtuellen GPUs.

Yandex hat YaFSDP öffentlich zugänglich gemacht auf GitHub .

Die Herausforderung der verteilten LLM-Ausbildung

Das Trainieren von LLMs über mehrere GPUs hinweg erfordert komplexe Vorgänge, die zu Ineffizienzen und hohem Speicherverbrauch führen. Eines der Hauptprobleme ist die Notwendigkeit, riesige Datenmengen zwischen GPUs zu senden und zu empfangen. Beispielsweise müssen bei einem typischen all_reduce-Vorgang doppelt so viele Gradientendaten wie Netzwerkparameter übermittelt werden. Im Fall eines Llama 70B-Modells bedeutet dies die Übertragung von 280 GB Daten pro Iteration.

Darüber hinaus werden Gewichte, Gradienten und Optimiererzustände über GPUs hinweg dupliziert, was zu einer enormen Speicherlast führt. Das Llama 70B-Modell und der Adam-Optimierer benötigen über 1 TB Speicher und überschreiten damit die typische Speicherkapazität von 80 GB der meisten GPUs bei weitem. Diese Redundanz verlangsamt den Trainingsprozess erheblich und macht es oft unpraktisch, selbst mittelgroße Modelle in den GPU-Speicher zu packen.

Einführung in YaFSDP

YaFSDP von Yandex bietet eine äußerst effektive Lösung für diese Herausforderungen. Durch die Fokussierung auf die Optimierung des Speicherverbrauchs und die Beseitigung von Kommunikationsengpässen verbessert YaFSDP die Effizienz des LLM-Trainings. Es funktioniert durch die Aufteilung von Schichten anstelle einzelner Parameter, wodurch eine effiziente Kommunikation aufrechterhalten und redundante Vorgänge vermieden werden. Darüber hinaus weist YaFSDP Puffer für alle erforderlichen Daten vorab zu und stellt so sicher, dass der Torch-Allocator keine Ineffizienzen verursacht.

YaFSDP funktioniert, indem zwei Puffer für Zwischengewichte und Gradienten verwendet werden, wobei die ungeraden Schichten einen Puffer und die geraden Schichten den anderen Puffer verwenden.

Die Gewichte aus verschiedenen Schichten werden im selben Speicher gespeichert. Wenn die Schichten dieselbe Struktur haben, sind sie immer identisch. Es ist wichtig sicherzustellen, dass der Puffer die Gewichte für Schicht X enthält, wenn Sie Schicht X benötigen. Alle Parameter werden im entsprechenden Speicherblock innerhalb des Puffers gespeichert.

Speicherverbrauch

Während des Trainings sind Gewichte, Gradienten, Optimiererzustände, Puffer und Aktivierungen die wichtigsten Speicherverbraucher. YaFSDP reduziert den Speicherverbrauch erheblich, indem es die Art und Weise optimiert, wie diese Elemente gespeichert und abgerufen werden.

• Gewichte, Gradienten und Optimiererzustände : Diese hängen von der Anzahl der Prozesse ab und ihr Speicherverbrauch tendiert dazu, sich Null zu nähern, wenn die Anzahl der Prozesse zunimmt. Durch die Aufteilung dieser Komponenten auf mehrere GPUs minimiert YaFSDP die Duplizierung und reduziert so den Speicherverbrauch.

• Puffer verbrauchen eine konstante Menge an Speicher und speichern Zwischenwerte während der Berechnungen.

• Aktivierungen hängen von der Modellgröße und der Anzahl der pro GPU verarbeiteten Token ab.

Aktivierungs-Checkpointing

Aktivierungs-Checkpointing ist eine Technik, die nur die notwendigen Aktivierungen während des Vorwärtsdurchlaufs speichert und sie während des Rückwärtsdurchlaufs neu berechnet. Dies reduziert den Speicherbedarf erheblich, da nur die wesentlichen Daten gespeichert werden. Beispielsweise kann beim Trainieren eines Llama 2 70B-Modells mit einer Batchgröße von 8192 Token der Aktivierungsspeicher von über 110 GB auf nur 5 GB reduziert werden.

Dieser Ansatz führt jedoch zu einem zusätzlichen Rechenaufwand, den YaFSDP vermeiden kann, indem die Aktivierungs-Checkpoints für einige Schichten nicht verwendet werden, was durch Speicheroptimierung möglich ist.

Kommunikationsoptimierung

YaFSDP verbessert die GPU-Kommunikationseffizienz, indem es sicherstellt, dass Daten nur bei Bedarf übertragen werden und indem die Kommunikation mit der Berechnung überlappt wird. Es nutzt CUDA-Streams, um gleichzeitige Berechnungen und Kommunikationen effektiv zu verwalten.

Das Tool verwendet zwei Streams: einen Berechnungsstream und einen Kommunikationsstream. Ereignisse synchronisieren diese Streams und stellen sicher, dass Vorgänge in der richtigen Reihenfolge ausgeführt werden, ohne dass es zu Deadlocks kommt.

Der Vorwärtsdurchlauf auf der dritten Ebene beginnt erst, wenn der all_gather-Vorgang abgeschlossen ist (Bedingung 1). Ebenso beginnt der all_gather-Vorgang auf der dritten Ebene erst, wenn der Vorwärtsdurchlauf auf der ersten Ebene, der denselben Puffer verwendet, abgeschlossen ist (Bedingung 2). Da es in diesem Schema keine Zyklen gibt, ist ein Deadlock unmöglich.

Experimentelle Ergebnisse und Leistungssteigerungen

Die Implementierung von YaFSDP hat bemerkenswerte Verbesserungen bei der Trainingseffizienz gezeigt. In einem Vortrainingsszenario mit einem Modell mit 70 Milliarden Parametern konnte YaFSDP die Ressourcen von etwa 150 GPUs einsparen. Dies führt zu erheblichen monatlichen Kosteneinsparungen, die je nach Anbieter oder Plattform des virtuellen GPUs zwischen 0,5 und 1,5 Millionen US-Dollar liegen.

YaFSDP reduziert die Trainingszeit um bis zu 26 % im Vergleich zu bestehenden Methoden wie FSDP und optimiert die Speichernutzung, wodurch es möglich wird, größere Modelle effizienter zu trainieren.

Yandex hat YaFSDP öffentlich zugänglich gemacht auf GitHub . ML-Ingenieure können dieses Tool nutzen, um die Effizienz ihrer LLM-Trainingsprozesse zu steigern. Durch die Open-Source-Veröffentlichung von YaFSDP möchte Yandex Innovation und Zusammenarbeit in der KI-Community fördern und es Entwicklern ermöglichen, Modelle schneller und kostengünstiger zu trainieren.

YaFSDP stellt einen bedeutenden Fortschritt im LLM-Training dar. Durch die Bewältigung der kritischen Herausforderungen des Speicherverbrauchs und der Kommunikationsineffizienzen wird ein schnelleres und effizienteres Training großer Sprachmodelle ermöglicht.