Optimierung von PyTorch-Modellen durch Profiling und Kernel-Fusion

Von nn.Linear zu einem Fused MLP: Tiefergehende Einblicke in das PyTorch Profiling

Im Bereich des Deep Learning ist die Optimierung der Modellleistung von entscheidender Bedeutung. Während der erste Teil dieser Serie eine Einführung in das PyTorch Profiling und die Analyse von grundlegenden Operationen wie torch.add(torch.matmul(x, w), b) bot, vertieft dieser Artikel die Analyse komplexerer Architekturen. Wir untersuchen die Leistungsmerkmale von nn.Linear-Layern und die Vorteile der Kernel-Fusion in einem Multilayer Perceptron (MLP) – sowohl durch den Einsatz von torch.compile als auch durch manuell optimierte Kernels.

Die Umstellung von Matmul-Add auf nn.Linear

nn.Linear ist ein fundamentaler Baustein in vielen neuronalen Netzen, der die Operation y = x @ w.T + b kapselt. Diese Operation beinhaltet eine Matrixmultiplikation und eine Bias-Addition. Um die Leistung zu analysieren, betrachten wir ein Beispiel mit nn.Linear, wobei bias=True gesetzt ist, um die Operationen aus dem ersten Teil der Serie zu emulieren.

Bei der Profilerstellung einer solchen nn.Linear-Operation zeigt sich eine interessante Beobachtung: Vor der eigentlichen Multiplikation (aten::addmm) wird eine Transponierungsoperation (aten::t) auf der CPU durchgeführt. Diese aten::t-Operation kopiert oder reorganisiert jedoch keine Daten. Stattdessen modifiziert sie lediglich die Metadaten des Tensors (Shape und Stride), um eine transponierte Sicht auf die Daten zu erzeugen, ohne einen Kernel auf der GPU zu starten. Dies ist ein wichtiges Detail, da es den Overhead auf der GPU minimiert.

Die Rolle des Epilogs: Warum es keine separaten Additions-Kernels gibt

Ein weiterer bemerkenswerter Punkt ist das Fehlen eines separaten aten::add-Kernels für die Bias-Addition. Dies liegt daran, dass die Bias-Addition in den Matrixmultiplikations-Kernel integriert wurde, ein Prozess, der als Epilog bezeichnet wird. Ein Epilog ist eine kleine Berechnung, die ein GEMM-Kernel (General Matrix Multiply) am Ende durchführt, bevor die Ergebnisse in den High Bandwidth Memory (HBM) geschrieben werden. Dies vermeidet unnötigen Speicherzugriff und macht die Operation effizienter. Wenn aten::linear feststellt, dass ein Bias übergeben wurde, wird aten::addmm(bias, x, weight) aufgerufen, welches einen cuBLAS GEMM-Kernel mit integrierter Bias-Addition verwendet. Die Addition erscheint daher nie als separater Kernel, da sie Teil des "Writebacks" des Matmul-Kernels ist.

Der Einfluss von torch.compile auf eine einzelne Linear-Schicht

Es stellt sich die Frage, ob torch.compile bei einer einzelnen nn.Linear-Schicht eine signifikante Verbesserung bewirken kann. Ein Vergleich der Profiler-Traces im Eager-Modus und im kompilierten Modus zeigt, dass der gleiche cuBLAS GEMM-Kernel auf der GPU und der gleiche aten::addmm-Operator auf der CPU verwendet werden. torch.compile fügt lediglich einige zusätzliche Zeilen im CPU-Bereich hinzu, die mit dem Kompilierungsprozess zusammenhängen. Dies deutet darauf hin, dass für eine einzelne GEMM-Operation mit Bias torch.compile nur wenig zu optimieren hat, da die Fusion bereits auf einer niedrigeren Ebene stattfindet. Der Hauptvorteil von torch.compile liegt in der Reduzierung des CPU-Overheads durch das Auflösen von Metadaten-Operationen wie Transpositionen zur Kompilierzeit, wodurch eine direkte aten::addmm-Aufruf mit vorab berechneten Strides ermöglicht wird.

Verständnis der Kernel-Layouts und Pre-Ops

Die Analyse der Kernel-Namen im Profiler-Trace ist entscheidend. Der Suffix _tn_ in Kernel-Namen wie cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8 gibt Aufschluss über das Layout der Eingaben. t steht für transponiert, n für nicht-transponiert. cuBLAS und CUTLASS kompilieren separate Kernel-Binärdateien für jede Kombination von Eingabelayouts. Der Dispatcher wählt den passenden Kernel basierend auf den Eingabe-Strides aus. Wenn im kompilierten Modus die Transpositionsoperation auf der CPU verschwindet, bedeutet dies nicht, dass die Transposition nicht stattfindet, sondern dass der Compiler die Strides zur Kompilierzeit ermittelt und den passenden Kernel direkt aufruft, wodurch CPU-Overhead reduziert wird.

Stacking von drei Linears: Das Multilayer Perceptron (MLP)

Um die Leistungsfähigkeit von torch.compile besser zu demonstrieren, betrachten wir ein Multilayer Perceptron (MLP) mit einer GeGLU-Aktivierung. Eine solche Architektur ist in der Praxis weit verbreitet und bietet mehr Möglichkeiten zur Optimierung.

Im Eager-Modus eines GeGLU-MLP erwarten wir drei aten::linear-Dispatches für die drei linearen Schichten sowie zwei punktweise Kernel-Launches für die GeLU-Aktivierung und die Multiplikation. Die Profiler-Traces bestätigen dies: Pro Forward-Pass werden fünf Kernel auf der GPU ausgeführt.

Zwei Arten von GEMM-Kernels

Interessanterweise zeigen die Profiler-Traces zwei unterschiedliche Typen von GEMM-Kernels, obwohl alle drei GEMMs die gleiche Anzahl an Floating Point Operations (FLOPs) aufweisen. Die down_proj-Schicht ist etwa 10% schneller als die gate_proj und up_proj-Schichten. Dies liegt daran, dass cuBLAS für unterschiedliche Shapes (N=768 statt 3072) einen anderen Tile-Algorithmus (128x256 mit einer tieferen stages_64x3-Pipeline) wählt, der eine bessere Wiederverwendung ermöglicht.

Was bewirkt torch.compile im MLP?

Beim Kompilieren des Forward-Passes des MLP mit torch.compile sind die Veränderungen deutlicher. Während im Eager-Modus jede nn.Linear-Schicht eine Kette von Dispatcher-Operationen durchläuft, werden diese im kompilierten Modus stark vereinfacht. Die einzelnen aten::linear-, aten::matmul- und Transpositions-Operationen werden zu drei nackten aten::mm-Aufrufen zusammengefasst. Die Kernel-Namen auf der GPU bleiben dabei identisch, was darauf hindeutet, dass die grundlegenden GEMM-Operationen unverändert bleiben.

Der fused Triton-Kernel

Der größte Gewinn durch torch.compile bei MLPs ist die Fusion der punktweisen Kernels. Die zwei separaten Kernels für GeLU und Multiplikation plus eine Reshape-Operation werden zu einem einzigen Triton-Kernel fusioniert: triton_poi_fused__unsafe_view_gelu_mul_0. Dieser Kernel wird von Inductor generiert und führt die Operationen effizienter aus, indem er Zwischenergebnisse in den Registern hält und somit den teuren Round-Trip durch den HBM vermeidet. Dies ist ein entscheidender Vorteil, da unnötiger Speicherverkehr eine Hauptursache für Leistungsengpässe ist.

Einsatz von handoptimierten Kernels

Neben torch.compile gibt es auch die Möglichkeit, manuell optimierte Kernels zu verwenden. Ein Beispiel hierfür sind die LigerGEGLUMLP-Kernels, die über die kernels-Bibliothek bezogen werden können. Diese Kernels sind von menschlichen Experten geschrieben und auf spezifische Hardware abgestimmt.

Vorteile der kernels-Bibliothek

Die kernels-Bibliothek vereinfacht den Einsatz optimierter Kernels erheblich, indem sie den Kompilierungsprozess von der lokalen Maschine verlagert. Die Kernels werden einmal in der CI für verschiedene Architekturen und Versionskombinationen vorkompiliert und als vorab erstellte, versionsgebundene Pakete bereitgestellt. Dies eliminiert Kompatibilitätsprobleme und gewährleistet, dass alle Nutzer die gleiche, stabile Kernel-Version verwenden. Darüber hinaus bieten diese Pakete Drop-in nn.Module-Implementierungen, die einen einfachen Austausch bestehender Module ermöglichen.

Warum getunte Kernels besser sein können

Manuell optimierte Kernels bieten zwei Hauptvorteile:

1. Integrierte Fusion: Die Fusion ist bereits im Kernel-Design verankert. Die LigerGELUMulFunction führt beispielsweise einen einzigen Triton-Kernel aus, der gelu(gate) * up in einem Durchgang berechnet. Dies erreicht die gleiche Effizienz wie torch.compile, jedoch ohne den Overhead von Dynamo-Guards, Kompilierungszeiten oder dem Risiko einer Neukompilierung.
2. Hardware-optimierte Launch-Parameter: Die Launch-Parameter der Kernels werden nicht zufällig gewählt, sondern basierend auf der Hardware und den Spaltenanzahlen spezifisch berechnet.

Es ist wichtig, die Kompromisse zu verstehen. Obwohl ein handgeschriebener Kernel auf den ersten Blick geringfügig langsamer erscheinen mag als ein von Inductor generierter Kernel (z.B. 92,8 µs vs. 89,4 µs), bietet er Robustheit gegenüber variablen Eingabeformen. torch.compile spezialisiert sich auf statische Formen. Bei Änderungen der Batch-Größe, Sequenzlänge oder Hidden-Dimension muss Dynamo neu kompilieren, was erneut Kompilierungskosten verursacht. Ein handgeschriebener Kernel kann hingegen mit einem Satz von Launch-Parametern für beliebige Formen ausgeführt werden, ohne Neukompilierung. Er verzichtet auf die letzten Mikrosekunden, die eine formspezifische Spezialisierung bieten würde, im Austausch für Flexibilität bei wechselnden Eingabeformen.

Fazit

Das Profiling von PyTorch-Modellen ist ein iterativer Prozess, der ein tiefes Verständnis der zugrunde liegenden Operationen und Optimierungstechniken erfordert. Die Analyse von nn.Linear-Layern und MLPs zeigt, dass sowohl PyTorch im Eager-Modus als auch torch.compile und manuell optimierte Kernels unterschiedliche Optimierungsansätze verfolgen und jeweils spezifische Vorteile bieten. Die Wahl der richtigen Technik hängt von den individuellen Anforderungen des Modells und den Leistungsprofilen ab.

Der wichtigste Rat, der aus dieser Analyse hervorgeht, ist, vor jeder Profiling-Sitzung eine Hypothese aufzustellen. Erraten Sie zuerst, was der Trace enthalten sollte, und suchen Sie dann nach Abweichungen. Jede Diskrepanz ist ein Hinweis auf ein potenzielles Optimierungspotenzial oder eine tiefere Einsicht in das Systemverhalten.

Zusammenfassung der Veränderungen

Setup	Was sich geändert hat	Was gleich geblieben ist
Eager nn.Linear	Baseline: Bias-Addition ist bereits in den GEMM-Epilog (`addmm`) integriert, es ist ein cuBLAS-Kernel, nicht eine Matmul plus eine Addition.	—
Kompiliertes nn.Linear	Einige CPU-Dispatch-Operationen (das aten::t View-Bookkeeping) verschwinden.	Derselbe einzelne cuBLAS GEMM-Kernel, Byte für Byte. Der Compiler hat nichts zu fusionieren.
Eager MLP	5 GPU-Kernels: 3 GEMMs + ein GeLU + eine Multiplikation. Das [8192, 3072]-Zwischenergebnis macht einen vollständigen Round-Trip durch den HBM.	Jeder GEMM ist immer noch derselbe bias-freie cuBLAS-Kernel wie ein eigenständiger Linear.
Kompiliertes MLP	GeLU + Multiplikation + Reshape kollabieren zu einem fusionierten Triton-Kernel; das Zwischenergebnis bleibt in Registern. Bezahlt Kompilierungs-Pre-Ops (Dynamo, Guards).	Die 3 GEMMs bleiben unberührt mit identischen cuBLAS-Kernel-Namen.
Liger MLP	Dieselbe Fusion, aber in einen handgeschriebenen Triton-Kernel mit hardwareoptimierten Launch-Parametern und ohne Dynamo, Guards oder Kompilierungs-Latenzzeiten integriert.	Die 3 GEMMs sind immer noch dieselben cuBLAS-Kernels.

Bibliografie

- "Profiling in PyTorch (Part 1): A Beginner's Guide to torch.profiler", Hugging Face Blog. - "Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP", Hugging Face Blog. - "torch.profiler", PyTorch Documentation. - "Profiling your PyTorch Module", PyTorch Tutorials. - "Empowering Models with Performance: The Art of Generalized Model Transformation Approach", PyTorch Blog. - "Quentin-Anthony/torch-profiling-tutorial", GitHub Repository. - "torch/csrc/jit/codegen/cuda/README.md at main · pytorch/pytorch", GitHub Repository. - "torch.compile End-to-End Tutorial", PyTorch Tutorials. - "Introduction to torch.compile", PyTorch Tutorials. - "PyTorch Profiler", PyTorch Tutorials. - "GLU Variants Improve Transformer", arXiv. - "Liger Kernels", Hugging Face Hub (kernels-community/liger-kernels).