Wie Amazon Search M5 durch den Einsatz von AWS Trainium 30 % der LLM-Schulungskosten einsparte | Amazon Web Services

Amazon leistet seit Jahrzehnten Pionierarbeit und Innovationen im Bereich des maschinellen Lernens (ML) und bietet seinen Kunden damit wunderbare Erlebnisse. Seit jeher nutzt Amazon ML für verschiedene Anwendungsfälle wie Buchempfehlungen, Suche und Betrugserkennung. Ähnlich wie im Rest der Branche haben die Fortschritte bei der beschleunigten Hardware es Amazon-Teams ermöglicht, Modellarchitekturen mithilfe neuronaler Netze und Deep Learning (DL) zu entwickeln.

Das M5-Programm innerhalb von Amazon Search besitzt die Discovery-Lernstrategie für Amazon und erstellt groß angelegte Modelle für mehrere Sprachen, mehrere Standorte, mehrere Entitäten, mehrere Aufgaben und mehrere modale Bereiche wie Text, Bild und Video. Das M5-Programm stellt Hunderten von ML-Teams auf Amazon universelle Einbettungen und groß angelegte Basismodelle zur Verfügung und behält dabei strenge Kontrollen über die Kostenoptimierung bei. Um dies zu erreichen, evaluiert das M5-Team regelmäßig neue Techniken zur Kostensenkung.

Wie viele ML-Organisationen werden Beschleuniger hauptsächlich zur Beschleunigung des DL-Trainings und der Inferenz eingesetzt. Als AWS mit der ersten Version von speziell entwickelte Beschleuniger auf den Markt brachte AWS-Inferenz Im Jahr 2020 begann das M5-Team schnell damit Nutzen Sie sie, um Produktions-Workloads effizienter bereitzustellenDies spart Kosten und reduziert die Latenz. Letztes Jahr hat AWS seine gestartet AWS-Training Beschleuniger, die die Leistung pro Kosten für die Entwicklung und den Bau von DL-Modellen der nächsten Generation optimieren. In diesem Beitrag diskutieren wir, wie M5 die Kosten für das Training seiner Modelle um 30 % senken konnte, und teilen einige der Best Practices, die wir dabei gelernt haben.

Trainium-Instanzen

Mit den Fortschritten bei speziell entwickelten Beschleunigern bietet Amazon mit AWS Inferentia und Trainium auch überzeugende Beschleuniger an. Wie der Name schon sagt, sind diese Chips so optimiert, dass sie die Anforderungen von Inferenz- bzw. Trainings-Workloads übertreffen. Für das groß angelegte Training von Grundlagenmodellen, die Milliarden von Parametern erreichen, Trainium Trn1- und Trn1n-Instanzen sind aufgrund ihrer Eigenschaften die ideale Wahl. Trn1-Instanzen basieren auf dem neuesten Stand der Technik NeuronCore-v2und verfügen über reichlich Beschleuniger-Rechenleistung und Speicher. Trn1n-Instanzen können auch für eine größere Netzwerkbandbreite (1,600 Gbit/s) ausgewählt werden und eignen sich daher ideal für leistungsstarkes Training unter Berücksichtigung der Kostenoptimierung.

Um Beschleuniger nutzen zu können, benötigen Sie eine Softwareschicht, die sie unterstützt. Mit Trn- und Inf-Chips ist das AWS Neuron-SDK schaltet mit Hilfe von PyTorch XLA speziell für Amazon entwickelte Beschleuniger frei. PyTorch XLA konvertiert den Eager-Modus von PyTorch in eine graphbasierte Implementierung im Lazy-Modus. Diese Diagramme werden dann verwendet und weiter kompiliert, um sie mit dem Beschleuniger zu verwenden. Mit PyTorch Neuron (Teil des Neuron SDK) können PyTorch-Benutzer ihre Modelle mit wenigen Codezeilen auf Trainium NeuronCores trainieren.

Modell und Arbeitsaufwand

Das M5-Team schult und implementiert grundlegende Modelle und universelle Darstellungen, um verschiedene Teams bei Amazon dabei zu unterstützen, Freude zu bereiten Amazon.com Kunden. Ein solches Modell ist ein Text-Encoder-Modell, gefolgt von einem mehrschichtigen Perzeptron (MLP) mit expliziten oder impliziten Merkmalsinteraktionen, die durch die neuronale Netzwerkarchitektur mit Hunderten Millionen trainierbaren Parametern definiert werden. Dieses Modell wird auf Milliarden von Token trainiert und zur Generierung von Millionen von Einbettungen in einer Offline-Batch-Inferenzeinstellung verwendet. Diese Einbettungen sind Eingaben für einen kundenorientierten Tier-1-Amazon-Dienst.

Die Infrastruktur für die Produktionspipeline nutzt AWS-Charge mit Fair-Share-Warteschlangenstrategien, wobei ein EFA-fähiger trn1.32xlarge-Cluster mit mehreren Knoten als Recheneinheit für das Modelltraining verwendet wird. Funktional führt die Produktionspipeline ein inkrementelles Modelltraining, eine Bewertung des trainierten Modells und eine Offline-Batch-Inferenz für das trainierte Modell durch, wobei alles PyTorch als zugrunde liegende DL-Bibliothek verwendet.

Ziele

Unsere Kunden zu begeistern ist unser oberstes Gebot. Angesichts des kundenorientierten Charakters der Pipeline ist es von entscheidender Bedeutung, dass alle Service-Level-Agreements (SLAs) ohne Rückschritte eingehalten werden. Wir haben zwei entscheidende Akzeptanzkriterien identifiziert, um unsere bestehende GPU-Produktionspipeline anzupassen und auf Trainium umzustellen:

• Modellqualität – Die Qualität unserer Modelle wirkt sich direkt auf das Kundenerlebnis aus. Wir verlangen, dass der Unterschied in der Modellqualität zwischen GPU und Trainium weniger als 0.1 % beträgt.
• Trainingsdurchsatz – Wir trainieren unsere Modelle regelmäßig iterativ, um unseren Kunden das aktuellste Erlebnis zu bieten. Wir verlangen, dass die Modellkonvergenz innerhalb eines vordefinierten Zeitraums (z. B. 1 Woche) erreicht werden muss, um unsere Produktions-SLAs zu erfüllen.

In den folgenden Abschnitten teilen wir unseren Weg der Rückwärtsarbeit anhand dieser Kriterien und unsere Erkenntnisse zur Unterstützung von Produktions-Workloads im Amazon-Maßstab.

Schulungsskript

Bevor wir mit dem Modelltraining beginnen, müssen wir Änderungen am Trainingsskript vornehmen, um es XLA-kompatibel zu machen. Angesichts der Größe des Modells verwenden wir Distributed Data Parallel (DDP), um das Modell zu trainieren. Mit DDP können wir den Durchsatz des Modelltrainings erhöhen, indem wir die Anzahl der für die Ausführung des Modelltrainings verwendeten Maschinen erhöhen, ohne dass Codeänderungen erforderlich sind. Wir haben die Anweisungen in der befolgt Neuron PyTorch MLP-Trainings-Tutorial um XLA-spezifische Konstrukte in unsere Trainingsskripte einzufügen. Diese Codeänderungen sind einfach zu implementieren. Im Folgenden sind einige wichtige technische Erkenntnisse aus der Übung aufgeführt, die unseren Modelldurchsatz erheblich verbessert haben:

• Platzierung von xm.mark_step() - xm.mark_step() kompiliert die träge gesammelten Berechnungsdiagramme und führt sie aus. Aufrufen mark_step Zu häufiges Aufrufen führt zu einer größeren Anzahl kleiner Diagramme, während ein zu häufiges Aufrufen zu wenigen, aber großen Diagrammen führt. Abhängig von Ihrer Anwendung variieren der Durchsatz und die Implementierung Ihres Modelltrainings je nach Ihrer Platzierung xm.mark_step(). Unsere Umsetzung setzt einen xm.mark_step() nach einem Vorwärts- und Rückwärtsdurchlauf und einen nach dem Optimierungsschritt.
• Datenladeprogramm mit XLA-Multiprozessor-Geräteladeprogramm umschließen – Dies ist ein entscheidender Schritt, der leicht übersehen werden kann. Der Multiprozessor-Gerätelader torch_xla.distributed.parallel_loader.MpDeviceLoader Lädt Trainingsdaten auf jedes XLA-Gerät mit Optionen zum Vorladen und Überlappen des Datenladens mit Geräteläufen, um den Durchsatz zu verbessern. Der Gerätelader wird ebenfalls aufgerufen xm.mark_step() und ist daher in der Lage, Diagramme für das Laden von Daten vom Host auf das Gerät zu erstellen.

Zusammenstellung für Trainium

Traditionell umfasst der Modellentwicklungszyklus mit GPUs das Vornehmen von Änderungen am Modell oder Trainingsskript und deren direkte Ausführung auf dem GPU-Gerät. Beschleuniger wie Trainium, die XLA verwenden, erfordern einen zusätzlichen Schritt, bevor das Modelltraining auf dem Beschleuniger ausgeführt werden kann. XLA-Berechnungsdiagramme können erst ausgeführt werden, nachdem sie kompiliert wurden. Im Allgemeinen gibt es zwei Möglichkeiten, diese Kompilierung durchzuführen: Ahead of Time (AOT), bei dem Sie zunächst alle Diagramme verfolgen und kompilieren und sie dann ausführen, oder Just In Time (JIT), bei dem Diagramme verfolgt, kompiliert und so ausgeführt werden, wie sie sind angetroffen werden. Das Neuron SDK bietet beides sofort. Normalerweise wird zuerst die AOT-Kompilierung durchgeführt. Nach dieser Kompilierung werden dann Diagramme ausgeführt. Wenn neue Diagramme gefunden werden, ruft die Neuron-Laufzeit eine JIT-Kompilierung auf, bevor sie ausgeführt wird. Um die AOT-Kompilierung durchzuführen, bietet das Neuron SDK neuron_parallel_compile, ein Kompilierungsdienstprogramm, das Diagramme aus einem Testlauf des Trainingsskripts extrahiert und eine parallele AOT-Kompilierung durchführt.

Ein wichtiger Aspekt der AOT-Kompilierung besteht darin, sicherzustellen, dass im Laufe des Trainings keine neuen Berechnungsdiagramme erstellt werden. Eine Quelle für neue Berechnungsdiagramme (und damit für Neukompilierungen) sind dynamische Formen der Trainingsstapel während des Modelltrainings. Wir haben festgestellt, dass die Verwendung statischer Formen und Stapel fester Größe die Kompilierung der Trainingszeit überflüssig macht und den Trainingsdurchsatz erheblich verbessert, ohne dass sich dies auf die Modellgenauigkeit auswirkt. Durch die Durchsetzung solcher Einschränkungen beim Training haben wir beobachtet, dass nur 4–5 Schritte des Modelltrainings, ein Schritt der Modellvalidierung und ein einmaliges Checkpointing des Modells erforderlich sind, um alle Diagramme während der AOT-Kompilierung zu verfolgen. Es ist wichtig zu beachten, dass das Neuron SDK ständig weiterentwickelt wird und in Zukunft auch dynamische Formen unterstützen wird.

Darüber hinaus werden die erstellten Diagramme im gespeichert Permanenter Neuronen-Cache auf der Festplatte oder in einem Amazon Simple Storage-Service (Amazon S3) Eimer. Dies ist besonders nützlich für Produktions-Workloads, bei denen sich die Modellarchitektur und die Trainingskonfiguration nicht ändern. Daher fällt der Kompilierungsaufwand nur einmal an. Die Verwendung des Caches ist so einfach wie das Setzen eines Umgebungsflags:

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Der Neuron-Compiler stellt auch drei bereit Optimierungsoptionen auf Compilerebene (O1, O2, O3), um Kompilierungszeit und Modelllaufdurchsatz auszugleichen. O1 ermöglicht Kernoptimierungen im Rechendiagramm und minimiert die Kompilierungszeit, O3 bietet einen verbesserten Modelllaufdurchsatz auf Kosten einer höheren Kompilierungszeit und O2 (Standardoption) ist ein Gleichgewicht zwischen beiden. Für unseren Anwendungsfall verwendeten wir die O1-Optimierung und beobachteten eine Reduzierung der Kompilierungszeit um 86 % ohne Änderung der Modellgenauigkeitsmetriken, während wir im Vergleich zur Standardoptimierung (O5) eine Reduzierung des Durchsatzes um etwa 7–2 % beobachteten. Je nach Anwendungsfall können Sie unterschiedliche Optimierungsstufen wählen.

Zusammenfassend haben wir die folgenden Flags für die Kompilierung verwendet:

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

Checkpoint-Kompatibilität

Wenn die Kompilierung erfolgreich abgeschlossen ist, können wir mit dem Training unserer Modelle auf Trainium fortfahren. Wie bereits erwähnt, trainieren wir unsere Modelle schrittweise, das heißt, wir laden einen zuvor trainierten Modellprüfpunkt und setzen das Training mit neuen Daten fort. PyTorch und PyTorch XLA ermöglichen einen nahtlosen Übergang zwischen Beschleunigern durch Checkpoint-Interoperabilität. Dank der Flexibilität, zwischen GPU und Trainium zu wechseln, konnten wir das vorherige GPU-Modell nahtlos laden und auf Trainium-Maschinen trainieren. Dies war von entscheidender Bedeutung, um sicherzustellen, dass wir unser Modell mit dem besten zuvor trainierten Modell initialisieren können, ohne dass es zu Produktionsausfällen oder Verlusten bei der Modellgenauigkeit kommt.

Da das GPU-Modell mit standardmäßigen PyTorch-Dienstprogrammen zum Speichern von Modellen gespeichert wurde, konnten wir das PyTorch-Dienstprogramm zum Laden von Prüfpunkten verwenden, um das GPU-Modell auf Trainium-Geräten zu laden.

Auf einer GPU/CPU können Sie das Modell beispielsweise mit dem folgenden Code speichern:

torch.save(model.state_dict(), PATH)

Anschließend laden Sie das Modell wieder auf Trainium:

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

Ebenso können Sie das Modell mit dem folgenden Code auf Trainium speichern:

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

Und laden Sie das Modell wieder auf GPU/CPU:

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

Da wir DDP für das Modelltraining verwenden, ist das Laden des Modells unabhängig von der Anzahl der Maschinen, die zum Trainieren des vorherigen Prüfpunkts verwendet wurden. Dadurch können wir die Trn1-Flotte horizontal skalieren, ohne Codeänderungen oder negative Auswirkungen auf das Modelltraining. Diese PyTorch-basierten Prüfpunkte können direkt verwendet oder sogar per Torch-Skript für Inferenz-Anwendungsfälle auf AWS Inferentia2 oder anderen Beschleunigern erstellt werden.

Betriebsstabilität

Es kann nicht genug betont werden, dass für die Ausführung von Workloads in der Produktion mehrere SLAs eingehalten werden müssen. Für unseren Anwendungsfall ist es neben den SLAs für Modellqualität und Trainingsdurchsatz unerlässlich, dass die Produktionspipeline betriebsstabil ist, was minimale Ausfallzeiten und Störungen während des Modelltrainings, der Bewertung und der Inferenz bedeutet.

Wie bei der bestehenden GPU-basierten Pipeline haben wir zahlreiche Mechanismen hinzugefügt, um die Pipeline betriebsstabil zu machen. Bevor wir mit dem Modelltraining beginnen, führen wir mehrere Integritätstests durch, um den Zustand der Maschinen zu beurteilen. Diese Tests umfassen im Allgemeinen einfache Tensoroperationen, um den Zustand der Beschleunigergeräte zu überprüfen. Wir haben beobachtet, dass es für verteiltes Training wichtig ist, Tests durchzuführen, um auch die kollektive Kommunikation zwischen Instanzen zu überprüfen. Wir haben das genutzt NCCOM-Testsuite aus dem Neuron SDK, um dies zu erreichen, indem eine Vielzahl von Operationen wie „All-Gather“, „All-Reduce“ und „Reduce-Scatter“ ausgeführt werden.

Selbst nachdem wir die von uns genannten Vorschläge befolgt haben, haben wir festgestellt, dass vorübergehende Probleme in jeder Pipeline unvermeidlich sind, unabhängig vom zugrunde liegenden Beschleuniger. Um die Ausfallsicherheit in jeder Trainingspipeline zu erhöhen, empfehlen wir den Einbau von Wiederholungsmechanismen, um diese potenziellen Probleme zu beheben. Wir gebrauchen Automatisierte AWS Batch-Wiederholungsversuche um Jobs erneut zu versuchen, bei denen während des Modelltrainings ein vorübergehender Fehler auftritt. Diese Neustarts können kostspielig sein, wenn gegen Ende des Trainings ein Fehler auftritt. Um diesem Problem entgegenzuwirken, haben wir unsere Trainingsskripte angepasst, um einen zuvor trainierten Modellprüfpunkt zu laden und das Training von diesem Punkt aus fortzusetzen. Mit dieser Funktionalität sind wir in der Lage, fehlgeschlagene Trainingsjobs mit minimalem Overhead aggressiv neu zu starten.

Mit diesen Ausfallsicherheitsmechanismen konnten wir für unsere Workloads auf Trn98.5 eine Erfolgsquote von 1 % erreichen, vergleichbar mit der Erfolgsquote unserer bestehenden GPU-Pipeline.

Die Ergebnisse

Um die Genauigkeit unserer Modelle zu validieren, haben wir zwei Modelle vom selben GPU-Checkpoint initialisiert und eines auf Trainium und das andere auf einer vergleichbaren GPU trainiert. Beide Modelle wurden mit den gleichen Trainingshyperparametern trainiert. Der für die Metrikberechnung verwendete Datensatz ist ein Holdout-Datensatz, und wir bewerten die Genauigkeit des Modells anhand dieses Datensatzes alle N globalen Schritte. Die X-Achse ist der globale Schritt und die Y-Achse ist die Modellgenauigkeit. Wir haben an jedem Punkt in der folgenden Grafik einen Unterschied in der Modellgenauigkeit von weniger als 0.1 % festgestellt.

Um die Kosteneffizienz des Modelltrainings zu bewerten, vergleichen wir außerdem lieber die Zeit, die benötigt wird, um die Modellkonvergenz zu erreichen. Wir glauben, dass dies eine praktischere Sicht auf Kosteneinsparungen im Vergleich zu Kennzahlen wie Kosten pro Token, erzielten FLOPS/Dollar und anderen Faktoren bietet. Unter Berücksichtigung der Trainingszeit von trn1.32xl und vergleichbar Amazon Elastic Compute-Cloud (Amazon EC2) haben wir beobachtet, dass Trainium bis zu 30 % günstigere Kosten für die Modellkonvergenz bietet.

Zusammenfassung

Bei der Bewertung verschiedener Beschleuniger für Ihre DL-Workloads sind viele Faktoren zu berücksichtigen. Zu den wichtigsten zählen Modellqualität, Durchsatz, Kosten und Verfügbarkeit. Es ist von größter Bedeutung, sicherzustellen, dass die Qualität und der Durchsatz Ihres Modells nicht aufgrund des von Ihnen gewählten Beschleunigers beeinträchtigt werden.

Dank unserer Partnerschaft und Zusammenarbeit mit dem Annapurna Neuron-Team konnte das Amazon Search M5-Team durch den Wechsel zu Trainium bis zu 30 % Kosten einsparen. Das Team ist in der Lage, Trainium zu nutzen und Modellqualität und Durchsatzparität mit vergleichbaren Beschleunigern auf dem Markt zu erreichen. Checkpoint-Interoperabilität und minimale Codeänderungen mit Unterstützung für XLA haben es M5 ermöglicht, zwischen mehreren Beschleunigern für ihre Arbeitslasten zu wählen. Dies hat es dem M5-Team ermöglicht, die große Rechenleistung von Trainium zu nutzen und beschleunigerunabhängige Lösungen zu entwickeln, um Amazon.com-Kunden zu begeistern. Aus betrieblicher Sicht hat sich Trainium als fähig erwiesen, Tier-1-Dienste im Amazon-Maßstab zu unterstützen. Das M5-Team verlagert weiterhin mehr Arbeitslasten auf Trainium, um Amazon die besten Modelle zu den niedrigsten Kosten bereitzustellen.

Zusammenfassend lässt sich sagen, dass das M5-Team durch die Aufnahme von Trainium in die Beschleunigerflotte in der Lage war, kostengünstige ML-Schulungen in Produktionsqualität durchzuführen. Wir empfehlen Ihnen, einen Blick auf Trainium und andere Neuron-Geräte wie AWS Inferentia zu werfen, um die Vorteile speziell entwickelter Amazon-Siliziumkomponenten für ML-Workloads zu nutzen. Beginnen Sie ganz einfach mit einem der vielen Tutorials zu verschiedenen Modellen, z Lama 2, verfügbar auf Trainium.

Über die Autoren

Abhinandan Patni ist Senior Software Engineer bei Amazon Search. Er konzentriert sich auf den Aufbau von Systemen und Werkzeugen für skalierbares verteiltes Deep-Learning-Training und Echtzeit-Inferenz.

James Park ist Lösungsarchitekt bei Amazon Web Services. Er arbeitet mit Amazon.com zusammen, um Technologielösungen auf AWS zu entwerfen, zu erstellen und bereitzustellen, und hat ein besonderes Interesse an KI und maschinellem Lernen. In seiner Freizeit erkundet er gerne neue Kulturen, neue Erfahrungen und bleibt über die neuesten Technologietrends auf dem Laufenden. Sie finden ihn auf LinkedIn.

Jerry Mannil ist Softwareentwickler bei Amazon Search. Er arbeitet an der Verbesserung der Effizienz, Robustheit und Skalierbarkeit der verteilten Trainingsinfrastruktur.

Ken Su ist Softwareentwickler bei Amazon Search. Er arbeitet an der Verbesserung der Trainingseffizienz und einem skalierbaren verteilten Trainingsworkflow. Außerhalb der Arbeit geht er gerne wandern und spielt Tennis.

RJ ist Ingenieur bei Amazon. Er erstellt und optimiert Systeme für verteilte Systeme für die Schulung und arbeitet an der Optimierung der Einführung von Systemen, um die Latenz für ML-Inferenz zu reduzieren. Außerhalb der Arbeit erforscht er den Einsatz generativer KI zum Erstellen von Lebensmittelrezepten.

• SEO-gestützte Content- und PR-Distribution. Holen Sie sich noch heute Verstärkung.
• PlatoData.Network Vertikale generative KI. Motiviere dich selbst. Hier zugreifen.
• PlatoAiStream. Web3-Intelligenz. Wissen verstärkt. Hier zugreifen.
• PlatoESG. Kohlenstoff, CleanTech, Energie, Umwelt, Solar, Abfallwirtschaft. Hier zugreifen.
• PlatoHealth. Informationen zu Biotechnologie und klinischen Studien. Hier zugreifen.
• Quelle: https://aws.amazon.com/blogs/machine-learning/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/