Código de error: E3000

Categoría: Tiempo de compilación: Falla de asignación de SparseCore

Este error se produce cuando el compilador XLA:SparseCore no puede asignar un bloque contiguo de memoria en el espacio de memoria especificado que requiere el programa SparseCore actual.

Ejemplos de mensajes de error:

INTERNAL:Failed to run pass pipeline. Hlo-Op: result.1:279:1: error: 'memref.alloca' op current allocation offset upper bound (140704 words) exceeds the legitimate user allocatable offset upper bound (131071 words) in memory space 201 when allocating 23440 words. result.1:279:1: note: see current operation: %232 = "memref.alloca"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<23440xf32, 201>

Backends de XLA: TPU

Descripción general

SparseCore (SC) es un procesador especializado para cargas de trabajo dispersas. Se basa en jerarquías de memoria específicas para administrar el movimiento de datos de manera eficiente. El compilador de XLA intenta dimensionar y asignar búferes de forma estática en función de los límites de hardware y las formas definidas por el usuario. Este error indica una condición de falta de memoria (OOM) durante esta fase de asignación.

Por lo general, el mensaje de error especifica un ID de espacio de memoria. A continuación, se muestran los espacios de memoria comunes y sus codificaciones de números enteros:

ID del espacio de memoria Nombre Descripción
0 Smem Memoria escalar local. Se usa para registros escalares y flujo de control.
201 TileSpmem Memoria de borrador específica de mosaico. SRAM local y rápida disponible para un mosaico de SC específico.
202 Spmem Memoria de borrador compartida. Se usa para organizar datos de forma oportunista (entradas, salidas, datos intermedios) para ocultar la latencia de HBM.
203 HBM Memoria de alto ancho de banda. Memoria grande y compartida que se usa para incorporar tablas, montones y pilas.
204 Marcas de sincronización Primitivas de sincronización que se usan para la coordinación.

Para obtener información detallada sobre SC y su jerarquía de memoria, consulta la documentación de SparseCore.

Depuración

La resolución depende del espacio de memoria que falló en la asignación.

Situación 1: Fallas de asignación de HBM (ID del espacio de memoria: 203)

Este error se produce si una sola asignación temporal solicitada por el programa SparseCore es demasiado grande para caber en la HBM disponible. En las cargas de trabajo de incorporación estándar y los colectivos descargados de SC, las particiones por núcleo extremadamente grandes o las especificaciones de fragmentación incorrectas pueden obligar al compilador a solicitar búferes masivos.

Acciones recomendadas:

  • Verifica la fragmentación: Asegúrate de que tus tablas de incorporación y los tensores de entrada/salida de SC estén particionados o fragmentados correctamente. Si un solo núcleo es responsable de demasiados datos, es posible que falle la asignación.
  • Ajusta los límites: Revisa max_ids_per_partition y max_unique_ids_per_partition. Si estos se establecen de forma innecesariamente alta, el compilador reserva más memoria de la necesaria. Consulta Cómo se traducen los límites a las tablas.

Situación 2: Fallas de memoria interna (IDs del espacio de memoria: 0, 201, 202, 204)

Por lo general, las fallas de asignación en Smem, TileSpmem, Spmem o marcas de sincronización se producen debido a errores del compilador o limitaciones en la estrategia de asignación, en la que el compilador no tiene en cuenta todos los requisitos de memoria.

Acciones recomendadas:

  1. Aísla la operación XLA con errores: Para identificar el kernel específico de SC HLO o Mosaic que causa la falla, genera las representaciones intermedias del compilador:
    • Volcado de SparseCore MLIR: Establece la marca --xla_sc_dump_mlir_to=/path/to/dump. Esto genera el MLIR del programa SparseCore, lo que te permite ver qué tamaño de asignación coincide con el mensaje de error.
    • Volcado de Mosaic LLO: Para los kernels personalizados, usa --xla_mosaic_dump_to=/path/to/dump para inspeccionar todos los programas del optimizador de bajo nivel (LLO) que emite Mosaic.
  2. Reduce los tamaños de borrador (usuarios de Pallas): Si la falla se produce dentro de un kernel de Mosaic, revisa la configuración de scratch_shapes. Asegúrate de que tus solicitudes de pltpu.SMEM se ajusten a las especificaciones de hardware de tu generación de TPU específica.
  3. Inhabilita la descarga colectiva: Si el error surge de una operación colectiva descargada de SC, intenta inhabilitar las funciones de descarga de SC:
    • --xla_tpu_enable_sparse_core_collective_offload_all_gather=false
    • --xla_tpu_enable_sparse_core_collective_offload_all_reduce=false
  4. Informa un error: Si los pasos anteriores no resuelven el problema, es probable que se trate de un error del compilador. Envía un informe de errores.