Generative KI auf AWS

Kapitel 4. Optimierungen für Speicher und Rechenleistung

Diese Arbeit wurde mithilfe von KI übersetzt. Wir freuen uns über dein Feedback und deine Kommentare: [e-mail]

In Kapitel 3 hast du bewährte Methoden für das Ausprobieren und Auswählen eines Basismodells für deinen Anwendungsfall kennengelernt. Der nächste Schritt besteht in der Regel darin, das Modell an deine spezifischen Bedürfnisse und Datensätze anzupassen. Dazu gehört auch die Anpassung des Modells an deine Datensätze mit Hilfe einer Technik namens Feinabstimmung, die du in Kapitel 5 genauer kennenlernen wirst. Wenn du große Basismodelle trainierst oder fein abstimmst, stehst du oft vor der Frage, wie du große Modelle in den GPU-Speicher einbauen kannst.

In diesem Kapitel lernst du Techniken kennen, mit denen du Speicherbeschränkungen überwinden kannst. Du lernst, wie du Quantisierung und verteiltes Training einsetzen kannst, um den benötigten GPU-RAM zu minimieren, und wie du das Modelltraining bei größeren Modellen horizontal auf mehrere GPUs verteilen kannst.

Das ursprüngliche Falcon-Modell mit 40 Milliarden Parametern wurde zum Beispiel auf einem Cluster von 48 ml.p4d.24xlarge Amazon SageMaker-Instanzen trainiert, die aus 384 NVIDIA A100-GPUs, 15 TB GPU-RAM und 55 TB CPU-RAM bestanden. Eine neuere Version von Falcon wurde auf einem Cluster von 392 ml.p4d.24xlarge SageMaker-Instanzen mit 3.136 NVIDIA A100 GPUs, 125 TB GPU-RAM und 450 TB CPU-RAM trainiert. Die Größe und Komplexität des Falcon-Modells erfordert einen Cluster von Grafikprozessoren, profitiert aber auch von der Quantisierung, wie du gleich sehen wirst.

Memory-Herausforderungen

Eines der häufigsten Probleme beim Training oder der Feinabstimmung von Basismodellen ist, dass dir der Speicher ausgeht. Wenn du schon einmal versucht hast, dein Modell auf NVIDIA-GPUs zu trainieren oder auch nur zu laden, kommt dir die Fehlermeldung in Abbildung 4-1 vielleicht bekannt vor.

Abbildung 4-1. CUDA out-of-memory Fehler

CUDA, die Abkürzung für Compute Unified Device Architecture, ist eine Sammlung von Bibliotheken und Tools, die für NVIDIA-Grafikprozessoren entwickelt wurden, um die Leistung gängiger Deep-Learning-Operationen, wie z. B. Matrixmultiplikation, zu steigern. Deep-Learning-Bibliotheken wie PyTorch und TensorFlow nutzen CUDA ausgiebig, um die hardwarespezifischen Low-Level-Details zu handhaben, einschließlich des Datenaustauschs zwischen CPU und GPU-Speicher. Da moderne generative Modelle mehrere Milliarden Parameter enthalten, bist du während der Entwicklung beim Laden und Testen eines Modells in deiner Forschungsumgebung wahrscheinlich schon auf diesen Out-of-Memory-Fehler gestoßen.

Ein einzelner Modellparameter, mit voller 32-Bit-Genauigkeit, wird durch 4 Bytes dargestellt. Ein Modell mit 1 Milliarde Parametern benötigt also 4 GB GPU-RAM, nur um das Modell mit voller Genauigkeit in den GPU-RAM zu laden. Wenn du das Modell auch trainieren willst, brauchst du mehr GPU-Speicher, um die Zustände des numerischen Optimierers, die Gradienten und Aktivierungen sowie alle temporären Variablen zu speichern, die von deinen Funktionen verwendet werden (siehe Tabelle 4-1).

Tabelle 4-1. Zusätzlich benötigter Arbeitsspeicher zum Trainieren eines Modells

Staaten	Bytes pro Parameter
Modellparameter (Gewichte)	4 Bytes pro Parameter
Adam Optimierer (2 Zustände)	8 Bytes pro Parameter
Farbverläufe	4 Bytes pro Parameter
Aktivierungen und temporärer Speicher (variable Größe)	8 Bytes pro Parameter (High-End-Schätzung)
GESAMT	= 4 + 20 Bytes pro Parameter

Tipp

Wenn du mit dem Training eines Modells experimentierst, ist es empfehlenswert, mit batch_size=1 zu beginnen, um die Speichergrenzen des Modells mit einem einzigen Trainingsbeispiel zu ermitteln. Dann kannst du die Stapelgröße schrittweise erhöhen, bis du den CUDA Out-of-Memory-Fehler erreichst. Dadurch wird die maximale Stapelgröße für das Modell und den Datensatz bestimmt. Mit einer größeren Stapelgröße kannst du das Training deines Modells oft beschleunigen.

Diese zusätzlichen Komponenten führen zu etwa 12-20 zusätzlichen Bytes GPU-Speicher pro Modellparameter. Um beispielsweise ein Modell mit 1 Milliarde Parametern zu trainieren, benötigst du etwa 24 GB GPU-RAM bei voller 32-Bit-Präzision, also sechsmal mehr Speicher als nur 4 GB GPU-RAM zum Laden des Modells, wie in Abbildung 4-2 dargestellt.

Abbildung 4-2. Vergleich des ungefähren GPU-RAM-Bedarfs zum Laden gegenüber dem Laden und Trainieren eines Modells mit 1 Milliarde Parametern bei voller 32-Bit-Präzision

Es ist erwähnenswert, dass die NVIDIA A100 und H100, die zum Zeitpunkt der Erstellung dieses Artikels verwendet wurden, nur bis zu 80 GB GPU-RAM unterstützen. Und da du wahrscheinlich Modelle mit mehr als 1 Milliarde Parametern trainieren willst, musst du eine Lösung finden, z. B. die Quantisierung deines Modells.

AWS hat außerdem speziell entwickelte ML-Beschleuniger, AWS Trainium, für das leistungsstarke und kosteneffiziente Training von generativen KI-Modellen mit mehr als 100B Parametern entwickelt. Du kannst die AWS Trainium-Chips über die Trn1 Instance-Familie nutzen. Die größte Instanz von Trn1 wird derzeit von 16 AWS Trainium-Chips angetrieben und verfügt über 512 GB gemeinsamen Beschleunigerspeicher. Außerdem sind die Trn1 Instanzen für Quantisierung und verteilte Modellschulung optimiert und unterstützen eine Vielzahl von Datentypen.

Quantisierung ist eine beliebte Methode , um deine Modellparameter von einer 32-Bit-Präzision auf eine 16-Bit-Präzision - oder sogar 8-Bit oder 4-Bit - umzurechnen. Indem du deine Modellgewichte von voller 32-Bit-Präzision auf halbe 16-Bit-Präzision quantisierst, kannst du den Speicherbedarf deines 1-Milliarde-Parameter-Modells schnell um 50 % auf nur 2 GB beim Laden und 12 GB beim Training reduzieren.

Doch bevor wir uns mit der Quantisierung beschäftigen, wollen wir uns mit den gängigen Datentypen für das Modelltraining und der numerischen Genauigkeit beschäftigen.

Datentypen und numerische Präzision

Im Folgenden sind die verschiedenen Datentypen aufgeführt, die von PyTorch und TensorFlow verwendet werden: fp32 für 32-Bit volle Genauigkeit, fp16 für 16-Bit halbe Genauigkeit und int8 für 8-Bit ganzzahlige Genauigkeit.

In jüngster Zeit hat sich bfloat16 zu einer beliebten Alternative zu fp16 für 16-Bit-Präzision in moderneren generativen KI-Modellen entwickelt. bfloat16 (oder bf16) ist die Abkürzung für "Brain Floating Point 16", da es bei Google Brain entwickelt wurde. Im Vergleich zu fp16 hat bfloat16 mit 8 Bits für den Exponenten einen größeren Dynamikbereich und kann daher eine große Bandbreite an Werten darstellen, die wir in generativen KI-Modellen finden.

Wir besprechen, wie sich diese Datentypen vergleichen lassen und warum bfloat16 eine beliebte Wahl für die 16-Bit-Quantisierung ist.

Angenommen, du möchtest pi mit voller 32-Bit-Genauigkeit auf 20 Dezimalstellen (3.14159265358979323846) speichern. Erinnere dich daran, dass Fließkommazahlen als eine Reihe von Bits gespeichert werden, die nur aus 0 und 1 bestehen. Zahlen werden in 32 Bits gespeichert, wobei 1 Bit für das Vorzeichen (negativ oder positiv), 8 Bits für den Exponenten (der den dynamischen Bereich darstellt) und 23 Bits für den Bruch, auch Mantisse oder Signifikant genannt, der die Genauigkeit der Zahl darstellt, verwendet werden. Tabelle 4-2 zeigt, wie fp32 den Wert von pi darstellt.

Tabelle 4-2. fp32 repräsentiert pi

Unterschreibe	Exponent	Bruch (Mantisse/Signifikant)
1 Bit 0	8 Bits 10000000	23 Bits 10010010000111111011011

fp32 kann Zahlen in einem Bereich von -3e38 bis +3e38 darstellen. Der folgende PyTorch-Code zeigt, wie man die Datentypinformationen für fp32 ausgibt:

import torch
torch.finfo(torch.float32)

Die Ausgabe ist:

finfo(resolution=1e-06, min=-3.40282e+38, max=3.40282e+38, eps=1.19209e-07, 
smallest_normal=1.17549e-38, tiny=1.17549e-38, dtype=float32)

Wenn du eine reelle Zahl in 32 Bit speicherst, führt das zu einem leichten Verlust an Genauigkeit. Du kannst das sehen, indem du pi als fp32 Datentyp speicherst und dann den Wert des Tensors mit Tensor.item() auf 20 Dezimalstellen druckst:

pi = 3.14159265358979323846
pi_fp32 = torch.tensor(pi, dtype=torch.float32)
print('%.20f' % pi_fp32.item())

Die Ausgabe ist:

3.14159274101257324219

Du kannst den leichten Verlust an Genauigkeit erkennen, wenn du diesen Wert mit dem realen Wert von Pi vergleichst, der mit 3.14159265358979323846 beginnt. Dieser leichte Verlust an Genauigkeit ist auf die Umrechnung in den Zahlenbereich fp32 zurückzuführen, wie in Abbildung 4-3 dargestellt.

Abbildung 4-3. fp32 Projektion von pi in den Bereich von -3e38 bis +3e38

Du kannst auch den Speicherverbrauch ausdrucken:

def show_memory_comsumption(tensor):
    memory_bytes = tensor.element_size() * tensor.numel()
    print("Tensor memory consumption:", memory_bytes, "bytes")
show_memory_comsumption(pi_fp32)

Die Ausgabe ist:

Tensor memory consumption: 4 bytes

Nachdem du nun die Datentypen und numerischen Darstellungen kennengelernt hast, wollen wir nun erörtern, wie die Quantisierung dir helfen kann, den Speicherbedarf für das Laden und Trainieren deines Modells mit mehreren Milliarden Parametern zu reduzieren.

Quantisierung

Wenn du versuchst, ein Modell mit mehreren Milliarden Parametern mit voller 32-Bit-Präzision zu trainieren, stößt du schnell an die Grenzen eines einzelnen NVIDIA A100 oder H100 Grafikprozessors mit nur 80 GB GPU-RAM. Deshalb musst du bei der Verwendung eines einzelnen Grafikprozessors fast immer eine Quantisierung durchführen.

Die Quantisierung reduziert den Speicherbedarf beim Laden und Trainieren eines Modells, indem sie die Genauigkeit der Modellgewichte verringert. Durch die Quantisierung werden die Modellparameter von einer 32-Bit-Präzision auf eine 16-Bit-Präzision umgerechnet - oder sogar auf 8- oder 4-Bit.

Wenn du deine Modellgewichte von 32-Bit-Vollpräzision auf 16-Bit- oder 8-Bit-Präzision quantisierst, kannst du den Speicherbedarf deines 1-Milliarde-Parameter-Modells schnell um 50 % auf nur 2 GB oder sogar um 75 % auf nur 1 GB reduzieren, wie in Abbildung 4-4 gezeigt.

Abbildung 4-4. Ungefährer GPU-RAM-Bedarf zum Laden eines Modells mit 1 Milliarde Parametern bei 32-Bit-, 16-Bit- und 8-Bit-Präzision

Die Quantisierung projiziert eine Quellmenge von Gleitkommazahlen höherer Genauigkeit in eine Zielmenge von Zahlen geringerer Genauigkeit. Anhand der Quell- und Zielbereiche berechnet der Mechanismus der Quantisierung zunächst einen Skalierungsfaktor, führt die Projektion durch und speichert dann die Ergebnisse in reduzierter Genauigkeit, wodurch weniger Speicherplatz benötigt wird und letztendlich die Trainingsleistung verbessert und die Kosten gesenkt werden.

fp16

Bei fp16 bestehen die 16 Bits aus 1 Bit für das Vorzeichen, aber nur 5 Bits für den Exponenten und 10 Bits für den Bruch, wie in Tabelle 4-3 gezeigt.

Tabelle 4-3. fp32 versus fp16

	Unterschreibe	Exponent	Bruch (Mantisse/Signifikant)
fp32 (verbraucht 4 Byte Speicherplatz)	1 Bit 0	8 Bits 10000000	23 Bits 10010010000111111011011
fp16 (verbraucht 2 Byte Speicherplatz)	1 Bit 0	5 Bits 10000	10 Bits 1001001000

Mit der reduzierten Anzahl von Bits für den Exponenten und den Bruch liegt der Bereich der darstellbaren fp16 Zahlen nur noch zwischen -65.504 und +65.504. Das kannst du auch sehen, wenn du die Datentypinformationen für fp16 ausdruckst:

torch.finfo(torch.float16)

Die Ausgabe ist:

finfo(resolution=0.001, min=-65504, max=65504, eps=0.000976562, 
smallest_normal=6.10352e-05, tiny=6.10352e-05, dtype=float16)

Speichern wir pi wieder mit 20 Nachkommastellen in fp16 und vergleichen wir die Werte:

pi = 3.14159265358979323846
pi_fp16 = torch.tensor(pi, dtype=torch.float16)
print('%.20f' % pi_fp16.item())

Die Ausgabe ist:

3.14062500000000000000

Beachte den Präzisionsverlust nach dieser Hochrechnung, da es jetzt nur noch sechs Nachkommastellen gibt. Der fp16 Wert von pi ist jetzt 3.140625. Erinnere dich daran, dass du schon allein durch das Speichern des Wertes in fp32 an Genauigkeit verloren hast, wie in Abbildung 4-5 gezeigt.

Abbildung 4-5. Quantisierung von fp32 auf fp16 spart 50% Speicherplatz

Der Verlust an Genauigkeit ist jedoch in den meisten Fällen akzeptabel. Die Vorteile eines um 50 % geringeren GPU-Speichers für fp16 im Vergleich zu fp32 sind den Kompromiss in der Regel wert, da fp16 nur 2 Byte Speicher gegenüber 4 Byte bei fp32 benötigt.

Das Laden eines Modells mit 1 Milliarde Parametern erfordert jetzt nur noch 2 GB GPU-RAM, während für das Training des Modells 12 GB GPU-RAM benötigt werden, wie in Abbildung 4-6 dargestellt.

Abbildung 4-6. Nur 12 GB GPU-RAM werden benötigt, um ein Modell mit 1 Milliarde Parametern bei 16-Bit halber Genauigkeit zu laden und zu trainieren

bfloat16

bfloat16 hat sich zu einer beliebten Alternative zu fp16 entwickelt, da es den gesamten Bereich von fp32 mit nur 16 Bits abdeckt. Dadurch werden numerische Instabilitäten während des Modelltrainings, die durch Überlauf verursacht werden, reduziert. Ein Überlauf tritt auf, wenn Zahlen bei der Umwandlung von einem hochpräzisen in einen niedrigpräzisen Bereich außerhalb des Darstellungsbereichs fließen und NaN (keine Zahl) Fehler verursachen.

Im Vergleich zu fp16 hat bfloat16 einen größeren Dynamikbereich, aber eine geringere Genauigkeit, was normalerweise akzeptabel ist. bfloat16 verwendet ein einziges Bit für das Vorzeichen und die vollen 8 Bits für den Exponenten. Allerdings wird der Bruch auf 7 Bits abgeschnitten, weshalb er oft als "abgeschnittene 32-Bit-Fließkommazahl" bezeichnet wird, wie in Tabelle 4-4 gezeigt.

Tabelle 4-4. fp32 versus bfloat16

	Unterschreibe	Exponent	Bruch (Mantisse/Signifikant)
fp32 (verbraucht 4 Byte Speicherplatz)	1 Bit 0	8 Bits 10000000	23 Bits 10010010000111111011011
bfloat16 (verbraucht 2 Byte Speicherplatz)	1 Bit 0	8 Bits 10000000	7 Bits 1001001

Der Bereich der darstellbaren Zahlen von bfloat16 ist identisch mit fp32. Drucken wir die Datentypinformationen für bfloat16 aus:

torch.finfo(torch.bfloat16)

Die Ausgabe ist:

finfo(resolution=0.01, min=-3.38953e+38, max=3.38953e+38, eps=0.0078125, 
smallest_normal=1.17549e-38, tiny=1.17549e-38, dtype=bfloat16)

Speichern wir pi wieder mit 20 Nachkommastellen in bfloat16 und vergleichen wir die Werte:

pi = 3.14159265358979323846
pi_bfloat16 = torch.tensor(pi, dtype=torch.bfloat16)
print('%.20f' % pi_bfloat16.item())

Die Ausgabe ist:

3.14062500000000000000

Ähnlich wie bei fp16 geht bfloat16 mit einem minimalen Verlust an Genauigkeit einher. Der bfloat16 Wert von pi ist 3.140625. Die Vorteile der Beibehaltung des dynamischen Bereichs von fp32 (siehe Abbildung 4-7) und der damit verbundenen Reduzierung des Überlaufs überwiegen jedoch in der Regel den Verlust an Genauigkeit.

Abbildung 4-7. Durch die Quantisierung von fp32 auf bfloat16 wird der Dynamikbereich von fp32 beibehalten und trotzdem 50% Speicherplatz eingespart

bfloat16 wird von neueren Grafikprozessoren wie NVIDIAs A100 und H100 nativ unterstützt. Viele moderne generative KI-Modelle wurden mit bfloat16 vortrainiert, darunter FLAN-T5, Falcon und Llama 2.

fp8

fp8 ist ein neuerer Datentyp und eine natürliche Weiterentwicklung von fp16 und bfloat16, um den Speicher- und Rechenaufwand für Modelle mit mehreren Milliarden Parametern weiter zu reduzieren.

fp8 ermöglicht es dem Benutzer, die Anzahl der Bits, die dem Exponenten und dem Bruch zugewiesen werden, je nach Aufgabe zu konfigurieren, z. B. für Training, Inferenz oder Quantisierung nach dem Training. NVIDIAs Grafikprozessoren unterstützen fp8 seit dem H100 Chip. AWS Trainium unterstützt auch fp8, genannt konfigurierbares fp8, oder einfach cfp8. Bei cfp8 wird 1 Bit für das Vorzeichen verwendet, und die restlichen 7 Bits sind für den Exponenten und den Bruch konfigurierbar, wie in Tabelle 4-5 dargestellt.

Tabelle 4-5. fp32 versus fp8

	Unterschreibe	Exponent	Bruch (Mantisse/Signifikant)
fp32 (verbraucht 4 Byte Speicherplatz)	1 Bit 0	8 Bits 10000000	23 Bits 10010010000111111011011
fp8 (verbraucht 1 Byte Speicher)	1 Bit 0	7 Bits 0000011 (konfigurierbar)

Empirische Ergebnisse zeigen, dass fp8 mit der Leistung von fp16 und bfloat16 mithalten kann, während der Speicherbedarf um weitere 50 % reduziert und die Modellbildung beschleunigt wird.

int8

Eine weitere Quantisierungsoption ist int8 8-Bit-Quantisierung. Wenn du 1 Bit für das Vorzeichen verwendest, werden die Werte von int8 durch die restlichen 7 Bits dargestellt, wie in Tabelle 4-6 gezeigt.

Tabelle 4-6. fp32 versus int8

	Unterschreibe	Exponent	Bruch (Mantisse/Signifikant)
fp32 (verbraucht 4 Byte Speicherplatz)	1 Bit 0	8 Bits 10000000	23 Bits 10010010000111111011011
int8 (verbraucht 1 Byte Speicherplatz)	1 Bit 0	k.A.	7 Bits 0000011

Der Bereich der darstellbaren int8 Zahlen ist -128 bis +127. Hier sind die Informationen zum Datentyp für int8:

torch.iinfo(torch.int8)

Die Ausgabe ist:

iinfo(min=-128, max=127, dtype=int8)

Speichern wir pi wieder mit 20 Dezimalstellen in int8 und schauen, was passiert:

pi = 3.14159265358979323846
pi_int8 = torch.tensor(pi, dtype=torch.int8)
print(pi_int8.item())

Die Ausgabe ist:

3

Es überrascht nicht, dass pi im 8-Bit-Niedrigpräzisionsraum nur auf 3 projiziert wird, wie in Abbildung 4-8 zu sehen ist.

Abbildung 4-8. Quantisierung von fp32 auf int8 spart 75% Speicherplatz

Dadurch sinkt der Speicherbedarf von ursprünglich 4 Byte auf nur noch 1 Byte, aber die Umwandlung von einer Fließkommadarstellung in einen Ganzzahlwert führt zu einem größeren Verlust an Genauigkeit.

Die Verringerung des Speicherbedarfs großer Basismodelle ist nicht nur beim Laden und Trainieren der Modelle hilfreich, sondern auch bei der Inferenz. Trotz des Präzisionsverlusts wird die 8-Bit-Quantisierung häufig verwendet, um den Durchsatz und die Latenzzeit bei der Inferenz von eingesetzten Modellen zu verbessern. Optimierte Implementierungen für die int8 Quantisierung, wie z. B. die Integration von LLM.int8() durch Hugging Face ( bitsandbytes), haben gezeigt, dass die Auswirkungen der Quantisierung auf die Modellleistung minimiert werden. Du wirst mehr über die Post-Training-Quantisierung (PTQ) und die Technik GPT Post-Training-Quantisierung (GPTQ) erfahren¹ genauer kennen, wenn du das Modell in Kapitel 8 für den Einsatz vorbereitest.

Tabelle 4-7 vergleicht die bisher besprochenen Datentypen.

Tabelle 4-7. Vergleich der für die Quantisierung verwendeten Datentypen

	Total Bits	Vorzeichen Bits	Exponent Bits	Bruchteile von Bits	Benötigter Speicherplatz, um einen Wert zu speichern
fp32	32	1	8	23	4 Bytes
fp16	16	1	5	10	2 Bytes
bf16	16	1	8	7	2 Bytes
fp8	8	1	7		1 Byte
int8	8	1	k.A.	7	1 Byte

Zusammenfassend lässt sich sagen, dass die Wahl des Datentyps für die Modellquantisierung von den spezifischen Anforderungen deiner Anwendung abhängen sollte. Während fp32 eine sichere Wahl ist, wenn die Genauigkeit im Vordergrund steht, stößt du bei Modellen mit mehreren Milliarden Parametern wahrscheinlich an die Grenzen der Hardware, z. B. den verfügbaren GPU-RAM.

In diesem Fall kann die Quantisierung mit fp16 und bfloat16 dazu beitragen, den benötigten Speicherplatz um 50 % zu reduzieren. bfloat16 wird in der Regel gegenüber fp16 bevorzugt, da es den gleichen dynamischen Bereich wie fp32 beibehält und den Überlauf reduziert. fp8 ist ein aufstrebender Datentyp, der die Speicher- und Rechenanforderungen weiter reduziert. Einige Hardware-Implementierungen ermöglichen es, die Bits für Exponent und Bruch zu konfigurieren; empirische Ergebnisse zeigen, dass die Leistung mit fp16 und bfloat16 beim Modelltraining gleichwertig ist. int8 ist zu einer beliebten Wahl geworden, um dein Modell für Inferenzen zu optimieren. fp8 wird immer beliebter, da sowohl die Hardware als auch die Deep-Learning-Frameworks immer mehr Unterstützung bieten.

Tipp

Es wird empfohlen, dass du immer einen Benchmark der Quantisierungsergebnisse durchführst, um sicherzustellen, dass der gewählte Datentyp deinen Anforderungen an Genauigkeit und Leistung entspricht.

Eine weitere Speicher- und Rechenleistung Optimierungstechnik ist FlashAttention. FlashAttention zielt darauf ab, den quadratischen Rechen- und Speicherbedarf ( O(ⁿ²)) der Self-Attention-Schichten in Transformer-basierten Modellen zu reduzieren.

Optimierung der Selbstaufmerksamkeitsschichten

Wie in Kapitel 3 erwähnt, wird die Leistung des Transformers oft durch die Rechen- und Speicherkomplexität der Selbstbeobachtungsschichten eingeschränkt. Viele Leistungsverbesserungen sind speziell auf diese Schichten ausgerichtet. Im Folgenden lernst du einige leistungsstarke Techniken kennen, um den Speicher zu reduzieren und die Leistung der Self-Attention-Schichten zu erhöhen.

FlashAttention

Die Aufmerksamkeitsschicht des Transformers ist ein Engpass, wenn es darum geht, längere Eingabesequenzen zu skalieren, denn die Berechnungs- und Speicheranforderungen skalieren quadratisch O(ⁿ²) mit der Anzahl der Eingabe-Token. FlashAttention, ursprünglich in einem Forschungspapier vorgeschlagen,² ist eine GPU-spezifische Lösung für dieses quadratische Skalierungsproblem.

FlashAttention, das derzeit in Version 2 vorliegt, reduziert die Anzahl der Lese- und Schreibvorgänge zwischen dem GPU-Hauptspeicher, dem sogenannten High-Bandwidth-Memory (HBM), und dem viel schnelleren, aber kleineren statischen On-Chip-GPU-RAM (SRAM). Trotz seines Namens ist der GPU-Hochbreitenspeicher um eine Größenordnung langsamer als der On-Chip-GPU-SRAM.

Insgesamt steigert FlashAttention die Leistung von Self-Attention um das 2-4-fache und reduziert den Speicherbedarf um das 10-20-fache, indem es die quadratischen O(ⁿ²)-Rechen- und Speicheranforderungen auf lineares O(n) reduziert, wobei n die Anzahl der Eingabe-Token in der Sequenz ist. Mit FlashAttention kann der Transformer viel längere Eingabesequenzen verarbeiten, was zu einer besseren Leistung bei größeren Eingabekontextfenstern führt.

Eine beliebte Implementierung lässt sich mit einem einfachen pip install flash-attn --no-build-isolation Befehl installiert werden, der die flash-attn Bibliothek als Ersatz für die ursprüngliche Aufmerksamkeit installiert.

Aufmerksamkeitsoptimierungen sind ein aktiver Bereich der Forschung, einschließlich der nächsten Generation von FlashAttention-2,³ die weiterhin GPU-spezifische Optimierungen implementiert, um die Leistung zu verbessern und den Speicherbedarf zu reduzieren.

Lass uns eine weitere Technik kennenlernen, um die Leistung der Selbstbeobachtungsschichten im Transformer zu verbessern.

Gruppierte Abfrage Aufmerksamkeit

Eine weitere beliebte Optimierung der Aufmerksamkeitsschichten ist die gruppierte Abfrageaufmerksamkeit (GQA). GQA verbessert die traditionelle mehrköpfige Aufmerksamkeit des Transformers, die in Kapitel 3 beschrieben wurde, indem sie für jede Gruppe von Abfrageköpfen(q) einen einzigen Schlüssel-(k) und Wertkopf(v) verwendet (im Gegensatz zu jedem Abfragekopf), wie in Abbildung 4-9 dargestellt.

Abbildung 4-9. Aufmerksamkeit durch gruppierte Abfragen im Vergleich zu traditioneller mehrköpfiger Aufmerksamkeit

GQA ermöglicht es, Abfragen in weniger Schlüssel- und Wertköpfe zu gruppieren und so den Speicherverbrauch der Aufmerksamkeitsköpfe zu reduzieren. Darüber hinaus verbessert GQA die Leistung, indem es die Anzahl der Lese- und Schreibvorgänge im Speicher reduziert.

Da diese Verbesserungen proportional zur Anzahl der Eingabe-Token sind, ist MQA besonders nützlich für längere Eingabe-Token-Sequenzen und ermöglicht ein größeres Kontextfenster. Das Modell Llama 2 von Meta nutzt beispielsweise GQA, um die Leistung zu verbessern und die Größe des Kontextfensters für die Eingabe-Token auf 4.096 zu erhöhen - das Doppelte der ursprünglichen Größe des Kontextfensters von 2.048 des LLaMA-Modells.

Verteiltes Rechnen

Für größere Modelle wirst du wahrscheinlich einen verteilten Cluster von Grafikprozessoren verwenden müssen, um diese massiven Modelle auf Hunderten oder Tausenden von Grafikprozessoren zu trainieren. Es gibt viele verschiedene Arten von verteilten Berechnungsmustern, wie z. B. Distributed Data Parallel (DDP) und Fully Sharded Data Parallel (FSDP). Der Hauptunterschied besteht darin, wie das Modell auf die GPUs im System aufgeteilt - oder gesharded - wird.

Wenn die Modellparameter in einen einzigen Grafikprozessor passen, wählst du DDP, um eine einzelne Kopie des Modells in jeden Grafikprozessor zu laden. Wenn das Modell zu groß für einen einzelnen Grafikprozessor ist - auch nach der Quantisierung -, musst du FSDP verwenden, um das Modell auf mehrere Grafikprozessoren aufzuteilen. In beiden Fällen werden die Daten in Stapel aufgeteilt und auf alle verfügbaren GPUs verteilt, um die GPU-Auslastung und die Kosteneffizienz zu erhöhen, allerdings auf Kosten eines gewissen Kommunikationsaufwands, den du gleich sehen wirst.

Verteilte Daten Parallel

PyTorch wird mit einer optimierten Implementierung von DDP geliefert, die dein Modell automatisch auf jede GPU kopiert (vorausgesetzt, es passt mit einer Technik wie der Quantisierung auf eine einzelne GPU), die Daten in Stapel aufteilt und die Stapel parallel an jede GPU sendet. Bei DDP wird jeder Datenstapel parallel auf jedem Grafikprozessor verarbeitet, gefolgt von einem Synchronisationsschritt, bei dem die Ergebnisse der einzelnen GPUs (z. B. Gradienten) kombiniert (z. B. gemittelt) werden. Anschließend wird jedes Modell - eines pro GPU - mit den kombinierten Ergebnissen aktualisiert und der Prozess wird fortgesetzt, wie in Abbildung 4-10 dargestellt.

Abbildung 4-10. Verteilte Daten parallel (DDP)

Beachte, dass DDP davon ausgeht, dass jeder Grafikprozessor nicht nur deine Modellparameter und Datenbatches speichern kann, sondern auch die zusätzlichen Daten, die für die Trainingsschleife benötigt werden, einschließlich Optimierungszustände, Aktivierungen, temporäre Funktionsvariablen usw., wie in Abbildung 4-15 dargestellt. Wenn deine GPU nicht alle diese Daten speichern kann, musst du dein Modell auf mehrere GPUs verteilen. PyTorch verfügt über eine optimierte Implementierung des Modell-Shardings, die du im Folgenden kennenlernen wirst.

Vollständig gesplittete Daten Parallel

FSDP wurde motiviert durch ein ZeRO-Papier von 2019.⁴ Das Ziel von ZeRO (Zero Redundancy Optimizer) ist es, die Datenredundanz von DDP zu reduzieren, indem das Modell - und seine zusätzlichen Gradienten, Aktivierungen und Optimierungszustände - auf die GPUs verteilt werden, um eine Nullredundanz im System zu erreichen. ZeRO beschreibt drei Optimierungsstufen (1, 2, 3), je nachdem, was auf die GPUs verteilt wird (siehe Abbildung 4-11).

Abbildung 4-11. ZeRO besteht aus drei Stufen, die von den GPU-Shards abhängen: Parameter, Gradienten und Optimierungsstatus

ZeRO Stufe 1 verteilt nur die Optimierungszustände auf die GPUs, reduziert aber trotzdem den Speicherbedarf deines Modells um bis zu 4x. ZeRO Stufe 2 verteilt sowohl die Optimierungszustände als auch die Gradienten auf die GPUs, um den GPU-Speicher um bis zu 8x zu reduzieren. ZeRO Stufe 3 verteilt alles - einschließlich der Modellparameter - auf die GPUs, um den GPU-Speicher um das n-fache zu reduzieren, wobei n die Anzahl der GPUs ist. Wenn du ZeRO Stage 3 zum Beispiel mit 128 Grafikprozessoren verwendest, kannst du deinen Speicherverbrauch um das 128-fache reduzieren.

Im Vergleich zu DDP, bei dem jeder Grafikprozessor über eine vollständige Kopie aller Daten verfügt, die für den Vorwärts- und Rückwärtsdurchlauf benötigt werden, muss FSDP vor dem Vorwärts- und Rückwärtsdurchlauf eine vollständige Schicht aus den gesplitteten Daten auf jedem Grafikprozessor dynamisch rekonstruieren, wie in Abbildung 4-12 dargestellt.

Abbildung 4-12. FSDP über mehrere GPUs

In Abbildung 4-12 siehst du, dass jeder Grafikprozessor vor dem Vorwärtsdurchlauf bei Bedarf Daten von den anderen Grafikprozessoren anfordert, um die gesharten Daten für die Dauer der Operation in ungesharte, lokale Daten umzuwandeln - in der Regel auf einer Basis pro Schicht.

Wenn der Vorwärtsdurchlauf abgeschlossen ist, gibt die FSDP die nicht gesharten lokalen Daten wieder an die anderen GPUs frei und setzt die Daten in ihren ursprünglichen Sharded-Zustand zurück, um GPU-Speicher für den Rückwärtsdurchlauf freizugeben. Nach dem Backward-Pass synchronisiert FSDP die Gradienten zwischen den GPUs, ähnlich wie DDP, und aktualisiert die Modellparameter in allen Modell-Shards, wobei verschiedene Shards auf verschiedenen GPUs gespeichert sind.

Durch die Materialisierung der Daten bei Bedarf gleicht die FSDP den Kommunikations-Overhead mit dem Gesamtspeicherbedarf der GPU aus. Du kannst den Sharding-Faktor manuell über die Konfiguration des verteilten Rechnens konfigurieren. Später in diesem Kapitel wirst du ein Beispiel sehen, bei dem der Konfigurationsparameter sharded_data_parallel_degree von Amazon SageMaker verwendet wird. Diese Konfigurationseinstellung hilft dir, den Kompromiss zwischen Leistung und Speicherauslastung in Abhängigkeit von deiner spezifischen Umgebung zu verwalten, wie in Abbildung 4-13 dargestellt.

Abbildung 4-13. Wähle einen Sharding-Faktor basierend auf den Ressourcen in deiner Umgebung

Ein Sharding-Faktor von 1 vermeidet das Sharding des Modells und repliziert das Modell auf alle GPUs, wodurch das System wieder auf DDP zurückgesetzt wird. Du kannst den Sharding-Faktor auf maximal n Anzahl der GPUs einstellen, um das Potenzial des Full Sharding auszuschöpfen. Vollständiges Sharding bietet die besten Speichereinsparungen - allerdings auf Kosten des GPU-Kommunikations-Overheads. Wenn du den Sharing-Faktor auf einen Wert dazwischen einstellst, wird hybrides Sharding möglich.

Leistungsvergleich von FSDP gegenüber DDP

Abbildung 4-14 ist ein Vergleich von FSDP und DDP aus einem PyTorch FSDP Paper von 2023.⁵ Diese Tests wurden auf unterschiedlich großen T5-Modellen mit 512 NVIDIA A100-GPUs durchgeführt, die jeweils über 80 GB Speicher verfügen. Sie vergleichen die Anzahl der FLOPs pro GPU. Ein TeraFLOP ist 1 Billion Fließkommaoperationen pro Sekunde.

Abbildung 4-14. Leistungsverbesserung mit FSDP gegenüber DDP (Quelle: angepasst an ein Bild in Zhao et al.)

Beachte, dass volle Replikation bedeutet, dass es kein Sharding gibt. Und da die vollständige Replikation dem DDP entspricht, ist die Leistung der Konfigurationen mit vollständiger Replikation und DDP fast identisch.

Bei den kleineren T5-Modellen mit 611 Millionen Parametern und 2,28 Milliarden Parametern erbringt FSDP die gleiche Leistung wie DDP. Bei 11,3 Milliarden Parametern geht dem DDP jedoch der GPU-Speicher aus, weshalb es in der Dimension 11,3 Milliarden keine Daten für das DDP gibt. FSDP unterstützt jedoch problemlos die höhere Parametergröße, wenn Hybrid- und Full-Sharding verwendet werden.

Außerdem zeigt das Training des 11-Milliarden-Parameter-Modells mit verschiedenen Clustergrößen von 8 bis 512 GPUs nur einen Rückgang von 7 % der TeraFLOPs pro GPU aufgrund des GPU-Kommunikations-Overheads. Diese Tests wurden mit Stapelgrößen von 8 (blau) und 16 (orange) durchgeführt, wie in Abbildung 4-15 zu sehen ist, die ebenfalls aus dem PyTorch FSDP Paper 2023 stammt.

Abbildung 4-15. Nur geringer Leistungsabfall aufgrund des GPU-Kommunikations-Overheads (Quelle: angepasst an ein Bild in Zhao et al.)

Dies zeigt, dass FSDP die Modellschulung sowohl für kleine als auch für große Modelle über verschiedene GPU-Clustergrößen hinweg skalieren kann. Als Nächstes erfährst du, wie du mit Amazon SageMaker verteiltes Rechnen und FSDP auf AWS durchführst.

Verteiltes Rechnen auf AWS

Amazon SageMaker verteiltes Training wurde verwendet, um einige der leistungsfähigsten Basismodelle der Welt zu trainieren, darunter Falcon und BloombergGPT. Falcon-180B zum Beispiel wurde mit einem verteilten Amazon SageMaker-Trainingscluster mit 512 ml.p4d.24xlarge -Instanzen trainiert - jede mit 8 NVIDIA A100-GPUs (je 40 GB GPU-RAM), insgesamt also 4.096 GPUs und etwa 164 TB GPU-RAM. BloombergGPT wurde auf 64 Instanzen von ml.p4d.24xlarge mit insgesamt 512 GPUs und ca. 20 TB GPU-RAM trainiert.

Mit der verteilten Recheninfrastruktur von SageMaker kannst du mit nur wenigen Codezeilen hoch skalierbare und kostengünstige generative KI-Arbeitslasten ausführen. Als Nächstes lernst du, wie du FSDP mit Amazon SageMaker implementieren kannst.

Vollständig gesplittete Daten parallel mit Amazon SageMaker

FSDP ist eine gängige Strategie für verteiltes Rechnen die von Amazon SageMaker unterstützt wird. Der folgende Code zeigt, wie ein verteilter FSDP-Schulungsauftrag unter Verwendung des PyTorch Estimators mit 2 ml.p4d.24xlarge SageMaker-Instanzen gestartet wird - jede mit 8 GPUs und 320 GB GPU-RAM:

# Choose instance type and instance count 
# based on the GPU memory requirements 
# for the model variant we are using 
# e.g. Llama2 7, 13, 70 billion
instance_type = "ml.p4d.24xlarge" # 8 GPUs each
instance_count = 2 
# Set to the number of GPUs on that instance
processes_per_host = 8
# Configure the sharding factor
# In this case, 16 is the maximum, fully-sharded configuration
# since we have 2 instances * 8 GPUs per instance
sharding_degree = 16
# Set up the training job
smp_estimator = PyTorch(
  entry_point="train.py", # training script
  instance_type=instance_type,
  instance_count=instance_count,
  distribution={
      "smdistributed": {
          "modelparallel": {
              "enabled": True,
              "parameters": {
                  "ddp": True,
                  "sharded_data_parallel_degree":
                       sharding_degree
              }
          }
      },
      ...
  },
  ...
)

Hier konfigurierst du den Auftrag so, dass er smdistributed verwendet, wobei modelparallel.enabled und ddp auf True gesetzt werden. Dadurch wird der SageMaker-Cluster so konfiguriert, dass er die FSDP-Strategie für verteiltes Rechnen verwendet. Beachte, dass wir den Parameter sharded_data_parallel_degree auf 16 setzen, weil wir zwei Instanzen mit jeweils acht GPUs haben. Dieser Parameter ist unser Sharding-Faktor, wie im Abschnitt "Fully Sharded Data Parallel" beschrieben . Hier wählen wir Full Sharding, indem wir den Wert auf die Gesamtzahl der GPUs im Cluster setzen.

Als Nächstes folgen einige interessante Ausschnitte aus der train.py, auf die im vorherigen PyTorch Estimator Code verwiesen wurde. Der vollständige Code befindet sich im GitHub-Repository zu diesem Buch:

from transformers import AutoConfig, AutoModelForCausalLM
import smp # SageMaker distributed library

# Create FSDP config for SageMaker
smp_config = {
      "ddp": True,
      "bf16": args.bf16,
      "sharded_data_parallel_degree": args.sharded_data_parallel_degree,
}

# Initialize FSDP
smp.init(smp_config)

# Load HuggingFace model
model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
# Wrap HuggingFace model in SageMaker DistributedModel class
model = smp.DistributedModel(
      model
)

# Define the distributed training step
@smp.step
def train_step(model, input_ids, attention_mask, args):
  if args.logits_output:
      output = model(input_ids=input_ids,
          attention_mask=attention_mask, 
          labels=input_ids)
      loss = output["loss"]
  else:
      loss = model(input_ids=input_ids, 
          attention_mask=attention_mask, 
          labels=input_ids)["loss"]
  model.backward(loss)
  if args.logits_output:
      return output

    return loss

Als Nächstes erfährst du, wie du ein Modell auf AWS Trainium-Hardware trainieren kannst, die speziell für Deep Learning-Arbeitslasten entwickelt wurde. Dazu lernst du das AWS Neuron SDK kennen - sowie die Hugging Face Optimum Neuron Bibliothek, die das Hugging Face Transformers Ökosystem mit dem Neuron SDK integriert.

AWS Neuron SDK und AWS Trainium

Das AWS Neuron SDK ist die Entwicklerschnittstelle zu AWS Trainium. Die Optimum Neuron-Bibliothek von Hugging Face ist die Schnittstelle zwischen dem AWS Neuron SDK und der Transformers-Bibliothek. Das folgende Beispiel zeigt die Klasse NeuronTrainer aus der Optimum Neuron-Bibliothek, die die Klasse Transformers Trainer beim Training mit AWS Trainium einfach ersetzen kann:

from transformers import TrainingArguments
from optimum.neuron import NeuronTrainer

def train():
    model = AutoModelForCausalLM.from_pretrained(
        model_checkpoint)

    training_args = TrainingArguments(
        ... 
    )

    trainer = NeuronTrainer(
        model=model,
        args=training_args,
        train_dataset=...,
        eval_dataset=...
    )

    trainer.train()

Zusammenfassung

In diesem Kapitel hast du dich mit den rechnerischen Herausforderungen beim Training großer Basismodelle aufgrund von GPU-Speicherbeschränkungen befasst und gelernt, wie du durch Quantisierung Speicherplatz sparen, Kosten senken und die Leistung verbessern kannst.

Du hast auch gelernt, wie du das Modelltraining über mehrere GPUs und Knoten in einem Cluster skalieren kannst, indem du verteilte Trainingsstrategien wie Distributed Data Parallel (DDP) und Fully Sharded Data Parallel (FSDP) einsetzt.

Durch die Kombination von Quantisierung und verteiltem Rechnen kannst du sehr große Modelle effizient und kostengünstig trainieren, mit minimalen Auswirkungen auf den Trainingsdurchsatz und die Modellgenauigkeit.

Du hast auch gelernt, wie du Modelle mit dem AWS Neuron SDK und der speziell für generative Deep Learning-Arbeitslasten entwickelten AWS Trainium-Hardware trainierst. Du hast gesehen, wie du die Hugging Face Optimum Neuron-Bibliothek verwendest, die in das AWS Neuron SDK integriert ist, um die Entwicklungserfahrung bei der Arbeit mit AWS Trainium zu verbessern.

In Kapitel 5 erfährst du, wie du bestehende generative Basismodelle mit einer Technik namens Feintuning an deine eigenen Datensätze anpassen kannst. Die Feinabstimmung eines bestehenden Basismodells kann eine kostengünstigere und dennoch ausreichende Alternative zum Pretraining eines neuen Modells sein.

¹ Elias Frantar et al., "GPTQ: Accurate Post-Training Quantization for Generative Pre-Trained Transformers", arXiv, 2023.

² Tri Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness", arXiv, 2022.

³ Tri Dao, "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning", arXiv, 2023.

⁴ Samyam Rajbhandari et al., "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models", arXiv, 2020.

⁵ Yanli Zhao et al., "PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel", arXiv, 2023.