Ein Vollständiger Überblick Über Die Von Meta Verwendeten FX-Tools: Optimierung Von PyTorch-Modellen Mit Graphtransformation

Der Graphmodus in PyTorch ist leistungsfähiger. Dieser Artikel stellt Torch.FX vor, ein leistungsstarkes Tool, das den Graphen von PyTorch-Programmen erfassen und optimieren kann.

1. Einleitung

PyTorch unterstützt zwei Ausführungsmodi: Eager-Modus und Graph-Modus.

Im Eager-Modus werden Operatoren im Modell sofort ausgeführt, wenn sie gelesen werden. Es ist einfach zu verwenden und benutzerfreundlicher für Anwender des maschinellen Lernens, daher ist es als Standardausführungsmodus festgelegt.

Im Graphmodus werden Operatoren zunächst zu einem Graphen synthetisiert und dann als Ganzes kompiliert und ausgeführt. Es weist eine höhere Leistung auf und wird daher in der tatsächlichen Produktion häufig eingesetzt.

Insbesondere unterstützt der Graphmodus die Operatorfusion. Durch die Zusammenführung zweier Operatoren kann der Gesamtaufwand für Speicherlesevorgänge und Kernelstarts reduziert oder lokalisiert werden.

Die Fusion kann horizontal erfolgen:Nimmt eine einzelne Operation (wie BatchNorm), die auf mehrere Operanden angewendet wird, und führt sie zu einem einzigen Array zusammen.

Die Fusion kann auch vertikal erfolgen:Kombinieren Sie einen Kernel mit einem anderen Kernel, der die Ausgabe des ersten Kernels benötigt (z. B. ReLU gefolgt von Faltung).

Torch.FX (abgekürzt FX) ist ein öffentlich verfügbares Toolkit, das die Ausführung im Graphmodus als Teil des PyTorch-Pakets unterstützt. Es kann:

1. Holen Sie sich den Graphen aus dem PyTorch-Programm

2. Ermöglichen Sie Entwicklern, Transformationen auf das erhaltene Diagramm zu schreiben

Meta hat FX zuvor verwendet, um den Trainingsdurchsatz von Produktionsmodellen zu optimieren. In diesem Artikel wird die von Meta entwickelte FX-basierte Optimierung vorgestellt, um zu zeigen, wie die Leistung von in PyTorch bereitgestellten Modellen mithilfe der Graphtransformation optimiert werden kann.

II. Hintergrund

Einbettungstabellen werden häufig in Empfehlungssystemen verwendet.In diesem Abschnitt werden die Hintergrundinformationen zu FX und Einbettungstabellen vermittelt.

2.1. Effekte

Abbildung 1 ist ein einfaches Beispiel, das zeigt, wie ein PyTorch-Programm mit FX konvertiert wird.Es besteht aus drei Schritten:

• Holen Sie sich das Diagramm aus dem Programm
• Ändern Sie den Graphen (in diesem Fall verwenden wir GELU statt RELU).
• Generieren Sie ein neues Programm aus dem geänderten Graphen

Die FX-API bietet viele weitere Funktionen zum Überprüfen und Transformieren von PyTorch-Programmdiagrammen.

2.2. Einbettungstabelle

Im EmpfehlungssystemSpärliche Merkmale (z. B. Benutzer-ID, Story-ID) werden durch eine Einbettungstabelle dargestellt.

Die Einbettungstabelle E ist eine HxD-Matrix, wobei H die Hash-Größe und D die Einbettungsvektordimension ist. Jede Zeile von E ist ein Vektor von Gleitkommazahlen.

Die Funktion des Feature-Hashings besteht darin, ein spärliches Feature auf eine Indexliste von E abzubilden, z. B. [S1, S2, …, Sk], wobei 0 ≤ Si

Um die GPU voll auszunutzen, werden spärliche Features normalerweise in Stapeln verarbeitet.Jede Entität im Stapel hat ihre eigene Indexliste. Wenn ein Stapel B Entitäten hat, kann er einfach als eine Darstellung mit B Indexlisten verstanden werden.

Eine strengere Darstellung wäre, die B-Indexlisten zu einer Indexliste zusammenzuführen und eine Liste mit Indexlängen hinzuzufügen (eine für jede Entität im Stapel).

Wenn ein Stapel beispielsweise 3 Entitäten enthält, sieht seine Indexliste wie folgt aus:

• Entität 1: Indizes = [10, 20]
• Entität 2: Indizes = [5, 9, 77, 81]
• Entität 3: Indizes = [15, 20, 45]

Dann lauten Index und Länge der vollständigen Batchgröße:

• Indizes = [10, 20, 5, 9, 77, 81, 15, 20, 45]
• Längen = [2, 4, 3]

Die Ausgabe der Einbettungstabellenabfrage für den gesamten Stapel ist eine BxD-Matrix.

3. 3 FX-Transformationen

PyTorch hat drei FX-Transformationen aktualisiert, um den Zugriff auf die Einbettungstabelle zu beschleunigen. Diese werden in diesem Abschnitt nacheinander vorgestellt.

Nachfolgend finden Sie 3.1 zur Transformation der Kombination mehrerer kleiner Eingabetensoren zu einem großen Tensor; 3.2 über die Transformation der Zusammenführung mehrerer paralleler Rechenketten zu einer Rechenkette; und 3.3 über die Transformation der sich überschneidenden Kommunikation und Informatik.

3.1 Kombinieren spärlicher Eingabefunktionen

Jedes spärliche Eingabemerkmal in einem Stapel kann als zwei Listen dargestellt werden: eine Indexliste und eine Liste der Länge B, wobei B die Stapelgröße darstellt.

In PyTorch können beide Listen als Tensoren existieren.Wenn ein PyTorch-Modell auf einer GPU ausgeführt wird, wird die Einbettungstabelle normalerweise im GPU-Speicher gespeichert (der näher an der GPU liegt und über eine höhere Lese- und Schreibbandbreite als der CPU-Speicher verfügt).

Wenn die spärlichen Eingabefunktionen verwendet werden müssen, müssen beide Tensoren zuerst von der CPU auf die GPU kopiert werden. Allerdings erfordert jede Speicherkopie vom Host auf das Gerät das Starten eines Kernels, was zeitaufwändiger ist als die eigentliche Datenübertragung.

Wenn ein Modell viele spärliche Eingabefunktionen verwendet, kann dieses Kopieren zu einem Leistungsengpass werden (beispielsweise erfordern 1000 spärliche Eingabefunktionen das Kopieren von 2000 Tensoren vom Host auf das Gerät).

Eine Optimierung zur Reduzierung der Anzahl von Memcpys vom Host zum Gerät besteht darin, mehrere spärliche Eingabefunktionen zu kombinieren, bevor sie an das Gerät gesendet werden.

Nehmen wir beispielsweise die folgenden drei Eingabefunktionen an:

• Feature_A: Indizes = [106, 211, 7], Längen = [2, 1]
• Feature_B: Indizes = [52, 498, 616, 870, 1013], Längen = [3, 2]
• Feature_C: Indizes = [2011, 19, 351, 790], Längen = [1, 3]

Die kombinierte Form lautet:

Features_A_B_C: Indizes = [106, 211, 7, 52, 498, 616, 870, 1013, 2011, 19, 351, 790], Längen = [2, 1, 3, 2, 1, 3]

Anstatt also 3 × 2 = 6 Tensoren vom Host auf das Gerät zu kopieren, müssen nur 2 Tensoren kopiert werden.

Abbildung 3(b) zeigt die Implementierung dieser Optimierung, die aus zwei Komponenten besteht:

• CPU-Seite:Die Eingabepipeline wird geändert, um die Indizes aller spärlichen Merkmale in einem Tensor und alle Längen in einem anderen Tensor zu kombinieren. Diese beiden Tensoren werden dann auf die GPU kopiert.
• GPU-Seite:Fügen Sie mithilfe von FX einen Permute_and_Split-Operator in den Modellgraphen ein, um die einzelnen Merkmalsindizes und Längentensoren aus den zusammengeführte Tensoren wiederherzustellen und sie an die entsprechenden Knoten weiter unten zu senden.

3.2 Horizontale Fusion von Berechnungsketten beginnend mit dem Zugriff auf die Einbettungstabelle

In einem Produktionsmodell sind üblicherweise 10 Einbettungstabellen pro GPU vorhanden. Aus LeistungsgründenAbfragen an diese Tabellen werden gruppiert, sodass ihre Ausgaben zu einem großen Tensor verkettet werden.(Siehe den roten Teil in Abbildung 4(a)).

Um eine einzelne Feature-Ausgabe zu berechnen,Verwenden Sie den Split-Operator, um einen großen Tensor in N kleine Tensoren aufzuteilen(wobei N die Anzahl der Merkmale ist) und wendet dann die gewünschte Berechnung auf jeden Tensor an.

Wie in Abbildung 4(a) dargestellt, wird für jede Merkmalsausgabe O die Berechnung Tanh(LayerNorm(O)) angewendet. Alle Berechnungsergebnisse werden zu einem großen Tensor verkettet und dann an den nachgelagerten Operator (Op1 in Abbildung 4(a)) übergeben.

Die Hauptlaufzeitkosten entstehen hier durch den Overhead beim Start des GPU-Kernels.Beispielsweise beträgt die Anzahl der Starts des GPU-Kernels in Abbildung 4(a) 2*N+3 (jede Ellipse in der Abbildung stellt einen GPU-Kernel dar). Dies wirkt sich auf die Leistung aus, da die Ausführungszeit von LayerNorm und Tanh auf der GPU im Vergleich zur Startzeit ihres Kernels sehr kurz ist.

Darüber hinaus kann der Split-Operator eine zusätzliche Kopie des Ausgabetensors des Einbettungsvektors erstellen und so zusätzlichen GPU-Speicher verbrauchen.

Die Verwendung von FX zur Implementierung einer Optimierung namens Horizontal Fusion kann die Anzahl der GPU-Kernelstarts erheblich reduzieren(In diesem Beispiel beträgt die Anzahl der GPU-Kernelstarts nach der Optimierung 5, siehe Abbildung 4(b)).

Verwenden Sie den Operator Add_middle_dim anstelle eines expliziten Splits, um den 2D-Einbettungstensor der Form (B, NxD) in einen 3D-Tensor der Form (B, N, D) umzuformen. Als nächstes wird eine einzelne LayerNorm auf die letzte Dimension angewendet. Wendet einen Tanh auf das Ergebnis von LayerNorm an. Schließlich wird der Operator Remove_middle_dim verwendet, um das Tanh-Ergebnis in einem 2D-Tensor wiederherzustellen.

Da Add_middle_dim und Remove_middle_dim lediglich den Tensor umformen,Es werden keine zusätzlichen Kopien erstellt, sodass auch der GPU-Speicherverbrauch reduziert werden kann.

3.3 Überschneidungen zwischen Berechnung und Kommunikation

Das Training von Empfehlungsmodellen für die Produktion erfolgt üblicherweise auf verteilten GPU-Systemen.Da die Gerätespeicherkapazität jeder GPU nicht ausreicht, um alle Einbettungstabellen im Modell zu speichern, müssen sie auf mehrere GPUs verteilt werden.

Während des Trainingsschritts muss eine GPU Merkmalswerte aus der Einbettungstabelle auf anderen GPUs lesen/schreiben.Dies wird als All-to-All-Kommunikation bezeichnet und kann zu erheblichen Leistungseinbußen führen.

Durch die Implementierung einer Transformation durch FX ist es möglich, die Berechnung mit der All-to-All-Kommunikation zu überlappen.Abbildung 5(a) zeigt ein Beispiel für einen Modellgraphen mit Zugriff auf die Einbettungsvektortabelle (EmbeddingAllToAll) und anderen Operatoren. Wie in Abbildung 5(b) gezeigt, werden sie ohne Optimierung sequenziell auf einem GPU-Stream ausgeführt.

Verwenden Sie FX, um EmbeddingAllToAll in EmbeddingAllToAll_Request und EmbeddingAllToAll_Wait aufzuteilen und unabhängige Operatoren dazwischen anzuordnen.

3.4 Zusammenfassung

Um herauszufinden, welche Modelle von diesen Transformationen profitieren würden, analysierten die Entwickler die von MAIProf gesammelten Leistungsdaten für im Meta Data Center ausgeführte Modelle.Wir zeigen, dass diese Transformationen bei einer Reihe von Produktionsmodellen im Vergleich zum Eager-Modus eine zwei- bis dreimal höhere Geschwindigkeit erreichen.

IV. Abschluss

Aus Leistungssicht ist der Graphmodus in PyTorch dem in Produktionsumgebungen verwendeten Eagermodus vorzuziehen. FX ist ein leistungsstarkes Tool zum Erfassen und Optimieren von PyTorch-Programmdiagrammen. Dieser Artikel stellt drei FX-Transformationen zur Optimierung von Produktionsempfehlungsmodellen innerhalb von Meta vor.

Abschließend hoffe ich, dass mehr PyTorch-Entwickler die Graphtransformation nutzen können, um die Leistung des Modells zu verbessern.

​—— Ende ——