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

Viele ML-Praktiker haben eine gewisse Intuition dafür, 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 fängt bei null an 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, und TorchNeuron, sein Open-Source-natives PyTorch-Backend. TorchNeuron unterstützt standardmäßige PyTorch-Workflows auf Trainium, einschließlich eager execution, torch.compileund standardmäßigen 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 identifizieren, welche Operationen Neuron-Gerätezeit verbrauchen, und deren Self CPU Time zu vergleichen.
• Wie man eine Operation durch Neuron Explorer verfolgt, vom Framework-Dispatch über die Laufzeitplanung 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 einen Workload profilieren, hilft es zu verstehen, wie Neuron sich in PyTorch integriert.

Das Neuron-Backend verwendet den PrivateUse1 -Erweiterungspunkt von PyTorch. 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 dem 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.

Profiler einrichten

Sie verwenden zwei Objekte, um einen Neuron-Workload zu profilieren:

• 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, neben 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äteanweisungen 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 Profile Modes Dokumentation.

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 über einen SSH-Tunnel oder die Visual Studio Code Extension darauf 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 dem Profiling enthält das Ausgabeverzeichnis den System-Trace und, wenn Sie Gerätedaten erfassen, die zugeordneten .neff und .ntff Dateien. Das Neuron Executable File Format (NEFF) enthält die kompilierte Neuron-Ausführungsdatei, während das zugehörige NTFF den Hardware-Befehlstrace 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. Bei einem reinen Geräteprofil laden Sie die zugeordneten .neff und .ntff Dateien einzeln hoch. Siehe Step 5: Upload your profile für die UI- und Befehlszeilenanweisungen.

Profilierung von RMSNorm

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

Sie profilieren drei Implementierungen derselben Operation:

1. Naives RMSNorm mit einzelnen PyTorch-Operationen: Multiplikation, Mittelwert, Addition, reziproke Quadratwurzel und Multiplikation.
2. torch.rms_norm(), die in PyTorch eingebaute Variante, die zu einem einzigen fusionierten Neuron-Kernel heruntergebrochen wird.
3. Naives RMSNorm mit torch.compile, mit denselben Operationen wie die erste Implementierung.

Die Skripte verwenden dieselben Eingaben, dieselbe Profiler-Konfiguration und denselben Zeitplan, damit 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 damit, RMSNorm mit einzelnen PyTorch-Operationen zu schreiben. 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 ein, 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. Damit wird eine Warmup-Iteration erfasst, ohne den anfänglichen Setup-Overhead 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 zu platzieren. 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 verbraucht 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 fügt den Epsilon-Wert hinzu und löst eine Host-zu-Device-Kopie aus, weil der Skalar auf der CPU beginnt. aten::rsqrt berechnet den reziproken Quadratwurzelwert.

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

Trace im Explorer lesen

Im 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, Scheduling 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 Instanz/PID/Subsystem/Bezeichner:

Im Eager-Modus läuft jede Operation unabhängig. Verwenden Sie den Dependency-Chain-Viewer, um aten::mul.Tensor von der Framework-Schicht über das Runtime-Scheduling 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 PyTorchs eingebaute rms_norm Funktion. Neuron wandelt die Operation in einen einzigen fused Kernel um:

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

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

Self NEURON gesamt: 403.3 μs

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

Absenkungspfad vergleichen

In der naiven Implementierung durchschreitet jede Operation unabhängig die Dispatch-Grenze und folgt ihrem eigenen Pfad von Framework zu Runtime zu Hardware. 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")„, TorchDynamo erfasst den Python-Bytecode als FX-Graph, und AOTAutograd senkt ihn zu ATen-Operationen ab. Das Neuron-Backend erhält den vollständigen Graphen, anstatt jede Operation einzeln zu dispatchen. Es fusioniert die Operationen, optimiert Speicherzugriffe 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 wichtigsten Durchschnittswerte.

Selbst NEURON gesamt: 375.5 μs

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

Die Tabelle zeigt eine Geräteausführung, wie in Profil 2, aber der Code nimmt einen anderen Pfad. Statt eine eigens erstellte verschmolzene 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 verglichen mit 403 μs.

Zusammenfassung im Überblick

Sowohl torch.rms_norm als auch torch.compile führen RMSNorm in einem einzigen Kernel-Dispatch zusammen. In diesen Läufen benötigt jeder der beiden Ansätze etwa ein Sechstel der Gerätezeit der naiven Implementierung. Fusion entfernt die Lücken zwischen Hardware-Ausführungen und vermeidet die Materialisierung von Zwischenwerten wie x_sq und mean_sq zwischen separat aufgerufenen Operationen.

Fazit

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

Denselben Workflow können Sie bei größeren Modellen anwenden: 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 das device-level Profile ansehen. Sie öffnen das NEFF, untersuchen die Aufschlüsselung von TensorEngine, VectorEngine und direktem Speicherzugriff (DMA) und platzieren den Kernel auf der Roofline, um zu bestimmen, ob er compute-bound oder memory-bound ist.

Besonderen Dank an Aritra Roy Gosthipaty für die Inspiration, diese Serie zu beginnen, und an meine Kollegen Pinak Panigrahi, Samir Araujo, Akhil Azhikodan, Esha Lakhotia und alle übrigen für das Review!