cj401-amd/investigating-jax-traces

Investigation des traces JAX

Aperçu

Une investigation en sept étapes d'une seule capture : le verdict, la réconciliation, les couloirs de nage (swimlanes), noyau → HLO → ligne Python, roofline, quoi modifier, et la liste de fiabilité.

Tout provient de la capture donnée par l'utilisateur : <host>.xplane.pb et/ou <host>.trace.json.gz. Ne réexécutez pas la charge de travail, ne videz pas XLA et n'exécutez pas un autre profileur. Là où la trace ne peut pas répondre, mettez-le dans la liste "ne peut pas être fiable" à la place.

Objectif

Après le rapport, le lecteur sait :

• (a) quoi modifier pour accélérer l'étape, sous forme d'estimation avec ses hypothèses ;
• (b) où cette trace ne peut pas soutenir l'attribution, avec des comptes.

Outils

Les chemins des scripts sont relatifs au répertoire de base de cette compétence, celui affiché lorsque la compétence se charge. Exécutez depuis un répertoire de travail : les deux scripts écrivent leurs sorties dans le répertoire courant.

B=<répertoire de base de cette compétence>
python3 -P $B/../analyzing-jax-launch-overhead/scripts/skip_stats.py <capture>/<host>.trace.json.gz
python3 -P $B/scripts/attribute_kernels.py <répertoire capture> [--program ID] [--user-re REGEX]

• skip_stats.py (bibliothèque standard uniquement) : le verdict, le budget de temps par flux, l'événement hôte pendant le temps d'inactivité, et les métriques du papier SKIP. Comment le lire : la compétence analyzing-jax-launch-overhead.
• attribute_kernels.py (nécessite les paquets jax et xprof) : chaque noyau GPU → instruction HLO → nom_op → ligne Python. Il lit uniquement le .xplane.pb, en le copiant d'abord dans ./attribution/xplane/, car XProf écrit les protos HLO à côté. Il écrit kernels.csv, ops.csv, lines.csv, how.csv, summary.txt et le texte HLO du programme dans ./attribution/.
  • exact : le propre hlo_op du traceur, pour les noyaux lancés un par appel.
  • nom : un noyau de graphe dont le nom est le seul candidat dans sa fenêtre.
  • ordre : noyaux partagés, GEMMs de bibliothèque, copies et collectifs, correspondant dans l'ordre de l'ordonnancement.
  • aucun : aucun candidat restant.
  • Les fenêtres sont les instructions entre les noyaux exacts voisins, dans l'ordonnancement du programme avec les boucles while déroulées par leur nombre de passages connu.
  • Validation : sur une capture MiniGPT, il a choisi la même instruction pour tous les 147 840 noyaux qu'un pipeline qui avait 100% d'accord avec une exécution sans command buffers du même programme.

Cinq techniques, chacune avec son outil

Domaine	Technique	Outil ici
Ingénierie de performance	Chemin critique ; roofline	Le budget de flux de skip_stats.py ; un roofline uniquement avec les bons pics (voir pièges)
Traçage distribué	Joindre les événements hôte et GPU par ID de corrélation ; trouver les retardataires	skip_stats.py → rows.csv ; l'étalement des temps de fin d'étape entre les GPU
Compilateurs	Provenance de la source à travers les passes du compilateur	attribute_kernels.py → lines.csv
Comptabilité forensique	Les totaux doivent se réconcilier à travers des enregistrements indépendants	Les vérifications de l'étape 2
Conception de l'information	Couloirs de nage alignés, un par flux GPU	Une étape, chaque GPU × flux, sur un axe de temps

Étapes

Le plan comprend ces sept étapes, toutes exécutées sur les fichiers donnés. Une question que seule une nouvelle capture pourrait répondre (ce qu'un appel hôte fait à l'intérieur, l'effet d'une correction, un balayage de lot) va dans la liste "nécessite une nouvelle capture" de l'étape 7, et décider de capturer est à l'utilisateur.

1. Verdict et cause côté hôte : exécutez skip_stats.py.
2. Réconcilier. Chaque vérification soit tient, soit va dans la liste de fiabilité :
   • Événements GPU par périphérique : json = xplane ;
   • chaque événement GPU rejoint un appel de lancement hôte ;
   • les noyaux par étape sont constants ;
   • tous les GPU font le même travail (comptes de noyaux égaux par étape, ou la différence expliquée) ;
   • pour chaque flux, les parts calcul + collectif + memcpy + inactivité s'additionnent pour couvrir la durée.
3. Couloirs de nage (swimlanes). Pour l'étape médiane, placez chaque GPU × flux sur un axe de temps, plus les événements hôte du thread de lancement. Cela montre le chemin critique et les retardataires : quel GPU termine en dernier, et ce que ce GPU attend.
4. Noyau → HLO → Ligne Python : exécutez attribute_kernels.py.
   • Rapportez la part de temps GPU de chaque classe de correspondance, et les noyaux restés sans candidat.
   • Utilisez lines.csv et ops.csv pour les tableaux par ligne et par opération.
   • --user-re choisit quelles trames comptent comme le code de l'utilisateur. Par défaut, c'est la première trame en dehors de site-packages et jax.
5. Roofline, uniquement pour les opérations dont vous pouvez obtenir les FLOPs et octets à partir de leurs formes HLO : le <programme>(<id>).hlo.txt que attribute_kernels.py écrit dans ./attribution/. Utilisez le pic pour le type de données de l'opération, cité depuis la spécification du fournisseur et étiqueté lu ; la trace ne peut pas mesurer les pics.
6. Quoi modifier. Classez les causes par leur part du chemin critique. Pour chacune, estimez le temps d'étape si elle disparaissait, par ex. étape − inactivité hôte sur le flux critique. Étiquetez-le inféré, non exécuté, et nommez le bouton comme une hypothèse pour l'utilisateur à tester (donate_argnums, command buffers désactivés, collectifs plus petits ou moins nombreux).
7. Où cette trace ne peut pas être fiable : listez chaque élément avec des comptes ou des parts. Ajoutez ensuite la liste "nécessite une nouvelle capture" : chaque question ouverte, et quelle capture y répondrait.

Pièges (mesurés sur des captures ROCm MI355X)

• Les vues d'opérations de XProf ignorent les noyaux de graphe. hlo_stats couvrait 37,6% du temps de noyau GPU sur une capture 4-GPU. Les autres 62,4% (noyaux command-buffer) apparaissent dans kernel_stats comme opération "XlaModule". Vérifiez la couverture avant de citer hlo_stats, op_profile ou roofline_model.
• Les pics roofline de XProf pour gfx950 sont 2 517 TFLOP/s et 3 624 Gio/s, quel que soit le type de données. Les GEMMs FP32 nécessitent le pic FP32.
• Le temps de noyau RCCL/NCCL inclut l'attente des pairs. Ce n'est pas du temps de transfert.
• Collectifs asynchrones : leur hlo_op (par ex. all-gather.34) n'est pas dans le HLO stocké dans la capture. attribute_kernels.py conserve l'étiquette et prend la ligne Python de l'instruction collective-start correspondante, par ordre.
• Les lignes memcpy ROCm peuvent afficher kind:Unknown size:0.
• Le traceur Python gonfle le temps hôte.
• Un programme peut apparaître deux fois (deux fichiers .hlo_proto.pb avec des ID différents). attribute_kernels.py utilise l'ID du programme avec le plus de temps GPU, à partir du xplane ; le json n'a pas de program_id.

Rapport

• Markdown : résultats en premier.
• Nombres : chacun vient avec sa plage entre les étapes et les GPU, étiqueté mesuré, lu ou inféré.
• Sections : verdict et cause, réconciliation, tableau par GPU/flux, opérations principales et lignes Python, quoi modifier (estimations), et la liste de fiabilité.
• Vérification des faits : avant d'envoyer, faites vérifier chaque affirmation par un sous-agent frais en consultant les tableaux.
• Page HTML : si l'utilisateur le souhaite, la compétence writing-trace-report-pages la construit à partir de ce rapport et de ses tableaux.