Profiling TorchNeuron (Teil 1): Systemweites Profiling auf AWS Trainium.

Viele ML-Praktiker haben eine gewisse Intuition darüber, wie man Modelle schnell laufen lässt, aber weniger wissen, wie man analysiert, warum ein Modell langsam ist. In Teil 1 dieser Serie lernen Sie, wie Sie die Inferenzleistung systematisch analysieren und herausfinden, wo die Zeit verbracht wird.

Dieser Beitrag beginnt bei null und setzt keine Vorerfahrung mit Profiling voraus. Die Anleitung verwendet AWS Trainium, einen speziell entwickelten KI-Chip für Deep-Learning-Trainings- und Inferenz-Workloads, sowie TorchNeuron, sein Open-Source-natives PyTorch-Backend. TorchNeuron unterstützt Standard-PyTorch-Workflows auf Trainium, darunter eager execution, torch.compile, und Standard-PyTorch-APIs, mit minimalen Codeänderungen. Darüber hinaus integriert sich TorchNeuron mit Neuron Kernel Interface (NKI) um benutzerdefinierte Operatoren und Leistungsoptimierung für Trainium-Hardware zu ermöglichen. Am Ende sollten Sie Folgendes wissen:

• Wie man torch.profiler mit dem Neuron-Backend konfiguriert und einen sauberen Trace erfasst.
• Wie man die key_averages() -Tabelle liest, um zu erkennen, welche Operationen Neuron-Gerätezeit verbrauchen, und deren Self CPU Time vergleicht.
• Wie man eine Operation durch Neuron Explorer verfolgt, vom Framework-Dispatch über die Runtime-Scheduling bis zur Hardware-Ausführung auf einem NeuronCore.
• Wie sich diese Signale über drei Implementierungen derselben Operation hinweg verändern.

Neuron als PyTorch-Backend

Bevor Sie eine Workload profilen, hilft es zu verstehen, wie Neuron sich in PyTorch integriert.

Das Neuron-Backend verwendet PyTorchs PrivateUse1 Erweiterungspunkt. Dieser Mechanismus ermöglicht es einem Beschleuniger, sich in das Framework zu integrieren, ohne einen eigenen Dispatch-Schlüssel wie cudahinzuzufügen. Das torch_neuronx -Paket arbeitet mit Upstream-PyTorch zusammen, um Neuron als natives Backend für PyTorch-Code zu registrieren. Einzelheiten zu diesem Erweiterungsmechanismus finden Sie unter Facilitating New Backend Integration by PrivateUse1.

Einrichten des Profilers

Sie verwenden zwei Objekte, um eine Neuron-Workload zu profilen:

• NeuronConfig steuert, welche Neuron-Datenströme und NeuronCores erfasst werden und wohin die Ausgabe geschrieben wird. Sie übergeben es an torch.profiler.profile als experimental_config.
• NeuronProfiler nimmt ein NeuronConfig entgegen und stellt einen export_trace -Callback bereit. Der Callback schreibt den Framework-Trace, trace.json, zusammen mit den Neuron-System-Traces.

nc = NeuronConfig(
    modes=[ProfileMode.DEVICE, ProfileMode.RUNTIME],
    profile_output_dir=OUTPUT_DIR,
    capture_enabled_for_nc='0',
    correlate_device_time=True,
)
exporter = NeuronProfiler(nc)

Der modes Parameter wählt die zu erfassenden Datenströme aus, z. B. Gerätebefehle und Runtime-Traces. Der capture_enabled_for_nc Parameter wählt die zu verfolgenden NeuronCores aus. Das Setzen von correlate_device_time=True überträgt die Gerätezeit aus dem Trace in die key_averages() -Tabelle, wo Sie CPU- und Neuron-Gerätezeit vergleichen können. Die vollständige Parameterreferenz und die verfügbaren ProfileMode -Werte finden Sie in der NeuronConfig-Dokumentation Profile Modes.

Tooling: Neuron Explorer

Neuron Explorer ist ein browserbasierter Viewer für den System-Trace, die Geräte-Timeline und die Aufschlüsselung pro Engine. Sie können darauf über einen SSH-Tunnel oder die Visual Studio Code-Erweiterung zugreifen. In diesem Beitrag verwenden Sie Explorer, um zu bestätigen, was key_averages() numerisch meldet, und um zu untersuchen, was zwischen den Dispatches passiert.

Nach der Profilerstellung enthält das Ausgabeverzeichnis den System-Trace und, wenn Sie Gerätedaten erfassen, passende .neff - und .ntff -Dateien. Das Neuron Executable File Format (NEFF) enthält die kompilierte Neuron-Executable, während das passende NTFF den Hardware-Befehls-Trace enthält. Die Dateien teilen sich einen numerischen Hash, der sie als Paar identifiziert. Um das Profil in Neuron Explorer zu öffnen, laden Sie das vollständige Ausgabeverzeichnis für ein System- oder kombiniertes Profil hoch. Für ein Nur-Geräte-Profil laden Sie die passenden .neff - und .ntff -Dateien einzeln hoch. Siehe Step 5: Upload your profile für die UI- und Befehlszeilenanweisungen.

Profile RMSNorm

Für diese erste Analyse profilieren Sie RMSNorm, eine kleine Operation mit einer klaren Abfolge von arithmetischen Operationen und Reduktionsoperationen.

Sie profilieren drei Implementierungen derselben Operation:

1. Naive RMSNorm mit einzelnen PyTorch-Operationen: multiplizieren, Mittelwert bilden, addieren, reziproke Quadratwurzel und multiplizieren.
2. torch.rms_norm(), die in PyTorch eingebaute Variante, die zu einem einzigen fusionierten Neuron-Kernel herabgesetzt wird.
3. Naives RMSNorm mit torch.compile, unter Verwendung derselben Operationen wie die erste Implementierung.

Die Skripte verwenden dieselben Eingaben, dieselbe Profiler-Konfiguration und denselben Zeitplan, sodass die Ergebnisse direkt vergleichbar sind.

Alle Skripte zur Reproduktion dieser Arbeit sind im torch-neuronx examples-Verzeichnis verfügbar.

Profil 1: Naive PyTorch-Operationen

Beginnen Sie mit der Implementierung von RMSNorm mit einzelnen PyTorch-Operationen. Jede Zeile entspricht einer oder mehreren ATen-Operationen und kann Zwischentensoren erzeugen.

device = torch.device('neuron')

x = torch.randn(4, 512, dtype=torch.bfloat16, device=device)
weight = torch.ones(512, dtype=torch.bfloat16, device=device)

def rms_norm_fn(t):
    x_sq = t * t                                    # square
    mean_sq = x_sq.mean(dim=-1, keepdim=True)       # mean of squares
    rms = torch.rsqrt(mean_sq + 1e-6)               # 1 / sqrt(mean_sq + eps)
    return t * rms * weight                         # normalize + scale

Rufen Sie den Profiler mit zwei Argumenten auf, die für das Neuron-Backend erforderlich sind:

1. Fügen Sie ProfilerActivity.PrivateUse1 hinzu, um Aktivitäten vom Neuron-Gerät zu erfassen.
2. Übergeben Sie NeuronConfig als experimental_config um die Neuron-spezifischen Kineto-Funktionen zu aktivieren.

Der Zeitplan verwendet wait=2, warmup=1, und active=1 über vier prof.step() Aufrufe hinweg. Dies erfasst eine Warmup-Iteration, ohne den anfänglichen Einrichtungsaufwand einzubeziehen.

Rufen Sie torch.neuron.synchronize() innerhalb der Profiler-Schleife auf. Ohne diese Synchronisierung kann der Profiler exportieren, bevor das Gerät fertig ist, und die Neuron-Zeitspalte leer lassen.

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.PrivateUse1],
    experimental_config=nc,
    record_shapes=False,
    on_trace_ready=exporter.export_trace,
    schedule=torch.profiler.schedule(wait=2, warmup=1, active=1),
) as prof:
    for step in range(4):
        with torch.no_grad():
            _ = rms_norm_fn(x)
        torch.neuron.synchronize()
        prof.step()

Nachdem die Profilerstellung abgeschlossen ist, betrachten Sie die key_averages() Tabelle. Sortieren Sie nach self_privateuse1_time_total um die Operationen mit den höchsten Gerätekosten zuerst anzuzeigen. Die hier gezeigte Tabelle ist auf die RMSNorm-Operationen gefiltert, die Neuron-Gerätezeit melden. Die vollständige key_averages() Ausgabe enthält auch CPU-seitige Ereignisse wie Allokationen und Host-zu-Gerät-Übertragungen.

print(prof.key_averages().table(sort_by="self_privateuse1_time_total", row_limit=4))

Die Key-Averages-Tabelle zeigt, dass drei aten::mul-Aufrufe den größten Anteil der Neuron-Gerätezeit ausmachen.

Die Tabelle zeigt, welche ATen-Operationen ausgeführt wurden, wie oft sie ausgeführt wurden und wie viel Zeit sie auf der CPU und dem Neuron-Gerät verbracht haben. aten::mul hat die höchste gesamte Gerätezeit, da die Funktion drei Multiplikationen durchführt.

Self NEURON gesamt: 2.355 ms

Jede Zeile ist direkt der Funktion zugeordnet. aten::mul ausgeführt wird x * x, x * rms, und die abschließende Multiplikation mit weight. aten::mean reduziert die letzte Dimension. aten::add addiert den Epsilon-Wert und löst eine Host-zu-Device-Kopie aus, weil der Skalar auf der CPU startet. aten::rsqrt berechnet den Kehrwert der Quadratwurzel.

Neuron Explorer verfolgt eine aten::mul-Operation vom Framework-Dispatch über die Runtime-Planung bis zur Ausführung auf den Neuron-Hardware-Lanes.

Den Trace in Explorer lesen

In Explorer fließt die Zeit von links nach rechts und der Ereignis-Lebenszyklus von oben nach unten. Jede ATen-Operation folgt demselben Muster: Erzeugung in der Framework-Lane, Planung in der Runtime-Lane als nrta_execute_schedule, und Geräteausführung in der Hardware-Lane als nc_exec_running.

Die System-Trace-Ansicht ordnet Ereignisse in Lanes an. Im Standard-Gruppierungsmodus folgen die Beschriftungen dem Muster instance/PID/subsystem/identifier:

Im Eager-Modus läuft jede Operation unabhängig. Verwenden Sie den Dependency-Chain-Viewer, um aten::mul.Tensor von der Framework-Schicht über die Runtime-Planung bis zur Geräteausführung zu verfolgen.

Der Dependency-Chain-Viewer hebt den Lebenszyklus einer Operation über die Framework-, Runtime- und Hardware-Lanes hinweg hervor.

Profil 2: torch.rms_norm

Ersetzen Sie die einzelnen Operationen durch die in PyTorch integrierte Funktion rms_norm . Neuron übersetzt die Operation in einen einzigen fused Kernel:

def rms_norm_fn(t):
    return torch.rms_norm(t, [512], weight)

Der fused RMSNorm erscheint als eine Geräteoperation in der Key-Averages-Tabelle.

Self NEURON gesamt: 403.3 μs

Das fusionierte RMSNorm folgt einem einzigen Framework-zu-Hardware-Pfad, anstatt jede arithmetische Operation separat zu dispatchen.

Den Lowering-Pfad vergleichen

In der naiven Implementierung überschreitet jede Operation unabhängig die Dispatch-Grenze und folgt ihrem eigenen Framework-zu-Runtime-zu-Hardware-Pfad. In der fusionierten Implementierung deckt ein einziger Python-Aufruf die vollständige Berechnung ab. Der Trace enthält einen Dispatch und ein Hardware-Ereignis, ohne Lücken zwischen den Komponentenoperationen.

Profil 3: Naives PyTorch mit torch.compile

Für das letzte Beispiel kompilieren Sie die ursprüngliche naive RMSNorm-Funktion mit torch.compile.

Mit @torch.compile(backend="neuron"), fängt TorchDynamo den Python-Bytecode als FX-Graph ab und AOTAutograd leitet ihn auf ATen-Operationen herunter. Das Neuron-Backend erhält den vollständigen Graphen, anstatt jede Operation separat zu dispatchen. Es fusioniert die Operationen, optimiert den Speicherzugriff und kompiliert den Graphen in ein einzelnes NEFF-Artefakt. Das Artefakt läuft als ein einziger Hardware-Dispatch.

@torch.compile(backend='neuron', fullgraph=True)
def naive_compiled(t):
    x_sq = t * t
    mean_sq = x_sq.mean(dim=-1, keepdim=True)
    rms = torch.rsqrt(mean_sq + 1e-6)
    return t * rms * weight

Die kompilierte Funktion erscheint als eine Torch-Compiled Region in der Tabelle der Schlüsseldurchschnittswerte.

Self NEURON gesamt: 375.5 μs

Der kompilierte Graph folgt einem einzigen Pfad durch die Laufzeitplanung und die Hardware-Ausführung.

Die Tabelle zeigt, wie in Profil 2, eine einzige Geräteausführung, aber der Code nimmt einen anderen Weg. Statt eine eigens entwickelte fusionierte Operation aufzurufen, entdeckt der Compiler die Fusion aus der ursprünglichen Python-Funktion. In diesem Lauf ist die Neuron-Zeit vergleichbar mit torch.rms_norm, mit 375 μs gegenüber 403 μs.

Zusammenfassung im Vergleich

Beide torch.rms_norm und torch.compile fuse in einer Kernel-Ausführung zusammen. In diesen Läufen benötigt jeder der beiden Ansätze etwa ein Sechstel der Gerätezeit der naiven Implementierung. Fusion beseitigt die Lücken zwischen den Hardware-Ausführungen und vermeidet die Materialisierung von Zwischenwerten wie x_sq und mean_sq zwischen separat ausgeführten Operationen.

Fazit

key_averages() zeigt, welche Operationen Gerätezeit verbrauchen, während die vollständige Ausgabe auch CPU-seitige Einrichtung und Transfers offenlegt. Explorer bietet den entsprechenden visuellen Trace. Sie können eine Operation vom Framework-Dispatch über die Runtime-Planung bis zur Hardware-Ausführung verfolgen und die Lücken sehen, die die Fusion beseitigt.

Diesen Workflow können Sie auch bei größeren Modellen verwenden: Finden Sie einen Hotspot in der Tabelle, verifizieren Sie ihn im Trace und bestimmen Sie dann, ob eine verschmolzene Operation oder torch.compile ihre Kosten senken kann.

In Teil 2 werden wir uns dem Geräte-Level-Profil zuwenden. Sie öffnen das NEFF, betrachten die Aufschlüsselung von TensorEngine, VectorEngine und Direct Memory Access (DMA) und platzieren den Kernel auf der Roofline, um festzustellen, ob er compute-bound oder memory-bound ist.

Besonderer Dank gilt Aritra Roy Gosthipaty für die Inspiration zu dieser Serie sowie meinen Kollegen Pinak Panigrahi, Samir Araujo, Akhil Azhikodan, Esha Lakhotia und allen anderen für das Review!