caj401-amd/debugging-jax-oom

Depuración de OOM en JAX a partir de una captura

Resumen

Un paso jitted de JAX obtiene toda su memoria temporal como un bloque planificado. Por lo tanto, la memoria del paso es el plan de XLA, y el .xplane.pb contiene tanto ese plan como los eventos de tiempo de ejecución del asignador. scripts/oom_report.py lee ambos en 2–5 segundos, usando solo la biblioteca estándar:

• plan: asignación de búferes de cada módulo, desde /host:metadata. Muestra los parámetros (y cuánto se dona), las salidas no aliasadas, el bloque temporal y el pico del heap. Para los búferes activos en ese pico, proporciona tamaño, forma, rol y línea de Python. El rol proviene de la pila de nombres de JAX: jvp( (directo, conservado para la backward), transpose( (backward), las salidas de custom-call s8 son espacio de trabajo de la biblioteca.
• tiempo de ejecución: los eventos del asignador BFC. Muestran la capacidad, los bytes en uso al inicio de cada ejecución, el pico, las esperas AllocatorRetry, las ejecuciones que nunca lanzaron un kernel y los eventos del host entre las dos últimas ejecuciones. También muestran lo que estaba activo cuando comenzó la última ejecución, agrupado por el evento del host que lo asignó: salidas de jit_step ejecución 3, BufferFromHostBuffer(s32[…]) (un device_put) o "asignado antes de que comenzara la captura".

El informe termina con líneas de diagnóstico, cada una etiquetada como medido o inferido. Ejecútelo primero:

python3 -P <directorio base de esta skill>/scripts/oom_report.py <directorio de captura o .xplane.pb> [--module 'jit_train_step(123)'] [--top 25] [--json out.json]
python3 -P <directorio base de esta skill>/scripts/oom_report.py --self-test

Por defecto, analiza el módulo cuya última ejecución nunca lanzó un kernel. Si no hay ninguno, toma el módulo más grande que se ejecutó. Use --module para elegir otro, por ejemplo un paso de evaluación.

De la línea de diagnóstico a la causa y la corrección

Línea de diagnóstico	Causa	Corrección a proponer; el ahorro proviene de los bytes del informe
el paso en sí no cabe	el plan por sí solo excede el pool	Siga la división de roles en el pico (abajo)
↳ principalmente forward, conservado para el backward	activaciones: crecen con batch × seq; la atención ingenua añade B·H·S·S por capa	jax.checkpoint / remat en el bloque (ahorra la mayoría de los bytes forward-conservados, cuesta aproximadamente un forward adicional); micro-batch más pequeño con acumulación de gradientes; secuencia más corta (los búferes S² se reducen 4× cuando S se reduce a la mitad). Ofrezca atención fusionada solo como "verificar que los búferes B·H·S·S salgan del plan"
↳ principalmente backward	gradientes y sus temporales	Fragmentar (FSDP) o acumular; verificar copias f32 de gradientes bf16
↳ espacio de trabajo	espacio de trabajo de GEMM / conv de la biblioteca	Generalmente pequeño; en ROCm son 64 MiB por GEMM hipBLASLt en gfx950
crece X MiB por ejecución, = 1.00x las salidas no aliasadas	el programa mantiene vivas las salidas de cada paso (una fuga entre pasos): la lista de activos muestra salidas de … ejecución 1, 2, …	Encuentre el titular en los eventos del host entre ejecuciones (ej. $<unknown> append = list.append). Luego suelte la referencia, reduzca dentro del paso (una norma, no el árbol) o haga jax.device_get. Un pool más grande solo retrasa la OOM
ningún parámetro se dona … puede ahorrar hasta X	los parámetros/estado del optimizador viejos y nuevos están vivos juntos	donate_argnums en los argumentos de estado
otros arreglos activos ocupan el espacio; el más grande … provino de X	arreglos fuera del paso: conjunto de datos o copia de evaluación en la GPU, precarga, otras salidas de jit	Manténgalo en el host y transmítalo, libérelo o reduzca su dtype. "Asignado antes de que comenzara la captura" significa perfil más temprano para nombrarlo
AllocatorRetry … más larga N s, no lanzó ninguna	la OOM ocurrió dentro de la captura	Solo evidencia. La solicitud fallida suele ser el bloque temporal de ese módulo, o una salida
XLA ya rematerializó N instrucciones	el remat propio de XLA ya se ejecutó contra el límite	Espere que la corrección provenga del programa, no de XLA

Proporcione a cada corrección su ahorro estimado y su costo, etiquetados como inferidos. Verifique la estimación frente a la capacidad. "Aumente XLA_PYTHON_CLIENT_MEM_FRACTION" corrige un plan que no cabe solo si la GPU tiene espacio. Una capacidad muy por debajo de la HBM sugiere que la fracción está establecida; la captura no registra variables de entorno.

Anatomía de la memoria del manual en términos de JAX

El Manual de Ultra-Escala ("Perfilando el uso de memoria") divide la memoria de entrenamiento en pesos, gradientes, estado del optimizador y activaciones. En el informe:

• Pesos + estado del optimizador = los parámetros del paso. Adam en fp32 son 12 B por parámetro (p, m, v). Donado significa que no hay segunda copia.
• Gradientes = el rol backward en el pico, más las salidas si el paso las devuelve.
• Activaciones = forward, conservado para el backward. Los modelos con Scan / nn.scan las mantienen como búferes apilados [capas, …] en un bucle while.
• "El paso 1 cabe, el paso N falla": en PyTorch esto es el estado del optimizador creado en el primer paso. En JAX el estado existe antes del paso 1, así que busque la línea de crecimiento. La primera llamada también compila y autoajusta, así que un pico antes de la primera ejecución es tiempo de compilación.

Reglas

• Use solo la captura. No reconstruya el modelo ni vuelva a ejecutar la carga de trabajo para probar una corrección: estímela a partir de los bytes del plan y diga cómo confirmarla (jit(...).lower(...).compile().memory_analysis() del lado del usuario).
• No escriba dentro del directorio de captura. Los convertidores de XProf (xspace_to_tool_names, memory_viewer, xprof CLI) escriben archivos de caché *.hlo_proto.pb junto al .xplane.pb. Si necesita vistas de XProf, copie el .xplane.pb a un directorio temporal primero.
• Etiquete cada número como medido, leído o inferido.

Referencia medida (MI355X único, pools pequeños, 2026-11-06)

Captura	Lo que muestra el informe
plan de atención demasiado grande	necesita 3.72 GiB > 2.79 GiB; pico del heap 81.6% forward-conservado; f32[8,8,1024,1024] ×5 en la línea de softmax; bloque temporal 3,495 MiB = el error "3.41GiB"; AllocatorRetry 10 s; 89 instrucciones de remat
gradientes se mantienen cada paso	en uso al inicio del paso 316 → 1,252 MiB, +104.0 MiB por paso = 1.00× las salidas no aliasadas; plan 788 MiB cabe; $<unknown> append entre pasos
sin donación	necesita 3.88 GiB > 2.79 GiB; 1,902 MiB parámetros + 1,902 MiB salidas, 0 donados; temp solo 172 MiB
MaxText llama2-7b, 8 GPUs, FSDP (cabe)	temp 61.9 GiB; pico del heap 57% forward-conservado, las pilas de scan bf16[32,4,4096,11008] ×2 vía while.50

Limitaciones

• Un host por archivo: las ejecuciones multi-host tienen un .xplane.pb por host, así que ejecute cada uno.
• Solo HBM: solo se desglosa el espacio de memoria 0. La descarga al host se muestra como memory space 5.
• Contigüidad: "cabe" compara bytes. BFC también necesita un fragmento contiguo, por lo que un ajuste cercano aún puede fallar por fragmentación.
• Rol "other": búferes definidos por operaciones creadas por XLA sin pila de nombres de JAX, y sin alias que la tenga.
• CUDA: no probado. Usa los mismos eventos BFC y XLA, pero los nombres no están verificados.
• TPU / eager: las capturas de TPU solo dan totales del plan, sin pico del heap ni ejecuciones. Los programas eager (op por op) se muestran como muchos módulos pequeños, así que lea la sección de tiempo de ejecución allí.