Categoría: CompileTime: SparseCore No Viable Logical Replica Count
Este error se produce cuando el compilador de XLA:SparseCore no puede determinar una configuración válida de recuento de réplicas lógicas que permita que la carga de trabajo quepa en la memoria de bloc de notas local de SparseCore (Tilespmem).
Ejemplos de mensajes de error:
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) ...
Back-ends 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 (LEM).
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 del bloc de notas. El compilador intenta dividir la carga de trabajo en fragmentos más pequeños (replicaciones) para que los búferes intermedios necesarios para cada fragmento quepan en la memoria de Scratchpad limitada de SparseCore. Por lo general, una mayor cantidad de réplicas lógicas reduce la presión de asignación, ya que procesa 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 de Tilespmem. El tamaño de la asignación se determina con una combinación de los siguientes factores:
sample_count: Es la cantidad de IDs de búsqueda de embeddings asignados a cada SparseCore (derivados del tamaño del lote).feature_width: Es el tamaño de la dimensión del embedding.max_nz_per_row: Es la cantidad máxima de IDs de búsqueda de incorporaciones no únicos en todos los SparseCores.max_unique_nz_per_row: Es la cantidad máxima de IDs de búsqueda de embeddings únicos.
Depuración
Para resolver este error, debes reducir la presión de memoria en el bloc de notas de SparseCore.
1. Mejora las estimaciones de metadatos
El compilador asigna memoria según 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 los IDs de tu conjunto de datos.
Puedes considerar la posibilidad de 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 valor de 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 bloc de notas necesarios.