Speichereffizientes Training von Sprachmodellen mit APOLLO

Speichereffizientes Training großer Sprachmodelle mit APOLLO

Das Training großer Sprachmodelle (LLMs) ist bekanntlich speicherintensiv, insbesondere bei der Verwendung des gängigen Optimierers AdamW. Dieser hohe Speicherbedarf erfordert den Einsatz von mehr oder leistungsstärkeren GPUs oder die Reduzierung der Batch-Größen, was die Skalierbarkeit und den Durchsatz des Trainings einschränkt. Um dem entgegenzuwirken, wurden verschiedene speichereffiziente Optimierer entwickelt, die den Speicherverbrauch des Optimierers reduzieren sollen. Diese stehen jedoch vor Herausforderungen: (i) Abhängigkeit von rechenintensiven SVD-Operationen; (ii) erhebliche Leistungseinbußen im Vergleich zu AdamW; und (iii) immer noch erheblicher Overhead im Optimizerspeicher, um eine wettbewerbsfähige Leistung zu gewährleisten.

Eine neue Forschungsarbeit stellt einen vielversprechenden Ansatz zur Lösung dieser Probleme vor: APOLLO (Approximated Gradient Scaling for Memory-Efficient LLM Optimization). Die Kernidee von APOLLO basiert auf der Erkenntnis, dass die Lernraten-Anpassungsregel von AdamW effektiv als strukturierte Lernratenaktualisierung vergröbert werden kann. APOLLO nutzt diese Erkenntnis, indem es die Skalierung der Lernrate mithilfe eines zusätzlichen Optimiererzustands mit niedrigem Rang approximiert, der auf einer reinen Zufallsprojektion basiert. Diese strukturierte Aktualisierungsregel für die Lernrate macht APOLLO sehr tolerant gegenüber weiteren Speicherreduzierungen, während gleichzeitig eine vergleichbare Pre-Training-Leistung erzielt wird.

APOLLO-Mini: SGD-ähnlicher Speicherverbrauch bei AdamW-Performance

Besonders hervorzuheben ist die Rank-1-Variante von APOLLO, genannt APOLLO-Mini. Diese Variante erreicht eine überlegene Pre-Training-Leistung im Vergleich zu AdamW, während der Speicherbedarf auf dem Niveau von SGD liegt. Dies ist ein bedeutender Fortschritt, da SGD zwar speichereffizient ist, aber oft nicht die gleiche Leistung wie AdamW erreicht. APOLLO-Mini kombiniert somit die Vorteile beider Ansätze.

Vorteile von APOLLO

Umfangreiche Experimente zeigen, dass die APOLLO-Serie eine Leistung erzielt, die mit AdamW vergleichbar oder sogar besser ist, während gleichzeitig größere Speichereinsparungen erzielt werden, indem die Optimierungszustände von AdamW nahezu eliminiert werden. Diese Einsparungen bieten erhebliche Vorteile auf Systemebene:

- Erhöhter Durchsatz: 3x höherer Durchsatz auf einem 8xA100-80GB-Setup im Vergleich zu AdamW durch Unterstützung von 4x größeren Batch-Größen. - Verbesserte Modellskalierbarkeit: Pre-Training von LLaMA-13B mit naivem DDP auf A100-80GB-GPUs ohne systemseitige Optimierungen. - Pre-Training für Low-End-GPUs: Pre-Training von LLaMA-7B auf einer einzelnen GPU mit weniger als 12 GB Speicher durch Gewichtsquantisierung.

Ausblick

APOLLO stellt einen vielversprechenden Ansatz für das speichereffiziente Training von LLMs dar. Durch die Approximation der Lernratenskalierung ermöglicht APOLLO eine deutliche Reduzierung des Speicherbedarfs, ohne die Leistung zu beeinträchtigen. Die Möglichkeit, größere Modelle auf weniger leistungsstarker Hardware zu trainieren, eröffnet neue Möglichkeiten für die Forschung und Entwicklung im Bereich der LLMs und könnte die Demokratisierung des Zugangs zu diesen mächtigen Modellen fördern.

Bibliographie Zhu, H., Zhang, Z., Cong, W., Liu, X., Park, S., Chandra, V., Long, B., Pan, D. Z., Wang, Z., & Lee, J. (2024). APOLLO: SGD-like Memory, AdamW-level Performance. arXiv preprint arXiv:2412.05270. Ma, X. (2022). Apollo: An adaptive parameter-wised diagonal quasi-newton method for nonconvex stochastic optimization. arXiv preprint arXiv:2009.13586. Kumar, A., Shen, R., Bubeck, S., & Gunasekar, S. (2022). How to fine-tune vision models with sgd. arXiv preprint arXiv:2211.09359. Liu, B., & Qian, Y. (2024). Memory-efficient training for deep speaker embedding learning in speaker verification. arXiv preprint arXiv:2412.01195v1. https://openreview.net/forum?id=WwKv20NrsfB https://www.reddit.com/r/MachineLearning/comments/16cgukc/rd_hey_lomo_paper_authors_does_sgd_have_optimizer/ https://openreview.net/pdf?id=rnFOPhTMB0Y https://github.com/XuezheMax/apollo/issues/1 https://www.chatpaper.com/chatpaper/zh-CN?id=5&date=1733673600&page=1 https://www.researchgate.net/publication/372201106_The_performance_analysis_of_Adam_and_SGD_in_image_classification_and_generation_tasks https://github.com/XuezheMax/apollo https://huggingface.co/papers/2211.09359