Código de error: E3001

Categoría: CompileTime: SparseCore No Viable Logical Replica Count

Este error se produce cuando el compilador XLA:SparseCore no puede determinar una configuración válida de recuento de réplicas lógicas que permita que la carga de trabajo se ajuste a la memoria de borrador local de SparseCore (Tilespmem).

Mensajes de error de muestra:

XLA:TPU compile permanent error. Compilation failure: No viable logical replica count for the embedding table with metadata: max_nz_per_row = 141352, max_unique_nz_per_row = 8, feature_width = 8, sample_count = 204800 (last tried split factor for vector splitting = 1, last tried split factor for sample dimension splitting = 1, fixed_size_allocation_bytes = 410880, row_dependent_size_allocation_bytes = 1696224, total_spmem_size_bytes = 524288) ...

Backends de XLA: TPU

Descripción general

Este error es específico para los casos de uso de SparseCore, en particular para los modelos de embedding grandes (LEMs).

El recuento de réplicas lógicas es un parámetro interno del compilador que determina cómo se particionan los lotes de entrada para administrar la presión de asignación de la memoria de borrador. El compilador intenta dividir la carga de trabajo en fragmentos más pequeños (réplicas) para que los búferes intermedios necesarios para cada fragmento quepan en la memoria de borrador limitada de SparseCore. Por lo general, un recuento de réplicas lógicas más alto reduce la presión de asignación procesando lotes de datos más pequeños a la vez.

Este error indica que, incluso después de intentar varias configuraciones de división, el compilador no pudo encontrar una configuración en la que los búferes requeridos quepan en la memoria Tilespmem. El tamaño de la asignación se determina mediante una combinación de lo siguiente:

  • sample_count: Es la cantidad de IDs de búsqueda de embedding asignados a cada SparseCore (derivados del tamaño del lote).
  • feature_width: Es el tamaño de la dimensión de embedding.
  • max_nz_per_row: Es la cantidad máxima de IDs de búsqueda de embedding no únicos en todos los SparseCores.
  • max_unique_nz_per_row: Es la cantidad máxima de IDs de búsqueda de embedding únicos.

Depuración

Para solucionar este error, debes reducir la presión de memoria en la memoria de borrador de SparseCore.

1. Mejora las estimaciones de metadatos

El compilador asigna memoria en función de max_nz_per_row y max_unique_nz_per_row. Si estos valores se estiman de forma conservadora (es decir, se establecen mucho más altos de lo que requieren los datos reales), el compilador reservará espacio innecesario, lo que provocará este error. Asegúrate de que estos parámetros reflejen con precisión la distribución real de IDs de tu conjunto de datos.

Puedes considerar aplicar la optimización dirigida por comentarios (FDO) para determinar los valores óptimos de estos parámetros.

2. Reduce el tamaño del lote

El parámetro sample_count se deriva directamente del tamaño del lote global. Reducir el tamaño del lote disminuye la cantidad de datos que cada SparseCore debe procesar por paso, lo que reduce el tamaño de los búferes de memoria de borrador requeridos.