Effizientes KI-Training mit COAT: Speicheroptimierung durch FP8-Kompression

Das Training großer KI-Modelle stellt hohe Anforderungen an die verfügbaren Ressourcen, insbesondere den Speicherbedarf. Die gängige Praxis, Berechnungen in FP16 oder BF16 durchzuführen, während Optimiererzustände und Aktivierungen in höherer Präzision gehalten werden, schöpft das Optimierungspotenzial nicht voll aus. Eine neue Methode namens COAT (Compressing Optimizer States and Activations for FP8 Training) verspricht hier Abhilfe, indem sie die Speichernutzung durch FP8-Kompression deutlich reduziert.

COAT: Ein innovativer Ansatz zur FP8-Kompression

COAT verfolgt einen neuartigen Ansatz, um den Speicherbedarf beim Training großer KI-Modelle zu minimieren. Im Gegensatz zu bisherigen Frameworks, die FP8-Berechnungen nur auf lineare Schichten anwenden, komprimiert COAT sowohl die Optimiererzustände als auch die Aktivierungen in das FP8-Format. Dies ermöglicht eine deutlich effizientere Speichernutzung und ebnet den Weg für das Training noch größerer Modelle.

Zwei Kerninnovationen zeichnen COAT aus:

1. Dynamische Bereichserweiterung (Dynamic Range Expansion): Diese Technik passt die Verteilung der Optimiererzustände an den Darstellungsbereich von FP8 an. Dadurch wird der Quantisierungsfehler reduziert und die Genauigkeit des Trainings verbessert.

2. Gemischtgranulare Aktivierungsquantisierung (Mixed-Granularity Activation Quantization): COAT optimiert den Speicherbedarf der Aktivierungen durch eine Kombination aus Per-Tensor- und Per-Gruppen-Quantisierungsstrategien. Diese flexible Herangehensweise ermöglicht eine fein abgestimmte Anpassung an die jeweiligen Anforderungen des Modells.

Beeindruckende Ergebnisse in der Praxis

Experimentelle Ergebnisse belegen die Effektivität von COAT. Im Vergleich zu BF16 konnte der Speicherbedarf beim End-to-End-Training um das 1,54-fache reduziert werden, ohne signifikante Leistungseinbußen hinzunehmen. Dies gilt für verschiedene Aufgaben, darunter das Vortraining und Finetuning von großen Sprachmodellen sowie das Training von Vision-Language-Modellen.

Darüber hinaus erzielt COAT eine 1,43-fache Beschleunigung des End-to-End-Trainings im Vergleich zu BF16 und liegt damit auf Augenhöhe oder sogar über der Leistungssteigerung von TransformerEngine. COAT ermöglicht somit das effiziente Training großer Modelle mit allen Parametern auf weniger GPUs und erlaubt die Verdoppelung der Batch-Größe in verteilten Trainingsumgebungen.

Ausblick und Bedeutung für das KI-Training

COAT stellt einen wichtigen Fortschritt im Bereich des effizienten KI-Trainings dar. Die signifikante Reduzierung des Speicherbedarfs und die gleichzeitige Leistungssteigerung eröffnen neue Möglichkeiten für das Training immer größerer und komplexerer Modelle. Durch die Skalierbarkeit des Trainings auf verteilten Systemen trägt COAT dazu bei, die Grenzen des Machbaren im Bereich der Künstlichen Intelligenz zu erweitern. Der Code für COAT ist öffentlich zugänglich und steht der Forschungsgemeinschaft zur Verfügung, um die Weiterentwicklung und Anwendung dieser vielversprechenden Technologie zu fördern.

Bibliographie - Xi, H., Cai, H., Zhu, L., Lu, Y., Keutzer, K., Chen, J., & Han, S. (2024). COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training. arXiv preprint arXiv:2410.19313. - OpenReview. (n.d.). Retrieved from https://openreview.net/pdf/c389ed5ff4aa2680725a26219fcd1150478ac2c7.pdf - Rajbhandari, S., Rasley, J., Ruwase, O., & He, Y. (2020). DeepSpeed: System optimizations enable training deep learning models with over 100 billion parameters. Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 1252–1260. - Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., & Catanzaro, B. (2019). Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053. - Anonymous Authors. (2023). Title of Paper. Advances in Neural Information Processing Systems, 36.