Microsoft Research bringt den ersten FP4-Genauigkeitsrahmen für das Training großer

Modelle heraus, die Effizienz wurde erheblich gesteigert
Am 29. Januar gab Microsoft Research die Einführung des ersten FP4-Genauigkeitsrahmens für das Training großer Modelle bekannt. Bei denselben Hyperparameter-Einstellungen kann dieser Rahmen ähnliche Trainingsergebnisse wie FP8 und BF16 erzielen. Dies bedeutet, dass der erforderliche Speicher- und Rechenaufwand verringert werden kann.
Modellgröße und Trainingseffekt
Die mit dieser Methode trainierten Modelle können eine Größe von bis zu 13 Milliarden Parametern erreichen, und die Anzahl der Trainingstoken liegt im Milliardenbereich. Obwohl derzeit FP4 durch FP8 simuliert wird, würde der Effekt mit echtem FP4 weiter verbessert.
FP4-Simulation und Trainingsmethoden
Da zum Zeitpunkt der Forschung keine native Unterstützung für FP4-Hardware verfügbar war, simulierte das Forschungsteam FP4 mithilfe von FP8 TensorCores. Genauer gesagt zeigte sich bei den LLaMA-Modellen mit 1,3B, 7B und 13B während des Trainings von 0 bis 10 Millionen Tokens, dass die Verlustkurve des FP4-Trainings nahezu identisch mit der von BF16 war.
Um FP4-Genauigkeitstraining zu ermöglichen, verwendete das Forschungsteam einen benutzerdefinierten FP4-Matrixmultiplikations-CUDA-Kernel (GeMM). Im Kernel werden zunächst FP4-A- und -B-Matrizen in FP16 in den gemeinsamen Speicher geladen und entsprechend transformiert. Anschließend wird die Blockmatrixmultiplikation mit FP4 durchgeführt und die Zwischenergebnisse mit FP16 reduziert, um eine Ausgabe-Matrix im FP16-Format zu erhalten.
Quantifizierungsstrategien und Hardware-Anpassung
Dieser Rahmen verwendet das E2M1 FP4-Format, bei dem 2 Bit für den Exponenten, 1 Bit für die Mantisse und 1 Bit für das Vorzeichen verwendet werden, insgesamt also 4 Bit. Dieses Format wurde gewählt, um mit dem Design der Quantifizierungseinheiten der aktuellen ML-Beschleunigerchips kompatibel zu sein. Für die Gewichtsmatrix W und die Aktivierungsmatrix A wurden unterschiedliche Quantifizierungsstrategien mit unterschiedlicher Granularität angewendet, um den Beschleunigungseffekt von FP4 bei der Matrixmultiplikation zu maximieren.
Innovationen in der Vorwärts- und Rückwärtspropagation
Zu Beginn der Vorwärtspropagation des Modells quantifiziert der Rahmen gleichzeitig die Gewichtsmatrix W und die Eingabeaktivierungsmatrix A jeder linearen Schicht in FP4. Bei der Quantifizierung werden die Werte in den Matrizen skaliert und verschoben, um sie in den Bereich zu überführen, den FP4 darstellen kann, und anschließend auf den nächsten FP4-Disretisierungswert durch eine Lookup-Tabelle gerundet.
Im Rückpropagationsprozess hat das Forschungsteam eine neuartige differenzierbare Gradienten-Schätzmethode entwickelt, um sowohl die Berechnungseffizienz als auch die Genauigkeit der Gradienten zu gewährleisten. Darüber hinaus wurde eine Strategie zur „Peak-Dämpfung und -Kompensation von Ausreißern“ für Probleme mit Ausreißern in der Aktivierungsmatrix vorgeschlagen.
Mixed Precision Design und Anwendung
In einigen Bereichen verwendete das Forschungsteam ein Mixed Precision Design, indem FP8 für die Gradientenkommunikation und FP16 für den Speicher des Optimiererzustands verwendet wurde. In anderen Teilen des Systems, wie z.B. bei nicht-Matrix-Multiplikationsoperationen und Loss Scaling, wurde ebenfalls FP16 verwendet. Diese Designs verringern den Rechen- und Speicheraufwand so weit wie möglich, während die numerische Stabilität des Trainings gewährleistet bleibt.
Forschungsteam und Hintergrund
Dieser Rahmen wurde von Microsoft Research Asia und dem SIGMA-Team entwickelt, wobei alle Forscher chinesischer Herkunft sind. Der erste Autor, Ruizhe Wang, ist Doktorand an der Universität für Wissenschaft und Technologie in China und derzeit als Praktikant bei MSRA tätig, wobei sein Forschungsschwerpunkt in der Quantifizierung mit niedriger Präzision liegt. Professor Zhejun Cha, der Leiter des Forschungsdepartments an der USTC, war ebenfalls an diesem Projekt beteiligt. Die Korrespondenzautoren sind Peng Cheng, Senior Research Manager bei MSRA, und Yeyun Gong, Chief Research Manager. Der herausragende MSRA-Wissenschaftler Baining Guo war ebenfalls an diesem Projekt beteiligt.
Zusammenfassung
Der von Microsoft Research eingeführte FP4-Genauigkeitsrahmen für das Training großer Modelle hat durch innovative Quantifizierungsstrategien und Mixed-Precision-Designs ein effizienteres Modelltraining ermöglicht. Dieser Rahmen steigert nicht nur die Trainingseffizienz, sondern reduziert auch die erforderlichen Rechen- und Speicherressourcen erheblich und eröffnet neue Wege für das Training großer Modelle.
Papier-Link: arXiv
Referenz-Link: Twitter