COAT: Komprimierung von Optimizer-Zuständen und Aktivierungen für speicher-effizientes FP8-Training

COAT: Compressing Optimizer states and Activation for Memory-Efficient FP8 Training

October 25, 2024

Autoren: Haocheng Xi, Han Cai, Ligeng Zhu, Yao Lu, Kurt Keutzer, Jianfei Chen, Song Han

cs.AI

Zusammenfassung

FP8-Training hat sich als vielversprechende Methode zur Verbesserung der Trainingseffizienz herausgestellt. Bestehende Frameworks beschleunigen das Training, indem sie FP8-Berechnungen auf lineare Schichten anwenden, während Optimizer-Zustände und Aktivierungen in höherer Präzision belassen werden, was jedoch nicht vollständig die Speicherauslastung optimiert. Dieses Paper stellt COAT (Compressing Optimizer States and Activations for FP8 Training) vor, ein neuartiges FP8-Trainingsframework, das darauf abzielt, den Speicherbedarf beim Training großer Modelle signifikant zu reduzieren. COAT begegnet aktuellen Einschränkungen durch zwei Schlüsselinnovationen: (1) Dynamische Bereichserweiterung, die Optimizer-Zustandsverteilungen enger an den FP8-Repräsentationsbereich anpasst, wodurch die Quantisierungsfehler reduziert werden, und (2) Gemischte Granularitätsaktivierungsquantisierung, die die Aktivierungsspeicherung mithilfe einer Kombination von pro-Tensor- und pro-Gruppen-Quantisierungsstrategien optimiert. Experimente zeigen, dass COAT den Gesamtspeicherbedarf für das Training um das 1,54-fache im Vergleich zu BF16 effektiv reduziert, während nahezu verlustfreie Leistung bei verschiedenen Aufgaben wie dem Vorabtraining und Feintuning großer Sprachmodelle sowie dem Training von Vision-Sprachmodellen erzielt wird. COAT erreicht auch eine 1,43-fache Beschleunigung des Gesamttrainings im Vergleich zu BF16 und liegt dabei auf Augenhöhe oder übertrifft die Beschleunigung von TransformerEngine. COAT ermöglicht effizientes Training großer Modelle mit allen Parametern auf weniger GPUs und erleichtert die Verdopplung der Batch-Größe in verteilten Trainingseinstellungen, was eine praktische Lösung für das Skalieren des Trainings großer Modelle darstellt. Der Code ist verfügbar unter https://github.com/NVlabs/COAT.

English

FP8 training has emerged as a promising method for improving training efficiency. Existing frameworks accelerate training by applying FP8 computation to linear layers while leaving optimizer states and activations in higher precision, which fails to fully optimize memory usage. This paper introduces COAT (Compressing Optimizer States and Activations for FP8 Training), a novel FP8 training framework designed to significantly reduce memory footprint when training large models. COAT addresses current limitations through two key innovations: (1) Dynamic Range Expansion, which aligns optimizer state distributions more closely with the FP8 representation range, thereby reducing quantization error, and (2) Mixed-Granularity Activation Quantization, which optimizes activation memory using a combination of per-tensor and per-group quantization strategies. Experiments demonstrate that COAT effectively reduces end-to-end training memory footprint by 1.54x compared to BF16 while achieving nearly lossless performance across various tasks, such as Large Language Model pretraining and fine-tuning and Vision Language Model training. COAT also achieves a 1.43x end-to-end training speedup compared to BF16, performing on par with or surpassing TransformerEngine's speedup. COAT enables efficient full-parameter training of large models on fewer GPUs, and facilitates doubling the batch size in distributed training settings, providing a practical solution for scaling large-scale model training. The code is available at https://github.com/NVlabs/COAT.