Das Training großer KI-Modelle (wie Transformer und Sprachmodelle) ist zu einem unverzichtbaren Bestandteil des KI-Bereichs geworden, steht aber vor hohen Rechenkosten, hohem Speicherverbrauch und hohem Energiebedarf. OpenAIs GPT-3 beispielsweise verfügt über 175 Milliarden Parameter und benötigt mehrere Wochen GPU-Training. Dieser enorme Ressourcenbedarf schränkt die Anwendung dieser Technologie auf Organisationen mit umfangreichen Rechenressourcen ein und verstärkt die Bedenken hinsichtlich Energieeffizienz und Umweltbelastung. Die Bewältigung dieser Herausforderungen ist entscheidend, um eine breitere Verfügbarkeit und Nachhaltigkeit der KI-Entwicklung zu gewährleisten.
Traditionelle Trainingsmethoden sind ineffizient und erfordern innovative Lösungen
Die Ineffizienz des Trainings großer Modelle beruht hauptsächlich auf ihrer Abhängigkeit von dichten Matrizen, die viel Speicher und Rechenleistung benötigen. Moderne GPUs unterstützen nur begrenzt optimierte Operationen mit niedriger Genauigkeit oder niedrigem Rang, was diese Anforderungen weiter verschärft. Obwohl einige Methoden wie Matrixzerlegung und heuristische Rangreduzierung vorgeschlagen wurden, um diese Probleme zu lindern, sind sie in der Praxis immer noch eingeschränkt. GaLore beispielsweise kann im Single-Batch-Setting trainieren, weist aber unrealistisch hohe Laufzeitkosten auf. Ähnlich verhält es sich mit LTE (Low-Rank Adapters), das bei großen Aufgaben Konvergenzprobleme aufweist. Es fehlt derzeit an einer Methode, die gleichzeitig den Speicherverbrauch, die Rechenkosten und die Trainingszeit reduziert, ohne die Leistung zu beeinträchtigen, was den Bedarf an innovativen Lösungen dringlicher macht.
CoMERA-Framework: Effizientes Training durch adaptive Tensoroptimierung
Forscher der Universität Albany (State University of New York), der Universität Kalifornien in Santa Barbara, Amazon Alexa AI und Meta haben ein neues Framework namens CoMERA (Computing-and Memory-Efficient training method via Rank-Adaptive tensor optimization) vorgestellt. Dieses Framework kombiniert Speichereffizienz und Rechengeschwindigkeit durch adaptive Rang-Tensor-Kompressionstechniken. Im Gegensatz zu traditionellen Methoden, die sich nur auf die Kompression konzentrieren, verwendet CoMERA eine Multi-Objective-Optimierung, um Kompressionsrate und Modellgenauigkeit auszubalancieren. Es nutzt tensorisierte Einbettungen und fortschrittliche Tensor-Netzwerk-Kontraktionen, um die GPU-Auslastung zu optimieren, wodurch die Laufzeitkosten reduziert und gleichzeitig eine hohe Leistung beibehalten wird. Das Framework führt außerdem CUDA-Graphen ein, um die Kernel-Startverzögerung während der GPU-Operationen zu minimieren, einen Hauptengpass bei traditionellen Tensor-Kompressionsmethoden.
CoMERA basiert auf einer adaptiven Tensor-Darstellung, die es den Modellschichten ermöglicht, ihren Rang dynamisch an die Ressourcenbeschränkungen anzupassen. Durch die Modifikation des Tensor-Ranges kann das Framework Kompression erreichen, ohne die Integrität der neuronalen Netzwerk-Operationen zu beeinträchtigen. Diese dynamische Optimierung erfolgt über einen zweistufigen Trainingsprozess:
Frühphase: Fokus auf stabile Konvergenz.
Spätphase: Feinabstimmung des Ranges, um spezifische Kompressionsziele zu erreichen.
In einem Sechs-Encoder-Transformer-Modell erreichte CoMERA in der Frühphase eine Kompressionsrate von bis zu 43-fach, während die Kompressionsrate in der späteren Optimierungsphase sogar 361-fach betrug. Im Vergleich zu GaLore reduzierte es den Speicherverbrauch um das Neunfache und beschleunigte das Training pro Runde um das Zwei- bis Dreifache.
Mehrere Testergebnisse zeigen die überragende Leistung von CoMERA
Bei der Anwendung auf ein auf dem MNLI-Datensatz trainiertes Transformer-Modell reduzierte CoMERA die Modellgröße von 256 MB auf nur noch 3,2 MB, während die Genauigkeit erhalten blieb. In großen Empfehlungssystemen wie DLRM komprimierte CoMERA das Modell um das 99-Fache und reduzierte den Spitzen-Speicherverbrauch um das Siebenfache. Das Framework zeigte auch bei der Vorab-Schulung von CodeBERT (einem domänenspezifischen großen Sprachmodell) eine hervorragende Leistung mit einer Gesamtkompressionsrate von 4,23-fach und einer Verdoppelung der Geschwindigkeit in einigen Trainingsphasen. Diese Ergebnisse unterstreichen seine Fähigkeit, verschiedene Aufgaben und Architekturen zu bewältigen und erweitern seine Anwendbarkeit in verschiedenen Bereichen.
Zusammenfassung der wichtigsten Vorteile des CoMERA-Frameworks
Die wichtigsten Schlussfolgerungen dieser Studie lauten:
CoMERA erreicht eine Kompressionsrate von bis zu 361-fach für einzelne Schichten und 99-fach für das gesamte Modell, wodurch der Speicher- und Speicherbedarf erheblich reduziert wird.
Das Framework verkürzt die Trainingszeit pro Runde für Transformer und Empfehlungssysteme um das Zwei- bis Dreifache und spart so Rechenressourcen und Zeit.
Durch die Verwendung tensorisierter Darstellungen und CUDA-Graphen reduziert CoMERA den Spitzen-Speicherverbrauch um das Siebenfache, wodurch das Training auf kleineren GPUs ermöglicht wird.
Die Methode von CoMERA unterstützt verschiedene Architekturen, darunter Transformer und große Sprachmodelle, und behält gleichzeitig die Genauigkeit bei oder verbessert sie.
Durch die Reduzierung des Energie- und Ressourcenbedarfs für das Training trägt CoMERA zu einer nachhaltigeren KI-Praxis bei und ermöglicht es einem breiteren Publikum, Spitzentechnologien zu nutzen.