In einem Beitrag auf der Plattform Hugging Face eröffnet Amazon eine Reihe zur Analyse von Inferenzleistung, also der Ausführung eines bereits trainierten Modells. Dieser erste Teil erklärt, wie man mit Profiling-Tools für den AWS Trainium-Chip, der für das Training und die Inferenz von Deep-Learning-Modellen entwickelt wurde, herausfindet, wo die Zeit verbraucht wird.
- Amazon startet auf Hugging Face eine Reihe zur Analyse der Inferenzleistung auf AWS Trainium.
- Der Beitrag zeigt, wie Profiling-Tools mit TorchNeuron die Zeitfresser auf dem Chip finden.
- Bei RMSNorm braucht die fusionierte Version nur etwa ein Sechstel der Rechenzeit der naiven Implementierung.
- Der zweite Teil behandelt die Profilierung auf Geräteebene anhand der NEFF-Datei.
Die Anleitung richtet sich an Leser ohne vorherige Profiling-Erfahrung. Sie basiert auf TorchNeuron, dem nativen Open-Source-Backend von PyTorch für Trainium, das Standard-PyTorch-Workflows, einschließlich der Ausführung im Eager-Modus und mit torch.compile, mit wenigen Codeänderungen unterstützt. TorchNeuron integriert sich auch in die Neuron Kernel Interface (NKI), die die Erstellung benutzerdefinierter Operatoren ermöglicht.
Den Profiler konfigurieren
Das Profiling basiert auf zwei Objekten. NeuronConfig definiert die zu erfassenden Datenströme, die zu verfolgenden NeuronCores und den Ausgabeordner und wird dann an torch.profiler.profile übergeben. NeuronProfiler stellt eine Exportfunktion bereit, die die Framework-Trace im Format trace.json neben den System-Traces von Neuron schreibt. Der Parameter correlate_device_time=True ermöglicht es, die auf dem Gerät verbrachte Zeit in der Tabelle key_averages() anzuzeigen, wo sie mit der CPU-Zeit verglichen werden kann.
Die Traces können dann in Neuron Explorer geöffnet werden, einem Visualisierungstool, das im Browser läuft. Es ermöglicht, eine Operation von der Verteilung im Framework über die Planung in der Laufzeitumgebung bis zu ihrer Ausführung auf der Hardware zu verfolgen. Für einen NeuronCore, also einen Rechenkern des Chips, enthält die NEFF-Datei das kompilierte ausführbare Programm, während die entsprechende NTFF-Datei die Trace der Hardwarebefehle enthält.
Drei Versionen derselben Operation
Der Beitrag untersucht die Funktion RMSNorm, eine einfache Operation, die aus einer Reihe arithmetischer Berechnungen und Reduktionen besteht. Drei Implementierungen werden verglichen: eine Version, die Operation für Operation in PyTorch geschrieben ist, der Aufruf der integrierten Funktion torch.rms_norm() und die Operation-für-Operation geschriebene, aber mit torch.compile kompilierte Version.
In der naiven Version zeigt die Tabelle, dass drei Aufrufe von aten::mul den größten Teil der auf dem Gerät verbrachten Zeit ausmachen. Die gesamte Neuron-Zeit erreicht 2.355 µs für sechs Operatoraufrufe. Laut Autor belegt die integrierte Funktion torch.rms_norm(), die in einen einzigen fusionierten Rechenkern umgewandelt wird, nur eine Operation, aten::_fused_rms_norm, mit einer Neuron-Zeit von 403,3 µs. Die mit torch.compile kompilierte Version erscheint als eine einzige kompilierte Region mit einer Dauer von 375,5 µs.
Der Beitrag präzisiert, dass in diesen Tests beide Ansätze etwa ein Sechstel der Rechenzeit der naiven Version benötigen. Die Fusion eliminiert Unterbrechungen zwischen Hardwareausführungen und vermeidet die Materialisierung von Zwischenwerten wie x_sq und mean_sq zwischen separat gesendeten Operationen.
Der Autor weist darauf hin, dass die Methode auf größere Modelle angewendet werden kann: einen Hotspot in der Tabelle identifizieren, ihn in der Trace überprüfen und dann bestimmen, ob ein fusionierter Operator oder torch.compile die Kosten reduzieren kann. Der zweite Teil der Serie wird sich mit der Geräteebenen-Profilierung, einschließlich der Analyse der NEFF-Datei, befassen.
