cj401-amd/investigating-jax-traces

Investigación de trazas JAX

Resumen

Una investigación en siete pasos de una captura: el veredicto, la reconciliación, los carriles, kernel → HLO → línea de Python, techo teórico, qué cambiar y la lista de confianza.

Todo proviene de la captura proporcionada por el usuario: <host>.xplane.pb y/o <host>.trace.json.gz. No vuelvas a ejecutar la carga de trabajo, no vuelques XLA ni ejecutes otro perfilador. Cuando la traza no pueda responder, ponlo en la lista de "no es de fiar" en su lugar.

Objetivo

Después del informe, el lector sabe:

• (a) qué cambiar para acelerar el paso, como una estimación con sus suposiciones;
• (b) dónde esta traza no puede respaldar la atribución, con recuentos.

Herramientas

Las rutas de los scripts son relativas al directorio base de este skill, el que se muestra cuando el skill se carga. Ejecuta desde un directorio temporal: ambos scripts escriben sus salidas en el directorio actual.

B=<directorio base de este skill>
python3 -P $B/../analisis-sobrecarga-lanzamiento-jax/scripts/skip_stats.py <captura>/<host>.trace.json.gz
python3 -P $B/scripts/atribuir_kernels.py <directorio de captura> [--program ID] [--user-re REGEX]

• skip_stats.py (solo biblioteca estándar): el veredicto, el presupuesto de tiempo por flujo, el evento del host durante el tiempo inactivo y las métricas del artículo SKIP. Cómo leerlo: el skill de análisis de sobrecarga de lanzamiento JAX.
• atribuir_kernels.py (necesita los paquetes jax y xprof): cada kernel de GPU → instrucción HLO → op_name → línea de Python. Lee solo el .xplane.pb, copiándolo primero a ./atribucion/xplane/, porque XProf escribe los protos HLO junto a él. Escribe kernels.csv, ops.csv, lines.csv, how.csv, summary.txt y el texto HLO del programa en ./atribucion/.
  • exacto: el propio hlo_op del trazador, para kernels lanzados uno por llamada.
  • nombre: un kernel de gráfico cuyo nombre es el único candidato en su ventana.
  • orden: kernels compartidos, GEMMs de bibliotecas, copias y colectivas, emparejados en orden de planificación.
  • ninguno: no queda ningún candidato.
  • Las ventanas son las instrucciones entre kernels exactos vecinos, en la planificación del programa con bucles while expandidos según su recuento conocido de iteraciones.
  • Validación: en una captura de MiniGPT eligió la misma instrucción para todos los 147,840 kernels, como una tubería que había concordado al 100% con una ejecución del mismo programa sin búferes de comandos.

Cinco técnicas, cada una con su herramienta

Campo	Técnica	Herramienta aquí
Ingeniería de rendimiento	Ruta crítica; techo teórico	El presupuesto de flujo de skip_stats.py; un techo solo con los picos correctos (ver trampas)
Trazado distribuido	Unir eventos del host y GPU por ID de correlación; encontrar rezagados	skip_stats.py → rows.csv; la dispersión de los tiempos de fin de paso entre GPUs
Compiladores	Procedencia de la fuente a través de pases del compilador	atribuir_kernels.py → lines.csv
Contabilidad forense	Los totales deben reconciliarse entre registros independientes	Las comprobaciones en el paso 2
Diseño de información	Carriles alineados, uno por flujo de GPU	Un paso, cada GPU × flujo, en un solo eje temporal

Pasos

El plan son estos siete pasos, todos ejecutados sobre los archivos proporcionados. Una pregunta que solo una nueva captura podría responder (qué hace internamente una llamada del host, el efecto de una corrección, un barrido de lotes) va en la lista de "necesita una nueva captura" del paso 7, y si capturar o no es decisión del usuario.

1. Veredicto y causa del lado del host: ejecuta skip_stats.py.
2. Reconciliar. Cada comprobación se cumple o pasa a la lista de confianza:
   • Eventos de GPU por dispositivo: json = xplane;
   • cada evento de GPU se une a una llamada de lanzamiento del host;
   • los kernels por paso son constantes;
   • todas las GPUs hacen el mismo trabajo (recuentos de kernel iguales por paso, o la diferencia explicada);
   • para cada flujo, las participaciones de cómputo + colectiva + memcpy + inactivo suman el intervalo.
3. Carriles. Para el paso mediano, coloca cada GPU × flujo en un solo eje temporal, más los eventos del host del hilo de lanzamiento. Esto muestra la ruta crítica y los rezagados: qué GPU termina último y qué espera esa GPU.
4. Kernel → HLO → línea de Python: ejecuta atribuir_kernels.py.
   • Informa la participación de tiempo de GPU de cada clase de coincidencia, y los kernels que quedan sin candidato.
   • Usa lines.csv y ops.csv para las tablas por línea y por operación.
   • --user-re selecciona qué marcos cuentan como código del usuario. Por defecto es el primer marco fuera de site-packages y jax.
5. Techo teórico (roofline), solo para operaciones cuyos FLOPs y bytes se pueden obtener de sus formas HLO: el <programa>(<id>).hlo.txt que atribuir_kernels.py escribe en ./atribucion/. Usa el pico para el tipo de datos de la operación, citado de la especificación del proveedor y etiquetado como leído; la traza no puede medir picos.
6. Qué cambiar. Clasifica las causas por su participación en la ruta crítica. Para cada una, estima el tiempo del paso si desapareciera, p. ej., paso − inactividad del host en el flujo crítico. Etiquétalo como inferido, no ejecutado, y nombra el parámetro como una hipótesis para que el usuario la pruebe (donate_argnums, búferes de comandos desactivados, colectivas menos o más grandes).
7. Dónde no se puede confiar en esta traza: lista cada elemento con recuentos o participaciones. Luego añade la lista de "necesita una nueva captura": cada pregunta abierta, y qué captura la respondería.

Trampas (medidas en capturas ROCm MI355X)

• Las vistas de operaciones de XProf omiten los kernels de gráficos. hlo_stats cubrió el 37.6% del tiempo de kernel de GPU en una captura de 4 GPUs. El otro 62.4% (kernels de búfer de comandos) aparece en kernel_stats como operación "XlaModule". Verifica la cobertura antes de citar hlo_stats, op_profile o roofline_model.
• Los picos de techo de XProf para gfx950 son 2,517 TFLOP/s y 3,624 GiB/s, independientemente del tipo de datos. Las GEMMs FP32 necesitan el pico FP32.
• El tiempo de kernel RCCL/NCCL incluye la espera por pares. No es tiempo de transferencia.
• Colectivas asíncronas: su hlo_op (p. ej., all-gather.34) no está en el HLO almacenado en la captura. atribuir_kernels.py conserva la etiqueta y toma la línea de Python de la instrucción de inicio de colectiva coincidente, por orden.
• Las filas de memcpy de ROCm pueden mostrar kind:Unknown size:0.
• El trazador de Python infla el tiempo del host.
• Un programa puede aparecer dos veces (dos archivos .hlo_proto.pb con IDs diferentes). atribuir_kernels.py usa el ID de programa con más tiempo de GPU, del xplane; el json no tiene program_id.

Informe

• Markdown: hallazgos primero.
• Números: cada uno viene con su rango entre pasos y GPUs, etiquetado como medido, leído o inferido.
• Secciones: veredicto y causa, reconciliación, tabla por GPU/flujo, principales operaciones y líneas Python, qué cambiar (estimaciones) y la lista de confianza.
• Verificación de datos: antes de enviar, haz que un subagente nuevo verifique cada afirmación contra las tablas.
• Página HTML: si el usuario quiere una, el skill de escribir páginas de informe de traza la construye a partir de este informe y sus tablas.