Wie der latente Raum die Amazon SageMaker-Modellparallelitätsbibliothek nutzte, um die Grenzen großer Transformatoren zu erweitern

Neuauflage von Plato

Verfolger: 0

Dieser Blog wurde von Sarah Jane Hong CSO, Darryl Barnhart CTO und Ian Thompson, CEO von Latent Space und Prem Ranga von AWS, gemeinsam verfasst.

Latenter Raum ist eine versteckte Darstellung abstrakter Ideen, die Modelle des maschinellen Lernens (ML) lernen. Zum Beispiel sind "Hund", "Blume" oder "Tür" Konzepte oder Orte im latenten Raum. Beim Latenter RaumWir arbeiten an einer Engine, mit der Sie diesen Bereich mit sprachlichen und visuellen Eingabeaufforderungen bearbeiten und erkunden können. Das Latent Space-Team kommt aus zwei Bereichen, die sich seit langem kaum überschneiden: Grafik und Verarbeitung natürlicher Sprache (NLP). Traditionell wurden die Modalitäten von Bildern und Text separat behandelt, wobei jede ihre eigene Geschichte komplexer, teurer und fragiler Feature-Engineering-Vorgänge aufweist. NLP-Aufgaben wie das Verstehen von Dokumenten oder das Beantworten von Fragen hatten normalerweise wenig mit Visionsaufgaben wie dem Verstehen oder Rendern von Szenen zu tun, und normalerweise verwenden wir für jede Aufgabe sehr unterschiedliche Ansätze und Modelle. Dies ändert sich jedoch schnell.

Diese Verschmelzung von Modalitäten in einem einzigen gemeinsamen latenten Raum eröffnet eine neue Generation kreativer und kommerzieller Anwendungen, vom Spielen bis zum Verständnis von Dokumenten. Das Freischalten dieser neuen Anwendungen in einem einzigen Modell eröffnet jedoch neue Herausforderungen bei der Skalierung, wie in „The Bitter Lesson“ von Richard Sutton und der aufregenden Arbeit der letzten Jahre an Skalierungsgesetzen hervorgehoben. Um dies zu ermöglichen, arbeitet Latent Space an modernster Forschung, um diese Modalitäten in einem einzigen Modell zusammenzuführen, aber auch zu skalieren und dies effizient zu tun. Hier kommt die Modellparallelität ins Spiel.

Amazon Sage MakerDie einzigartige automatisierte Modellpartitionierung und der effiziente Pipelining-Ansatz ermöglichten die Einführung der Modellparallelität mit geringem Engineering-Aufwand und wir haben unser Training von Modellen über 1 Milliarde Parameter hinaus skaliert (wir verwenden die p4d.24xlarge A100-Instanzen), was für uns eine wichtige Voraussetzung ist. Darüber hinaus stellten wir fest, dass wir beim Training mit einem 16-Knoten- und acht GPU-Trainingsaufbau mit der SageMaker-Modellparallelitätsbibliothek eine Effizienzsteigerung von 38% im Vergleich zu unseren vorherigen Trainingsläufen verzeichneten.

Herausforderungen bei der Ausbildung von Großtransformatoren

Bei Latent Space verbinden wir Sprache und Vision in Transformatormodellen mit Milliarden von Parametern, um Anwendungsfälle zu unterstützen, die aus der Vorstellungskraft eines Benutzers stammen oder die in der realen Welt auftreten würden, jedoch nicht in unseren Trainingsdaten. Wir bewältigen die Herausforderungen bei der Skalierung auf Milliarden von Parametern und darüber hinaus auf zwei verschiedene Arten:

Techniken zum Abrufen von Informationen sind seit langem eine Schlüsselkomponente von Suchmaschinen und QS-Aufgaben. In jüngster Zeit wurden aufregende Fortschritte bei der Kombination klassischer IR-Techniken mit modernen Transformatoren erzielt, insbesondere für Aufgaben zur Beantwortung von Fragen, bei denen ein Modell gemeinsam mit einem neuronalen Retriever trainiert wird, der lernt, relevante Dokumente abzurufen, um Fragen zu beantworten. Eine Übersicht finden Sie in der aktuellen Arbeit von FAIR in Retrieval Augmented Generation: Optimierung der Erstellung intelligenter Modelle zur Verarbeitung natürlicher Sprache und Fusion-in-Decoder, Google Brain's REICHund Nvidias Neuronaler Retriever zur Beantwortung von Fragen.

Retrieval-Augmented-Techniken helfen zwar bei Kosten und Effizienz, aber wir können das Modell für unser größtes Modell immer noch nicht auf eine einzige GPU anpassen. Dies bedeutet, dass wir Modellparallelität verwenden müssen, um es zu trainieren. Aufgrund der Art unserer Abrufarchitektur war das Entwerfen unserer Modellaufteilung jedoch aufgrund der Interdependenzen zwischen abgerufenen Kontexten über Trainingseingaben hinweg eine Herausforderung. Selbst wenn wir bestimmen, wie wir unser Modell aufteilen, war die Einführung der Modellparallelität für uns eine wichtige technische Aufgabe, die wir über unseren gesamten Forschungs- und Entwicklungslebenszyklus hinweg manuell durchführen mussten.

Die SageMaker-Modellparallelitätsbibliothek

Bei der Modellparallelität wird ein Modell auf mehrere Geräte oder Knoten (z. B. mit GPU ausgestattete Instanzen) aufgeteilt und eine effiziente Pipeline erstellt, um das Modell über diese Geräte hinweg zu trainieren und die GPU-Auslastung zu maximieren. Das Modellparallelitätsbibliothek In SageMaker wird die Modellparallelität durch die Bereitstellung einer automatisierten Modellaufteilung, auch als Modellaufteilung bezeichnet, leichter zugänglich automatisierte Modellpartitionierung und ausgefeilte Pipeline-Laufplanung. Die Modellaufteilungsalgorithmen können hinsichtlich Geschwindigkeit oder Speicherverbrauch optimiert werden. Die Bibliothek verwendet einen Partitionierungsalgorithmus, der den Speicher ausgleicht, die Kommunikation zwischen Geräten minimiert und die Leistung optimiert.

Automatisierte Modellpartitionierung

Für unseren PyTorch-Anwendungsfall führt die Modellparallelbibliothek intern einen Verfolgungsschritt (im ersten Trainingsschritt) aus, der den Modellgraphen erstellt und die Tensor- und Parameterformen bestimmt. Anschließend wird ein Baum erstellt, der aus den verschachtelten besteht nn.Module Objekte im Modell sowie zusätzliche Daten aus der Ablaufverfolgung, z. B. die Menge der gespeicherten Daten nn.Parametersund Laufzeit für jeden nn.Module.

Die Bibliothek durchläuft diesen Baum dann vom Stamm aus und führt einen Partitionierungsalgorithmus aus, der die Rechenlast und die Speichernutzung in Einklang bringt und die Kommunikation zwischen Instanzen minimiert. Wenn mehrere nn.Module denselben nn.Parameter verwenden, werden diese Module auf demselben Gerät platziert, um zu vermeiden, dass mehrere Versionen desselben Parameters verwaltet werden. Nachdem die Partitionierungsentscheidung getroffen wurde, werden die zugewiesenen Module und Gewichte auf ihre Geräte geladen.

Pipeline-Laufplanung

Ein weiteres Kernmerkmal der parallelen Bibliothek für verteilte Modelle von SageMaker ist Pipeline-Läufe, die die Reihenfolge bestimmen, in der Berechnungen durchgeführt und Daten während des Modelltrainings geräteübergreifend verarbeitet werden. Das Pipelining basiert auf der Aufteilung eines Mini-Batches in Mikrobatches, die einzeln in die Trainingspipeline eingespeist werden und einem durch die Laufzeit der Bibliothek definierten Ausführungsplan folgen.

Die Mikrobatch-Pipeline stellt sicher, dass alle GPUs voll ausgelastet sind, was wir selbst erstellen müssten. Mit der Modellparallelitätsbibliothek geschieht dies jedoch ordentlich hinter den Kulissen. Zuletzt können wir verwenden Amazon FSxDies ist wichtig, um sicherzustellen, dass unsere Lesegeschwindigkeiten angesichts der Anzahl der Dateien, die während des Trainings eines multimodalen Modells mit Abruf gelesen werden, schnell sind.

Trainingsarchitektur

Das folgende Diagramm zeigt, wie wir unsere Trainingsarchitektur einrichten. Unser Hauptziel war es, die Trainingsgeschwindigkeit zu verbessern und die Kosten zu senken. Die Bild- und Sprachtransformatoren, die wir trainieren, sind sehr komplex, mit einer beträchtlichen Anzahl von Ebenen und Gewichten im Inneren, die Milliarden von Parametern umfassen, was dazu führt, dass sie nicht in den Speicher eines einzelnen Knotens passen. Jeder Knoten enthält eine Teilmenge des Modells, über die die Datenflüsse und Transformationen gemeinsam genutzt und kompiliert werden. Wir richten 16 ein p4d.24xlarge Instanzen mit jeweils acht GPUs unter Verwendung der folgenden Architekturdarstellung:

Wenn wir unsere Modelle vergrößern, besteht ein allgemeiner Trend darin, alles in den Gewichten des Netzwerks zu speichern. Aus praktischen Gründen möchten wir unsere Modelle jedoch erweitern, um zu lernen, wie man nach relevanten Kontexten sucht, um beim Rendern zu helfen. Dies ermöglicht es uns, unsere Servierkosten niedrig zu halten, ohne die Bildqualität zu beeinträchtigen. Wir verwenden ein großes transformatorbasiertes NLP-Modell und haben, wie bereits erwähnt, mit der SageMaker-Modellparallelitätsbibliothek eine Steigerung der Trainingseffizienz um 38% beobachtet, wie im Folgenden gezeigt:

Bei Parallelität auf Tensorebene benötigen wir für jede Berechnung eine Allreduktion. Dies erfordert O (log_2 n) parallele Schritte. Das sind n Maschinen, die O (n) Schritte ausführen, für O (n log_2 n) Gesamtoperationen.
Für die Pipeline-Parallelität benötigen wir O (1) parallele Schritte, um Daten durch die Pipeline zu leiten
Bei 16 Maschinen mit acht GPUs haben wir O (1) Kosten für die parallele Pipeline und O (log_2 (8)) = O (3) Kosten für die parallele Tiefenmodellierung.
In diesem Fall sehen wir, dass die Netzwerkkosten durch Umschalten auf Pipeline parallel zu dem, was wir mit der SageMaker-Modellparallelität verwenden, auf 1/3 reduziert werden und die Gesamtschulungskosten auf 1/2 + 1/2 * 1 / log_2 (16) reduziert werden ) = 0.625 der ursprünglichen Kosten, was zu einer entsprechenden Effizienzverbesserung führt.

Wenn der Bedarf verteiltes Training erfordert (Probleme mit der Skalierung der Modellgröße oder der Trainingsdaten), können wir im Allgemeinen eine Reihe von Best Practices befolgen, um festzustellen, welcher Ansatz am besten funktioniert.

Best Practices für verteiltes Training

Basierend auf unseren Erfahrungen empfehlen wir, mit einem verteilten datenparallelen Ansatz zu beginnen. Verteilte Datenparallelität wie die SageMaker verteilte parallele Datenbibliothek Behebt die meisten Netzwerkprobleme mit Modellreplikaten. Daher sollten Sie Modelle in die kleinste Anzahl von Knoten einpassen und dann replizieren, um die Stapelgröße nach Bedarf zu skalieren.

Wenn Ihnen während des Trainings der Speicherplatz ausgeht, wie in diesem Szenario, möchten Sie möglicherweise zu einem modellparallelen Ansatz wechseln. Berücksichtigen Sie jedoch diese Alternativen, bevor Sie ein modellparalleles Training ausprobieren:

Verwenden Sie auf mit NVIDIA Tensor Core ausgestatteter Hardware Training mit gemischter Präzision um die Geschwindigkeit zu erhöhen und den Speicherverbrauch zu reduzieren.
Reduzieren Sie die Stapelgröße (oder reduzieren Sie die Bildauflösung oder die NLP-Sequenzlänge, wenn möglich).

Darüber hinaus bevorzugen wir Modellkonstruktionen ohne Batch-Normalisierung, wie in beschrieben Hochleistungs-Bilderkennung in großem Maßstab ohne Normalisierung. Wenn dies nicht vermieden werden kann, stellen Sie sicher, dass die Batch-Normalisierung geräteübergreifend synchronisiert wird. Wenn Sie verteiltes Training verwenden, wird Ihr Stapel auf GPUs aufgeteilt, sodass genaue Stapelstatistiken eine Synchronisierung auf allen Geräten erfordern. Ohne dies hat die Normalisierung einen erhöhten Fehler und beeinträchtigt dadurch die Konvergenz.

Beginnen Sie mit dem modellparallelen Training, wenn Sie die folgenden Einschränkungen haben:

Ihr Modell passt nicht auf ein einzelnes Gerät
Aufgrund Ihrer Modellgröße sind Sie bei der Auswahl größerer Stapelgrößen mit Einschränkungen konfrontiert, z. B. wenn Ihre Modellgewichte den größten Teil Ihres GPU-Speichers beanspruchen und Sie gezwungen sind, eine kleinere, suboptimale Stapelgröße zu wählen

Gehen Sie bei der Leistungsoptimierung wie folgt vor:

Verwenden Sie Pipelining für die Kommunikation zwischen Knoten, um die Latenz zu minimieren und den Durchsatz zu erhöhen
Halten Sie die Rohrleitungen so kurz wie möglich, um Blasenbildung zu minimieren. Die Anzahl der Mikrobatches sollte so eingestellt werden, dass die Recheneffizienz mit der Blasengröße in Einklang gebracht wird und mindestens die Pipeline-Länge beträgt. Bei Bedarf können Sie Mikrobatches auf Token-Ebene erstellen, wie in beschrieben TeraPipe: Token Level Pipeline Parallelity zum Trainieren umfangreicher Sprachmodelle

Verwenden Sie zur Kostenoptimierung von SageMaker verwaltete Spot-Instanzen für Schulungen. Dadurch können die Kosten für Schulungsmodelle gegenüber On-Demand-Instanzen um bis zu 90% optimiert werden. SageMaker verwaltet die Spot-Unterbrechungen in Ihrem Namen.

Weitere zu berücksichtigende Faktoren:

Innerhalb eines Knotens ist eine schnelle Verbindung nuancierter. Wenn ausreichend Netzwerkkapazität innerhalb des Knotens vorhanden ist, kann das erneute Mischen von Daten für eine optimalere Berechnung einen Vorteil zeigen.
Wenn die Aktivierungen viel größer als die Gewichtstensoren sind, kann auch ein Sharded-Optimierer hilfreich sein. Bitte beziehen Sie sich auf Null für weitere Informationen an.

Die folgende Tabelle enthält einige gängige Scaleup-Szenarien für Schulungen und deren Konfiguration in AWS.

Szenario	Wann gilt es?	Lösung
Skalierung von einer einzelnen GPU auf viele GPUs	Wenn die Menge der Trainingsdaten oder die Größe des Modells zu groß ist	Wechseln Sie zu einer Instanz mit mehreren GPUs wie p3.16xlarge mit acht GPUs, wobei die Daten und die Verarbeitung auf die acht GPUs aufgeteilt sind, und führen Sie in der Zeit, die zum Trainieren Ihres Modells benötigt wird, zu einer nahezu linearen Beschleunigung.
Skalierung von einer einzelnen Instanz auf mehrere Instanzen	Wenn die Skalierungsanforderungen über das Ändern der Instanzgröße hinausgehen	Skalieren Sie die Anzahl der Instanzen mit der Schätzfunktion des SageMaker Python SDK, indem Sie Ihren Instanztyp auf p3.16xlarge und instance_count auf 2 setzen. Anstelle der acht GPUs auf einer einzelnen p3.16xlarge verfügen Sie über 16 GPUs auf zwei identischen Instanzen. Erwägen Sie die Verwendung der SageMaker verteilte parallele Datenbibliothek.
Auswahl eines modellparallelen Ansatzes für das Training	Bei Fehlern aufgrund von Speichermangel während des Trainings	Wechseln Sie mit dem zu einem modellparallelen Ansatz Parallele Bibliothek für verteilte Modelle von SageMaker.
Netzwerkleistung für die Kommunikation zwischen Knoten	Für verteiltes Training mit mehreren Instanzen (z. B. Kommunikation zwischen den Knoten im Cluster bei einer AllReduce-Operation)	Ihre Instanzen müssen sich in derselben Region und derselben Verfügbarkeitszone befinden. Wenn Sie das SageMaker Python SDK verwenden, wird dies für Sie erledigt. Ihre Trainingsdaten sollten sich ebenfalls in derselben Verfügbarkeitszone befinden. Erwägen Sie die Verwendung der SageMaker verteilte parallele Datenbibliothek.
Optimierte GPU, Netzwerk und Speicher	Für groß angelegte verteilte Schulungsbedürfnisse	Der Instanztyp p4d.24xlarge wurde für schnellen lokalen Speicher und eine schnelle Netzwerk-Backplane mit bis zu 400 Gigabit entwickelt. Wir empfehlen ihn als leistungsstärkste Option für verteiltes Training.

Zusammenfassung

Mit der Modellparallelbibliothek in SageMaker erhalten wir viele Vorteile aus der Box, wie z. B. automatisierte Modellpartitionierung und effizientes Pipelining. In diesem Beitrag haben wir unsere Herausforderungen mit unserem ML-Anwendungsfall, unseren Überlegungen zu verschiedenen Trainingsansätzen und der Verwendung der Amazon SageMaker-Modellparallelitätsbibliothek zur Beschleunigung unseres Trainings geteilt. Das Beste ist, dass es nur noch wenige Stunden dauern kann, um die hier beschriebenen Best Practices für Modellparallelität und Leistungsverbesserungen zu übernehmen. Wenn dieser Beitrag Ihnen hilft oder Sie zur Lösung eines Problems inspiriert, würden wir gerne davon hören! Bitte teilen Sie Ihre Kommentare und Feedback.

Bibliographie

Weitere Informationen finden Sie unter:

Über die Autoren

Prem Ranga ist ein Enterprise Solutions Architect mit Sitz in Atlanta, GA. Er ist Teil der Technical Field Community für maschinelles Lernen und liebt es, mit Kunden auf ihrer ML- und AI-Reise zusammenzuarbeiten. Prem hat eine Leidenschaft für Robotik, ist ein Forscher für autonome Fahrzeuge und hat auch die von Alexa kontrollierten Beer Pours in Houston und anderen Orten gebaut.

Sarah Jane Hong ist Mitbegründer und Chief Science Officer von Latent Space. Ihr Hintergrund liegt an der Schnittstelle von Mensch-Computer-Interaktion und maschinellem Lernen. Zuvor leitete sie die NLP-Forschung bei Sonar (von Marchex übernommen), das Unternehmen im Bereich der KI-Konversation dient. Sie ist auch eine angesehene AR / VR-Entwicklerin, die Auszeichnungen und Stipendien von Oculus, Mozilla Mixed Reality und Microsoft Hololens erhalten hat.

Darryl Barnhart ist Mitbegründer und Chief Technology Officer von Latent Space. Er ist ein erfahrener Entwickler mit Erfahrung in den Bereichen GPU-Beschleunigung, Computergrafik, umfangreiche Daten und maschinelles Lernen. Andere Leidenschaften sind Mathematik, Spieleentwicklung und das Studium von Informationen.

Ian Thompson ist der Gründer und CEO von Latent Space. Ian ist ein Ingenieur und Forscher, der von den „benachbarten Möglichkeiten“ inspiriert ist - Technologien, die einen großen Einfluss auf unser Leben haben werden. Derzeit liegt der Schwerpunkt auf der Vereinfachung und Skalierung des Lernens multimodaler Repräsentation, um eine sichere und kreative KI aufzubauen. Zuvor war er am Aufbau von Unternehmen in den Bereichen Grafik / Virtual Reality (AltspaceVR, von Microsoft übernommen) und Education / NLP (HSE) beteiligt.