Megatron-LM: Training von Sprachmodellen mit mehreren Milliarden Parametern unter Verwendung von Modellparallelismus

Neuere Arbeiten im Bereich der Sprachmodellierung zeigen, dass das Training großer Transformer-Modelle den Stand der Technik in Anwendungen der Natürlichen Sprachverarbeitung (Natural Language Processing, NLP) verbessert. Allerdings können sehr große Modelle aufgrund von Speicherbeschränkungen recht schwierig zu trainieren sein. In dieser Arbeit stellen wir unsere Techniken zur Ausbildung sehr großer Transformer-Modelle vor und implementieren einen einfachen, effizienten inneren Schichtparallellismus (intra-layer model parallel approach), der es ermöglicht, Transformer-Modelle mit Milliarden von Parametern zu trainieren. Unser Ansatz erfordert keinen neuen Compiler oder Bibliotheksänderungen, ist orthogonal und ergänzend zum Pipeline-Schichtparallellismus und kann vollständig durch die Einfügung einiger Kommunikationsoperationen in nativen PyTorch umgesetzt werden. Wir veranschaulichen diesen Ansatz durch das Konvergieren von Transformer-basierten Modellen mit bis zu 8,3 Milliarden Parametern unter Verwendung von 512 GPUs. Im Vergleich zu einer starken Einzels-GPU-Basislinie, die 39 TeraFLOPs aufrechterhält (was 30% des maximalen FLOPs entspricht), erreichen wir eine Skalierungseffizienz von 76% und halten 15,1 PetaFLOPs über die gesamte Anwendung hinweg aufrecht. Um zu demonstrieren, dass große Sprachmodelle den Stand der Technik weiter verbessern können, trainieren wir ein 8,3-Milliarden-Parameter-Transformer-Sprachmodell ähnlich dem GPT-2 sowie ein 3,9-Milliarden-Parameter-Modell ähnlich dem BERT. Wir zeigen, dass eine sorgfältige Beachtung der Platzierung der Schichtnormalisierung in BERT-artigen Modellen entscheidend für eine bessere Leistung bei wachsendem Modellumfang ist. Mit dem GPT-2-Modell erreichen wir Stand-der-Techik-Ergebnisse (State-of-the-Art, SOTA) auf den Datensätzen WikiText103 (Perplexität von 10,8 im Vergleich zur SOTA Perplexität von 15,8) und LAMBADA (Genauigkeit von 66,5% im Vergleich zur SOTA Genauigkeit von 63,2%). Unser BERT-Modell erreicht SOTA-Ergebnisse auf dem RACE-Datensatz (Genauigkeit von 90,9% im Vergleich zur SOTA Genauigkeit von 89,4%).