cj401-amd/investigating-jax-traces

Untersuchung von JAX-Traces

Übersicht

Eine siebenstufige Untersuchung einer Aufnahme: das Urteil, der Abgleich, Swimlanes, Kernel → HLO → Python-Zeile, Roofline, Änderungsvorschläge und die Liste der nicht vertrauenswürdigen Punkte.

Alles stammt aus der vom Benutzer bereitgestellten Aufnahme: <host>.xplane.pb und/oder <host>.trace.json.gz. Führen Sie den Workload nicht erneut aus, dumpen Sie XLA oder starten Sie einen anderen Profiler. Wo der Trace keine Antwort liefern kann, wird dies in die Liste "kann nicht vertraut werden" aufgenommen.

Ziel

Nach dem Bericht weiß der Leser:

• (a) was geändert werden muss, um den Schritt schneller zu machen, als Schätzung mit seinen Annahmen;
• (b) wo dieser Trace die Zuordnung nicht stützen kann, mit Zählern.

Werkzeuge

Skriptpfade sind relativ zum Basisverzeichnis dieses Skills, dem Verzeichnis, das beim Laden des Skills angezeigt wird. Führen Sie die Skripte aus einem temporären Verzeichnis aus: Beide Skripte schreiben ihre Ausgaben in das aktuelle Verzeichnis.

B=<Basisverzeichnis dieses Skills>
python3 -P $B/../analyzing-jax-launch-overhead/scripts/skip_stats.py <capture>/<host>.trace.json.gz
python3 -P $B/scripts/attribute_kernels.py <capture dir> [--program ID] [--user-re REGEX]

• skip_stats.py (nur Standardbibliothek): das Urteil, das Zeitbudget pro Stream, das Host-Ereignis während der Leerlaufzeit und die SKIP-Paper-Metriken. Wie man es liest: das Skill "analyzing-jax-launch-overhead".
• attribute_kernels.py (benötigt die Pakete jax und xprof): jeder GPU-Kernel → HLO-Anweisung → op_name → Python-Zeile. Es liest nur die .xplane.pb, kopiert sie zuerst nach ./attribution/xplane/, weil XProf die HLO-Protos daneben schreibt. Es schreibt kernels.csv, ops.csv, lines.csv, how.csv, summary.txt und den HLO-Text des Programms nach ./attribution/.
  • exact: das eigene hlo_op des Tracers, für Kernel, die pro Aufruf gestartet werden.
  • name: ein Graph-Kernel, dessen Name der einzige Kandidat in seinem Fenster ist.
  • order: gemeinsame Kernel, Bibliotheks-GEMMs, Kopien und Kollektive, abgeglichen in der Plan-Reihenfolge.
  • none: kein Kandidat mehr übrig.
  • Fenster sind die Anweisungen zwischen benachbarten exakten Kernel, im Plan des Programms mit While-Schleifen, die um ihre bekannte Anzahl von Durchläufen erweitert wurden.
  • Validierung: Bei einer MiniGPT-Aufnahme wählte es dieselbe Anweisung für alle 147.840 Kernel aus, als eine Pipeline, die zu 100 % mit einem Lauf desselben Programms ohne Befehlsbuffer übereinstimmte.

Fünf Techniken, jede mit ihrem Tool

Feld	Technik	Tool hier
Performance-Engineering	Kritischer Pfad; Roofline	Das Streambudget von skip_stats.py; eine Roofline nur mit den richtigen Spitzenwerten (siehe Fallstricke)
Verteiltes Tracing	Host- und GPU-Ereignisse nach Korrelations-ID zusammenführen; Nachzügler finden	skip_stats.py → rows.csv; die Streuung der Schritt-Endzeiten über GPUs hinweg
Compiler	Quellherkunft durch Compiler-Durchläufe	attribute_kernels.py → lines.csv
Forensische Buchhaltung	Summen müssen über unabhängige Aufzeichnungen hinweg übereinstimmen	Die Überprüfungen in Schritt 2
Informationsdesign	Ausgerichtete Swimlanes, eine pro GPU-Stream	Ein Schritt, jede GPU × Stream, auf einer Zeitachse

Schritte

Der Plan besteht aus diesen sieben Schritten, die alle auf den angegebenen Dateien ausgeführt werden. Eine Frage, die nur eine neue Aufnahme beantworten könnte (was ein Host-Aufruf im Inneren tut, die Wirkung einer Korrektur, eine Batch-Durchlaufanalyse), wird in die Liste "benötigt eine neue Aufnahme" von Schritt 7 aufgenommen, und ob eine Aufnahme durchgeführt wird, liegt im Ermessen des Benutzers.

1. Urteil und Host-seitige Ursache: Führen Sie skip_stats.py aus.
2. Abgleich. Jede Überprüfung hält entweder oder geht in die Vertrauensliste:
   • GPU-Ereignisse pro Gerät: json = xplane;
   • jedes GPU-Ereignis verbindet einen Host-Startaufruf;
   • Kernel pro Schritt sind konstant;
   • alle GPUs erledigen die gleiche Arbeit (gleiche Kernel-Anzahlen pro Schritt, oder der Unterschied wird erklärt);
   • für jeden Stream addieren sich die Anteile von Berechnung + Kollektiv + memcpy + Leerlauf zur Spanne.
3. Swimlanes. Legen Sie für den medianen Schritt jede GPU × Stream auf eine Zeitachse, plus die Host-Ereignisse des Start-Threads. Dies zeigt den kritischen Pfad und die Nachzügler: welche GPU zuletzt fertig wird und worauf diese GPU wartet.
4. Kernel → HLO → Python-Zeile: Führen Sie attribute_kernels.py aus.
   • Berichten Sie den Anteil jeder Übereinstimmungsklasse an der GPU-Zeit und die Kernel, die ohne Kandidaten übrig bleiben.
   • Verwenden Sie lines.csv und ops.csv für die Tabellen pro Zeile und pro Operation.
   • --user-re wählt aus, welche Frames als Code des Benutzers gelten. Standardmäßig ist es der erste Frame außerhalb von site-packages und jax.
5. Roofline, nur für Operationen, deren FLOPs und Bytes Sie aus ihren HLO-Formen ermitteln können: die <program>(<id>).hlo.txt, die attribute_kernels.py nach ./attribution/ schreibt. Verwenden Sie den Spitzenwert für den Datentyp der Operation, zitiert aus der Herstellerspezifikation und als gelesen gekennzeichnet; der Trace kann keine Spitzenwerte messen.
6. Was zu ändern ist. Ordnen Sie die Ursachen nach ihrem Anteil am kritischen Pfad. Für jede schätzen Sie die Schrittzeit, wenn sie wegfallen würde, z. B. Schritt − Host-Leerlauf auf dem kritischen Stream. Kennzeichnen Sie es als geschlossen, nicht ausgeführt, und nennen Sie den Stellhebel als Hypothese, die der Benutzer testen kann (donate_argnums, Befehlsbuffer aus, weniger oder größere Kollektive).
7. Wo dieser Trace nicht vertrauenswürdig ist: Listen Sie jeden Punkt mit Anzahlen oder Anteilen auf. Fügen Sie dann die Liste "benötigt eine neue Aufnahme" hinzu: jede offene Frage und welche Aufnahme sie beantworten würde.

Fallstricke (gemessen an ROCm MI355X-Aufnahmen)

• XProfs Op-Ansichten lassen Graph-Kernel aus. hlo_stats deckte 37,6 % der GPU-Kernelzeit auf einer 4-GPU-Aufnahme ab. Die anderen 62,4 % (Command-Buffer-Kernel) erscheinen in kernel_stats als Operation "XlaModule". Überprüfen Sie die Abdeckung, bevor Sie hlo_stats, op_profile oder roofline_model zitieren.
• XProfs Roofline-Spitzenwerte für gfx950 betragen 2.517 TFLOP/s und 3.624 GiB/s, unabhängig vom Datentyp. FP32-GEMMs benötigen den FP32-Spitzenwert.
• RCCL/NCCL-Kernelzeit beinhaltet Warten auf Peers. Es ist keine Übertragungszeit.
• Asynchrone Kollektive: Ihre hlo_op (z. B. all-gather.34) befindet sich nicht im HLO, das in der Aufnahme gespeichert ist. attribute_kernels.py behält das Label und übernimmt die Python-Zeile aus der passenden Kollektiv-Start-Anweisung, in der Reihenfolge.
• ROCm-Memcpy-Zeilen können kind:Unknown size:0 anzeigen.
• Der Python-Tracer bläht die Hostzeit auf.
• Ein Programm kann zweimal erscheinen (zwei .hlo_proto.pb-Dateien mit verschiedenen IDs). attribute_kernels.py verwendet die Programm-ID mit der meisten GPU-Zeit, aus dem xplane; die json hat keine program_id.

Bericht

• Markdown: Ergebnisse zuerst.
• Zahlen: Jede kommt mit ihrer Spanne über Schritte und GPUs hinweg, gekennzeichnet als gemessen, gelesen oder geschlossen.
• Abschnitte: Urteil und Ursache, Abgleich, Tabelle pro GPU/Stream, Top-Operationen und Python-Zeilen, Änderungsvorschläge (Schätzungen) und die Vertrauensliste.
• Faktencheck: Lassen Sie vor dem Senden einen frischen Subagenten jede Behauptung anhand der Tabellen überprüfen.
• HTML-Seite: Wenn der Benutzer eine wünscht, erstellt das Skill "writing-trace-report-pages" eine Seite aus diesem Bericht und seinen Tabellen.