Leistungsoptimierungen für TPU7x (Ironwood)

In diesem Leitfaden werden mehrere Methoden zur Leistungsoptimierung mit TPU7x (Ironwood) beschrieben, indem die Datenübertragung zwischen dem mehrstufigen Speichersystem effizient verwaltet wird. Dazu gehören Techniken wie Training mit niedriger Genauigkeit, Sharding, Kommunikationsoptimierung, Aktivierungsrematerialisierung, Scoped Vector Memory (VMEM)-Tuning und benutzerdefinierte Beschleuniger-Kernels.

Wenn Sie die Leistung mit TPU7x optimieren möchten, müssen Sie sich zuerst mit der Ironwood-Architektur vertraut machen, insbesondere mit der Speicherhierarchie und der Interconnect-Topologie. Weitere Informationen finden Sie unter TPU7x (Ironwood).

Training mit niedriger Genauigkeit mit FP8

FP8 (8-Bit-Gleitkomma) ist ein effizientes numerisches Datenformat, das hauptsächlich zur Beschleunigung von Modelltraining und ‑inferenz verwendet wird. Durch die Darstellung von Zahlen mit 8 Bit anstelle der standardmäßigen 16-Bit-Formate (FP16 oder BF16) und 32-Bit-Formate (FP32) können TPUs Daten deutlich schneller verarbeiten und weniger Speicher verbrauchen.

TPU7x unterstützt die integrierte Hardwarebeschleunigung für FP8-Datentypen und bietet eine theoretische Spitzenleistung von 4.614 TFLOPS pro Chip. Diese Funktion kann zu deutlich schnelleren End-to-End-Trainingszeiten führen. Bei kompatiblen Vorgängen, insbesondere bei dichten Matrixmultiplikationen, die bei KI-Arbeitslasten üblich sind, kann die Verwendung von FP8 zu einer Leistungssteigerung um das 1,3-Fache im Vergleich zum Standard-BF16-Training führen. Im Vergleich zu BF16 verdoppelt FP8 die Spitzen-FLOPs und halbiert den Speicherbedarf für Gewichtungen und Aktivierungen. FP8 sollte ein primärer Optimierungshebel für rechenintensive Arbeitslasten und Szenarien sein, die durch Speicherkapazität oder Bandbreite eingeschränkt sind.

Die Verwendung von FP8 bietet die folgenden Leistungsvorteile:

• Geringere Belastung des High-Bandwidth Memory (HBM): Durch einen kleineren Speicherbedarf können größere Modelle oder Modelle mit größeren KV-Caches während der Inferenz vollständig in die 192 GiB HBM passen. So wird ein kostspieliges Offloading in den langsameren Hostspeicher vermieden.
• Erhöhte effektive Batchgröße: Durch die Reduzierung des für Aktivierungen erforderlichen Arbeitsspeichers ermöglicht FP8 die Verwendung größerer Batchgrößen. Dadurch wird die Datenparallelität verbessert und es kann zu einem höheren Durchsatz und einer besseren Auslastung der Recheneinheiten kommen.
• Geringere Anforderungen an die Speicherbandbreite: Wenn für jeden Vorgang nur die Hälfte der Datenmenge verschoben wird, sinken die Anforderungen an den Datenpfad von HBM zu MXU. Auf Systemen, auf denen die Datenverschiebung ein häufiger Engpass ist, trägt dies dazu bei, dass die MXUs ausgelastet sind.

Die Verwendung von FP8 mit keiner oder nur geringer Leistungseinbuße erfordert eine sorgfältige Auswahl der Quantisierungstechniken. Hier sind einige Best Practices für das FP8-Training:

• Skalierungsgranularität: Beginnen Sie mit der Skalierung pro Tensor als Referenz. Wenn es Qualitäts- oder Leistungsprobleme gibt, wechseln Sie zur Skalierung pro Achse. Eine Skalierung von Unterkanälen ist möglicherweise nicht erforderlich.
• Skalierungsmodus: Die dynamische Skalierung, bei der Skalierungsfaktoren zur Laufzeit berechnet werden, ist eine gute Standardeinstellung, um die Qualität beizubehalten. Die statische Skalierung kann durch das Eliminieren von Berechnungen eine erhebliche Leistungssteigerung bieten, erfordert jedoch eine sorgfältige Profilerstellung, um die richtigen Skalierungsfaktoren zu bestimmen. Sie ist möglicherweise nicht für alle Anwendungsfälle geeignet, insbesondere wenn sich die Modellkonfigurationen ändern. Umgekehrt können einige robuste Modelle und Konfigurationen die Skalierung auf das FP8-Limit für Gewichtungen oder Aktivierungen festlegen. So können Sie den Quantisierungsaufwand reduzieren, während die Genauigkeit beibehalten und die Leistung verbessert wird.
• FP8-Formate (E4M3 und E5M2): Ein gängiger und effektiver Ansatz ist die Verwendung einer Mischung aus FP8-Formaten. Verwenden Sie beispielsweise E4M3 für Gewichtungen und Aktivierungen im Vorwärtsdurchlauf, um die höhere Genauigkeit von E4M3 zu nutzen, und E5M2 für Gradienten im Rückwärtsdurchlauf, um den größeren dynamischen Bereich der Gradienten zu berücksichtigen.
• Rundung: Wenn Sie „Rundung auf die nächste gerade Zahl“ (RNE) anstelle von stochastischer Rundung für Gradienten verwenden, können Sie die Qualität beibehalten und gleichzeitig eine bessere Leistung und Reproduzierbarkeit erzielen.
• FP8 in MaxText aktivieren: MaxText unterstützt FP8-Training über die QWIX-Quantisierungsbibliothek. Wenn Sie die Quantisierung aktivieren möchten, legen Sie in Ihrer Konfiguration das folgende Flag fest: use_qwix_quantization=true.

Fragmentierung und Parallelität

Beim Sharding wird ein großes Modell oder seine Trainingsdaten in kleinere Teile aufgeteilt und auf mehrere TPU-Chips oder -Kerne verteilt. Die Wahl der richtigen Sharding-Strategie ist wichtig, um eine hohe Leistung auf TPU7x zu erzielen.

Ein naiver Ansatz, der nur den Grad der Parallelität maximiert, führt oft zu einer schlechten Leistung, da er kommunikationsgebunden ist. Der beste Ansatz ist oft, die einfachste Sharding-Strategie auszuwählen, die die Arbeitsspeicherbeschränkungen erfüllt, da so der Kommunikationsaufwand minimiert wird und die Recheneinheiten effizient genutzt werden können.

Bevor Sie eine Sharding-Strategie auswählen, sollten Sie zuerst eine Analyse der arithmetischen Intensität durchführen. Mit dieser Analyse wird ermittelt, ob eine bestimmte Berechnung durch die Rechenleistung, die Speicherbandbreite oder die Interconnect-Bandbreite eingeschränkt ist. Sie wird als Verhältnis von Gleitkommaoperationen zu den zu verschiebenden Datenbyte berechnet.

Eine hohe arithmetische Intensität weist auf eine rechengebundene Arbeitslast hin. Eine niedrige arithmetische Intensität deutet auf eine speicher- oder kommunikationsgebundene Arbeitslast hin, bei der die Leistung durch die Geschwindigkeit begrenzt ist, mit der Daten aus dem HBM oder über das ICI-Netzwerk verschoben werden können. Diese Analyse liefert Informationen zur idealen Batchgröße und Sharding-Strategie. Eine kommunikationsgebundene Arbeitslast profitiert beispielsweise nicht von einer Sharding-Strategie, die noch mehr Kommunikation einführt, z. B. Tensorparallelität mit hohem Grad.

Entscheidungs-Framework für die Sharding-Strategie

MaxText bietet eine Vielzahl von Fragmentierungsstrategien. Die optimale Wahl hängt von der Modellarchitektur, der Sequenzlänge und der Notwendigkeit ab, die Rechenlast mit dem Kommunikationsoverhead in Einklang zu bringen.

• Fully Sharded Data Parallelism (FSDP): Dies ist die bevorzugte Standardstrategie für die Datenparallelität. Bei FSDP werden die Modellgewichtungen, Gradienten und Optimierungszustände über die datenparallelen Geräte hinweg fragmentiert. Während der Berechnung führt jedes Gerät einen All-Gather-Vorgang aus, um die erforderlichen vollständigen Gewichtungen für seinen lokalen Microbatch abzurufen. FSDP ist sehr effektiv, solange die Batchgröße pro Gerät groß genug ist, um die Latenz dieser All-Gather-Kommunikation zu verbergen. Bei MoE-Modellen (Mixture of Experts) muss die Berechnung der arithmetischen Intensität die dünne Besetzung berücksichtigen.
• Tensor-Parallelität (TP): Bei der Tensor-Parallelität werden einzelne Tensoren auf mehrere Geräte aufgeteilt. In der Regel sind die Tensoren Gewichtungsmatrizen in mehrschichtigen Perzeptronen (MLP) und Aufmerksamkeitsblöcken. Die hohe arithmetische Intensität der Hardware (11.500) stellt sehr hohe Anforderungen an die Dimensionen des Modells, damit TP über ICI sinnvoll ist. Der Versuch, TP zu verwenden, kann dazu führen, dass das System kommunikationsgebunden ist.
• Expert Parallelism (EP): Dies ist die Standard- und erforderliche Strategie für das Training von MoE-Modellen. EP fragmentiert die „Experten“-Ebenen auf eine Reihe von Geräten und ein All-to-All-Kommunikationskollektiv wird verwendet, um Tokens an das dafür vorgesehene Expertengerät weiterzuleiten. EP kann effizient sein, wenn die MLP-Dimension des Modells groß genug ist, um sich der Roofline zu nähern.
• Kontextparallelität (CP): CP ist eine spezielle Strategie, die für das Training von Modellen mit sehr langen Sequenzlängen unerlässlich ist. Die Hauptfunktion besteht darin, den Arbeitsspeicherverbrauch von Aktivierungen zu verwalten, der quadratisch mit der Sequenzlänge wächst und die HBM-Kapazität überschreiten kann. CP fragmentiert die Sequenzdimension der Aktivierungstensoren, was die Verwendung einer fraktionalen Batchgröße pro Gerät ermöglicht. Da CP mehr Kommunikation als FSDP erfordert, gilt die allgemeine Regel, den minimalen CP-Grad zu verwenden, der erforderlich ist, um die Speicherbeschränkungen zu erfüllen und dafür zu sorgen, dass der Batchachsen-Shard eine Ganzzahl bleibt.

In der folgenden Tabelle werden gängige Arbeitslasttypen der optimalen Sharding-Strategie zugeordnet:

Arbeitslasttyp	Empfohlenes primäres Sharding	Sekundäre Fragmentierung	Wichtige Engpässe	Begründung
Dichtes Modell – kurze Sequenz	FSDP	–	Rematerialisierung, FF-Matmuls	FSDP bietet die beste Balance. Bei kurzen Sequenzen ist der Aktivierungsspeicher möglicherweise kein großes Problem. Der Schlüssel ist ein ausreichend großer globaler Batch, um die All-Gather-Gewichtung von FSDP zu verbergen. Mit zunehmender Batchgröße nimmt auch die Aktivierungsgröße zu. Daher ist eine geeignete Rematerialisierungsrichtlinie erforderlich, damit dieser Konfiguration nicht der Arbeitsspeicher ausgeht.
Dichtes Modell – lange Sequenz	FSDP	CP	Flash Attention, Aktivierungsspeicher	Der Aktivierungsspeicher wird zur primären Einschränkung. CP ist erforderlich, um fraktionale Batchgrößen pro Gerät zu aktivieren und Probleme mit fehlendem Speicher (OOMs) zu vermeiden. Flash Attention ist die Hauptquelle für Rechenleistung und Zeitverschwendung.
MoE-Modell – kurze Sequenz	FSDP + EP	–	All-to-All (Expertenrouting), Rematerialisierung	MoE-Modelle erfordern EP, um die Experten zu fragmentieren. Die All-to-All-Kommunikation für das Token-Routing ist ein großer Engpass, der überlappt werden muss. Die Rematerialisierung ist auch eine erhebliche Quelle für Verschwendung.
MoE-Modell – sehr groß	FSDP + EP + PP	Modellparallelität (MP)	Alle zuvor erwähnten Engpässe sowie Pipeline-Bubbles	Für Modelle, die den Arbeitsspeicher eines einzelnen Pods überschreiten, ist PP erforderlich, um Ebenen auf Pods aufzuteilen. Dadurch entstehen Overheads für die DCN-Kommunikation und Pipeline-Blasen. Dies ist eine hochkomplexe Konfiguration, die eine sorgfältige Abstimmung erfordert.

Kommunikationsoptimierung

Der primäre Mechanismus für die Überlappung von Kommunikation und Berechnung auf TPU7x wird als SparseCore Collective Offloading bezeichnet. Die Ironwood-Architektur umfasst dedizierte SparseCore-Einheiten, die als unabhängige Kontrollthreads fungieren und die Datenübertragung über die ICI-Struktur verwalten können. Dadurch können kollektive Kommunikationsvorgänge (wie All-Gather oder Reduce-Scatter) parallel zu den Hauptberechnungen auf den TensorCores ausgeführt werden. Dies ist die empfohlene Methode für asynchrone Kollektive auf TPU7x. Verwenden Sie die empfohlenen Flags, um die Auslagerung für die gängigsten Kollektive zu aktivieren.

Aktivierungs-Rematerialisierung

Die Aktivierungsrematerialisierung, auch als Gradienten-Prüfpunkterstellung bezeichnet, ist eine grundlegende Technik zur Reduzierung des HBM-Speicherbedarfs eines Modells. Anstatt alle Zwischenaktivierungen aus dem Vorwärtsdurchlauf im HBM zu speichern, um sie während des Rückwärtsdurchlaufs zu verwenden, werden nur einige wichtige Aktivierungen (Prüfpunkte) gespeichert und die anderen bei Bedarf während des Rückwärtsdurchlaufs neu berechnet. Dadurch wird eine erhebliche Menge an Arbeitsspeicher eingespart, allerdings auf Kosten einer höheren Rechenleistung (ca. 25–30% zusätzliche FLOPs für einen Standard-Transformer-Block).

Die Entscheidung, wie aggressiv die Rematerialisierung angewendet werden soll, ist ein wichtiger Abstimmungsparameter, der vollständig vom primären Engpass abhängt, der oft mit der Sequenzlänge variiert.

Bei Arbeitslasten mit langen Sequenzen (z. B. 128k): In diesen Fällen ist die Größe der Aktivierungstensoren der Hauptverbraucher von HBM. Die Arbeitslast ist in der Regel speichergebunden. Daher ist die Anwendung einer aggressiven Rematerialisierungsrichtlinie sehr vorteilhaft. Durch die Speichereinsparungen kann das Training ohne Fehler aufgrund mangelnden Arbeitsspeichers fortgesetzt werden. Außerdem sind größere Batchgrößen möglich. Der Rechenaufwand für die Neuberechnung ist ein lohnender Kompromiss.

Bei Arbeitslasten mit kurzen Sequenzen (z. B. 8k): In diesen Fällen ist der Aktivierungsspeicher viel weniger ein Problem und die Arbeitslast ist eher rechengebunden. Der Rechenaufwand für die Rematerialisierung kann die größte Quelle für Ineffizienz sein.

Rematerialisierungsrichtlinien in MaxText optimieren

MaxText bietet eine detaillierte Kontrolle über die Rematerialisierung durch eine Reihe vordefinierter und benutzerdefinierter Richtlinien, die mit dem Flag remat_policy konfiguriert werden.

Voreingestellte Richtlinien

MaxText bietet die folgenden integrierten Richtlinien:

• full: Die aggressivste Richtlinie, bei der fast alles rematerialisiert wird. Dadurch wird die HBM-Nutzung minimiert, aber der Aufwand für die Neuberechnung maximiert. Ideal für Szenarien mit extrem begrenztem Arbeitsspeicher und langen Sequenzen.
• minimal: Die am wenigsten aggressive Richtlinie, bei der die meisten Aktivierungen gespeichert werden. Dadurch wird die HBM-Nutzung maximiert, aber die Neuberechnung minimiert. Am besten für rechenintensive Arbeitslasten mit kurzen Sequenzen geeignet, bei denen der Arbeitsspeicher keine Rolle spielt.
• Zwischenrichtlinien: Optionen wie save_dot_with_context_except_mlp, save_qkv_proj und save_out_proj bieten verschiedene Kompromisse, indem sie die Ausgaben teurer Punktproduktvorgänge selektiv prüfen, während günstigere elementweise Vorgänge rematerialisiert werden.

Benutzerdefinierte Richtlinien

Wenn Sie mehr Kontrolle wünschen, können Sie remat_policy auf custom festlegen. So können Sie das Verhalten für einzelne Ebenen im Decodierungsmodul des Modells festlegen. Jeder Ebene kann eines von drei Verhaltensweisen zugewiesen werden:

• device: Die Aktivierung wird im HBM auf dem TPU-Gerät gespeichert.
• remat: Die Aktivierung wird verworfen und während des Rückwärtsdurchlaufs neu materialisiert.
• offload: Die Aktivierung wird vom HBM in den Arbeitsspeicher des CPU-Hosts verschoben. Dadurch wird HBM freigegeben, was jedoch zu einer PCIe-Übertragungslatenz führt.

Bereichsbezogene VMEM-Abstimmung

Die Kernel-Leistung, z. B. Flash Attention, hängt von den ausgewählten Tile-Größen im Kernel ab, deren Größe durch den verfügbaren Vektorspeicher (VMEM) begrenzt ist. Jeder der beiden TensorCores in einem TPU7x-Chip hat 64 MiB Vektorspeicher (VMEM). Diese VMEM-Kapazität kann zwischen dem aktuellen Bereich (bereichsbezogener VMEM) und dem zukünftigen Gewichtungs-Prefetch aufgeteilt werden. Durch Erhöhen des VMEM-Bereichs können die Tile-Größen im Kernel erhöht werden, was möglicherweise Speicherstaus reduziert und die Leistung der Kernel erhöht. Sie können die Größe des VMEM-Bereichs ändern, indem Sie xla_tpu_scoped_vmem_limit_kib (in LIBTPU_INIT_ARGS) festlegen. Damit können Sie die Kernel-Leistung sowie die End-to-End-Leistungsgrenzen untersuchen. Die Optimierung der VMEM-Größe kann sich indirekt auf die Leistung des benutzerdefinierten Pallas-Kernels auswirken, da durch die Erhöhung der VMEM-Größe ein größerer Hyperparameter-Suchbereich für die Tile-Größen im Kernel freigeschaltet wird.

Tokamax-Kernel

Tokamax, eine leistungsstarke JAX-Kernel-Bibliothek mit vielen hochoptimierten TPU-Kerneln, behebt mehrere gängige hardwarespezifische Engpässe:

• Splash Attention: Splash Attention wird als primäre Attention-Implementierung verwendet, um den HBM-Engpass der Standard-Attention zu beseitigen. Dabei wird die effizienteste Attention-Implementierung auf TPUs verwendet.
• Megablox Grouped Matrix Multiplication (GMM): Bei MoE-Arbeitslasten verarbeitet Megablox gruppierte Matrixmultiplikationen effizient, indem Berechnungen über die unregelmäßige Aktivierungsdarstellung durchgeführt werden. Sie wird effizient über die unregelmäßige Dimension abgebildet und berechnet Matrixmultiplikationen zwischen unregelmäßigen Zeilengruppen in LHS und der entsprechenden Expertenmatrix. So müssen Batches nicht auf eine feste Größe aufgefüllt werden.
• Empirische Abstimmung mit tune-jax: Die tune-jax-Bibliothek enthält Dienstprogramme, mit denen Sie empirische Suchen nach optimalen Blockgrößen durchführen können. Die Standardkernelgrößen sind oft nicht optimal. Durch die Abstimmung können hardwarefreundliche VMEM-Tile-Größen ausgewählt werden, um die Hardwareauslastung zu maximieren.
• Max. Logits-Schätzung: Der Tokamax Splash Attention-Kernel kann weiter optimiert werden, indem ein Wert für max_logit_const festgelegt wird. Wenn festgelegt, wird die Reduktionsberechnung des maximalen Logits während des Softmax-Vorgangs der Aufmerksamkeit (softmax(Q * KT)) ersetzt, wodurch ein Teil des Rechen- und Synchronisierungsaufwands reduziert wird. In MaxText wird dies durch die Konfiguration use_max_logits_estimate implementiert, die auf None (deaktiviert) oder einen Gleitkommawert festgelegt werden kann. Prüfen Sie, ob der Logit-Bereich Ihres spezifischen Modells mit der Schätzung kompatibel bleibt, um einen numerischen Überlauf zu verhindern. Wenn dieser Wert festgelegt ist, wird ein Konvergenztest empfohlen.