Análisis detallado de SparseCore para modelos de embedding grandes (LEM)

SparseCore es un procesador especializado en mosaicos diseñado para la aceleración de alto rendimiento de cargas de trabajo que implican acceso y procesamiento irregulares y dispersos de la memoria, en particular en conjuntos de datos grandes almacenados en memoria de alto ancho de banda (HBM). Si bien se destaca en tareas como las búsquedas de embeddings, sus capacidades se extienden a la aceleración de una variedad de otras cargas de trabajo dinámicas y dispersas.

1. Introducción a SparseCore

Características arquitectónicas clave:

• Arquitectura en mosaicos: Comprende varios mosaicos de procesamiento (cada mosaico es una unidad completa de flujo de datos con su propia memoria local y unidad de procesamiento) que permiten el procesamiento paralelo.
• Ejecución dinámica: Admite de forma nativa el flujo de control y los accesos a la memoria que dependen de los datos, lo que es fundamental para los datos dispersos.
• Procesamiento vectorial: Utiliza tareas de vectores pequeños (de 8 o 16 elementos, según la versión de hardware) para realizar cálculos eficientes.
• Control centralizado: Un solo secuenciador de SparseCore organiza las tareas en todas las tarjetas, lo que garantiza operaciones sincronizadas.
• Compatibilidad con el resumen de datos: Incluye operaciones especializadas entre carriles que son beneficiosas para tareas como la clasificación, el filtrado y las sumas de prefijos.
• Jerarquía de memoria: Aprovecha estratégicamente la HBM para almacenar grandes conjuntos de datos y la memoria local de trabajo (SPMEM) para organizar los datos a los que se accede con frecuencia, lo que reduce significativamente la latencia de la HBM.

Especificaciones de un vistazo:

2. Preprocesamiento del host de SparseCore

La preparación de datos eficaz es fundamental para el rendimiento de SparseCore, y es aquí donde el procesamiento previo del host desempeña un papel vital. Abarca varias funcionalidades clave:

• Transformación de datos:
  • Aplica las transformaciones necesarias a los datos de entrada sin procesar.
  • Administrar las transformaciones de ID, lo que es especialmente importante cuando se trabaja con el apilamiento de atributos o tablas
  • Convierte los datos de entrada al formato disperso de coordenadas (COO), que se detalla en la siguiente sección.
  • Particiona los datos para distribuirlos de manera eficiente entre los diferentes SparseCores disponibles en el chip.
• Validación de límites:
  • Asegúrate de que las características de los datos de entrada (por ejemplo, la cantidad de IDs) cumplan con los límites operativos predefinidos de SparseCore, como max_ids_per_partition y max_unique_ids_per_partition.
  • Si los datos de entrada superan estos límites, la capa de preprocesamiento del host puede intentar segmentar los datos en lotes más pequeños que sí se ajusten a las restricciones.
• Transferencia de datos:
  • Copiar de manera eficiente los datos procesados y validados en la memoria de ancho de banda alto (HBM) de la TPU, lo que los prepara para la ejecución de SparseCore

Información sobre el apilamiento de tablas:

El apilamiento de tablas es una técnica de optimización significativa en la que se combinan lógicamente varias tablas de incorporación para mejorar la eficiencia de la búsqueda de incorporaciones. Por lo general, el framework de AA subyacente controla este proceso de forma automática.

• Apilamiento de funciones: Esto ocurre cuando varias funciones distintas comparten la misma tabla de incorporación subyacente. Un ejemplo común es usar un solo diccionario de incorporaciones para varios atributos categóricos, como los códigos postales de diferentes contextos.
• Apilamiento de tablas: En este caso, se apilan varias tablas de incorporación distintas. Las tablas que comparten la misma dimensión de incorporación y configuración del optimizador suelen agruparse.

La principal ventaja de la combinación de tablas es la creación de un tamaño de lote efectivo más grande para las operaciones en estas tablas combinadas. Esto reduce la sobrecarga computacional y puede ser eficaz para ocultar las latencias de comunicación entre chips (ICI). Para obtener un rendimiento óptimo, se recomienda una cantidad moderada de tablas apiladas (generalmente, entre 5 y 100).

3. Conversión a tensores de COO

Antes de que SparseCore pueda procesar los datos, estos suelen convertirse en un formato de tensor disperso de coordenadas (COO). El formato COO es una forma de representar matrices dispersas de manera eficiente, por lo general, con tres arrays:

• row_ids: Es un array que contiene los índices de fila para cada elemento distinto de cero. En el contexto del procesamiento por lotes, esto suele corresponder a la dimensión del lote.
• col_ids: Es un array que contiene los índices de columna para cada elemento distinto de cero. En el caso de las incorporaciones, suelen ser los valores de ID o de atributos.
• values (opcional): Es un array que contiene los valores reales de los elementos distintos de cero en las coordenadas (row, col) correspondientes. Para los cálculos de límites (que se explican más adelante) relacionados con los recuentos de IDs, estos valores (ganancias) a menudo no se tienen en cuenta.

Ejemplo ilustrativo:

Considera una matriz dispersa de entrada que representa lotes de IDs:

[
    [id_A],                 // Sample 0
    [id_A, id_B, id_C],     // Sample 1
    [id_B, id_B, id_D],     // Sample 2 (note duplicate id_B)
]

Después de la conversión al formato de COO (y, posiblemente, después de quitar los IDs duplicados dentro de la misma muestra):

row_ids = [0, 1, 1, 1, 2, 2]
col_ids = [id_A, id_A, id_B, id_C, id_B, id_D]

Esta conversión es fundamental para la forma en que SparseCore procesa y distribuye el trabajo. En particular, los col_ids son fundamentales para determinar a qué partición específica de SparseCore pertenece un ID, lo que permite una fragmentación y una búsqueda eficientes.

4. SparsecoreConfig: La API de alto nivel

APIs de embeddings específicas del framework:

• JAX: https://github.com/jax-ml/jax-tpu-embedding
• TensorFlow: https://www.tensorflow.org/recommenders/api_docs/python/tfrs/layers/embedding/TPUEmbedding
• Keras: https://keras.io/keras_rs/api/embedding_layers/distributed_embedding

El parámetro SparsecoreConfig, o los mecanismos equivalentes, como las marcas de XLA, sirven como una interfaz de alto nivel para controlar una amplia variedad de comportamientos de SparseCore. Comprender a fondo estos parámetros es fundamental para ajustar el rendimiento de manera eficaz y garantizar el funcionamiento correcto de tus modelos.

• disable_table_stacking: bool = False
  • Explicación: Esta marca controla si se evita el apilamiento automático de tablas, lo que podría reducir el rendimiento debido al aumento de la sobrecarga y a una menor capacidad para ocultar la latencia de interconexión entre chips (ICI).
  • Predeterminado: False (lo que implica que, por lo general, la función de apilamiento de tablas está habilitada de forma predeterminada cuando el framework la admite).
• max_ids_per_chip_per_sample: int = 64
  • Explicación: Este parámetro establece un límite superior global en la cantidad total de IDs de incorporación que un solo chip puede procesar a partir de una muestra en el lote de entrada, agregada en todas las tablas. Es un mecanismo para administrar recursos a nivel del chip, antes de que se tengan en cuenta los límites más detallados por tabla o por partición. Por lo general, el ajuste de este valor depende de las características específicas del modelo y de la capacidad general del sistema.
  • Predeterminado: 64.
• max_ids_per_table: Optional[Dict[str, int]] = None
  • Explicación: Este parámetro especifica la cantidad máxima de IDs de incorporación (que pueden incluir duplicados) que se pueden procesar para cada tabla lógica, teniendo en cuenta todas sus particiones en todos los SparseCores. Este es un límite más amplio que max_ids_per_partition. Si una tabla T se divide en P particiones, este límite se aplica a la suma de los IDs dirigidos a todas las particiones de P. A menudo, se relaciona con max_ids_per_partition_per_sample y el tamaño general del lote.
  • Configuración: Por lo general, se configura con un archivo de límites (por ejemplo, con la marca xla_sparse_core_max_ids_file), en el que se define max_ids_per_partition. Este concepto a nivel de la tabla es un método para establecer esos límites a nivel de la partición (max_ids y max_uniques).
  • Valor predeterminado: None (el valor se puede inferir a partir de los límites por partición o de otras configuraciones si no se proporciona de forma explícita).
• max_unique_ids_per_table: Optional[Dict[str, int]] = None
  • Explicación: Es análogo a max_ids_per_table, pero este parámetro especifica la cantidad máxima de IDs únicos para cada tabla lógica. Este es un parámetro de configuración fundamental para dimensionar de forma adecuada los búferes integrados en el dispositivo que se usan en el procesamiento de IDs únicos y las operaciones vectoriales posteriores.
  • Parámetro de configuración: También se define comúnmente en un archivo de límites o se deriva de max_unique_ids_per_partition_per_sample.
  • Predeterminado: None.
• allow_id_dropping: bool = False
  • Explicación: Esta marca booleana controla el descarte de IDs cuando la cantidad de IDs encontrados en los datos de entrada (límites observados) supera los límites establecidos durante la compilación (por ejemplo, max_ids_per_partition).
    • Si True: Los IDs que harían que se superen los límites se descartan de forma silenciosa. Por lo general, los IDs dentro de una partición se procesan en orden de clasificación, y se descarta cualquier ID que supere el límite de recuento en ejecución para su minibatch designado. Esto permite que el programa continúe su ejecución, pero puede tener un impacto negativo en la precisión del modelo.
    • Si False: Se activa un error y es probable que el proceso finalice si los límites observados superan los límites compilados. Este enfoque garantiza que se procesen todos los datos, pero requiere que los límites se configuren de forma más conservadora.
  • Valor predeterminado: False (lo que provoca un error por desbordamiento en lugar de una pérdida de datos silenciosa).

• initialize_tables_on_host: bool = True

  • Explicación: Esta marca determina si las tablas de incorporación se inicializan en la CPU del host antes de transferirse posteriormente a la memoria con ancho de banda alto (HBM) de la TPU. La práctica estándar es que las tablas se inicialicen en el host. Si se establece en True, se sigue esta convención. Si se configurara como False, implicaría un mecanismo de inicialización en el dispositivo, que podría tener diferentes implicaciones en el rendimiento o requisitos previos de inicialización específicos.

• enable_fast_table_initialization: bool = False

  • Explicación: Inicializa las tablas directamente en la TPU. Esto puede ayudar a reducir los tiempos de inicio del modelo.

5. Canalización para el rendimiento

El procesamiento en canalización es una técnica de optimización del rendimiento que permite la ejecución simultánea de operaciones en el TensorCore (TC) y el SparseCore (SC). Si se superponen estos cálculos, se puede mejorar significativamente la capacidad de procesamiento general.

• Mecanismo: En un paso de entrenamiento estándar que involucra búsquedas de incorporación dispersas (controladas por SC) y cálculos de capas densas (controlados por TC), la canalización permite que SC trabaje en su parte del paso i (por ejemplo, pase hacia adelante o hacia atrás) mientras que TC procesa de forma simultánea una parte diferente del mismo paso i, o incluso partes de pasos adyacentes como i-1 o i+1.
• Impacto en los gradientes: Es posible que SparseCore opere con gradientes "obsoletos". Por ejemplo, es posible que los gradientes calculados durante la fase de retropropagación del paso i no se actualicen por completo y no sean visibles para el SC hasta el paso i+2.
• Compensación entre rendimiento y valores numéricos: Esta ejecución superpuesta puede generar aceleraciones sustanciales, posiblemente hasta una mejora del doble en el tiempo de paso del dispositivo. Sin embargo, los cambios sutiles en los valores numéricos (embedding_weights) que resultan del uso de gradientes obsoletos pueden influir en el comportamiento de convergencia del modelo o en la exactitud final alcanzada. La aceptabilidad de esta compensación depende en gran medida del modelo y, a menudo, requiere validación empírica.
• Marca de control: La canalización se puede controlar con tf_xla_disable_full_embedding_pipelining. Si se configura esta marca como true, se inhabilita el procesamiento en canalización completo (cálculo superpuesto de TensorCore y SparseCore), mientras que, si se configura como false (o si la semántica de la marca implica la habilitación cuando es falsa), se activa.

Flujo conceptual de la canalización:

• Sin segmentación (flujo secuencial simplificado):

  Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i

• Con segmentación (flujo superpuesto simplificado):

  Time ->
  Step i:   SC/F_i | TC/F_i | TC/B_i | SC/B_i
  Step i+1:          SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1
  

  Nota: Las etapas de canalización reales implementadas en el hardware y el compilador pueden ser más complejas, y, a menudo, incluyen bucles previos, bucles de ejecución principales y bucles posteriores para administrar las dependencias de datos y garantizar la corrección.

6. El rol de XLA

XLA (Accelerated Linear Algebra) es el compilador específico del dominio que traduce los grafos computacionales de alto nivel, por lo general, de frameworks como TensorFlow, en código máquina altamente optimizado y adaptado para las TPU. Esto incluye generar las instrucciones para las operaciones destinadas a SparseCore.

Funciones clave en el contexto de SparseCore:

• Compilación de operaciones dispersas: XLA es responsable de compilar operaciones de búsqueda de incorporaciones (como SparseDenseMatmulOp) y otros cálculos dispersos en programas ejecutables de SparseCore de bajo nivel.
• Integración de límites: Utiliza los límites operativos configurados (por ejemplo, max_ids_per_partition, max_unique_ids_per_partition, que a menudo se proporcionan a través de un archivo de límites especificado por marcas como xla_sparse_core_max_ids_file) para determinar de forma estática los tamaños de los búferes de memoria en el dispositivo y asignarlos, en particular, dentro de SPMEM.
• Optimizaciones segmentadas: XLA realiza un conjunto de optimizaciones diseñadas específicamente para la arquitectura de SparseCore. Estas pueden incluir la programación de instrucciones, las transformaciones del diseño de la memoria y la fusión de operaciones para maximizar la eficiencia.
• Control con marcas: Muchos aspectos del comportamiento de SparseCore, los parámetros de ajuste y las estrategias de optimización se exponen y controlan a través de marcas de XLA (por ejemplo, xla_sparse_core_estimate_max_ids para la estimación de límites o xla_sc_detect_nan para la depuración).

Estado de código abierto:

Actualmente, la implementación de Sparsecore es interna y se realiza con libtpu.so.

Informes y diagnósticos de errores:

Las fallas de compilación relacionadas con la configuración de SparseCore o las restricciones de recursos a menudo se manifiestan como errores de tiempo de compilación de XLA:TPU. Estos mensajes de error pueden proporcionar información valiosa sobre problemas como límites demasiado altos para el SPMEM disponible o el uso de configuraciones no admitidas.

7. Cómo se traducen los límites en tablas en SparseCore

En SparseCore, los "límites" son parámetros de configuración fundamentales que se refieren principalmente a dos parámetros de configuración por partición para cada tabla que se fragmenta (distribuye) en los SparseCores disponibles:

• max_ids_per_partition: Define la cantidad máxima de IDs totales (incluidos los duplicados) que se espera que cualquier SparseCore individual envíe a una partición específica de una tabla determinada o procese para ella en un solo paso computacional.
• max_unique_ids_per_partition: Define la cantidad máxima de IDs únicos que se espera que cualquier SparseCore envíe o procese para un

Traducción al diseño y procesamiento de la tabla física:

• Estrategia de fragmentación de tablas: Por lo general, las tablas de incorporación se fragmentan con "mod" en todos los SparseCores del sistema. Esto significa que cada SparseCore se hace responsable de un subconjunto distinto del vocabulario (filas) de cada tabla. Por lo general, se asignaría un ID j a SparseCore_k según una fórmula como k = j % num_total_sparse_cores.
• Definición de "partición": En este contexto, una "partición" hace referencia al segmento específico de una tabla de incorporación para el que un solo SparseCore controla las búsquedas.
• Asignación de búferes de SPMEM: El compilador de XLA usa estos límites para asignar y dimensionar de forma estática los búferes dentro de la memoria de trabajo (SPMEM) del dispositivo. Los búferes se dimensionan de modo que todos los datos necesarios relacionados con los IDs de una partición determinada (hasta los límites especificados de max_ids y max_unique_ids) se puedan cargar en la SPMEM para su procesamiento. Esto es especialmente importante para los cálculos que no son elemento por elemento, como la reducción de IDs duplicados dentro de una partición (por ejemplo, al crear una representación de fila dispersa comprimida [CSR]), en los que todo el conjunto de datos pertinente para los IDs de esa partición debe estar disponible en la memoria rápida.

• Límites compilados en comparación con los límites observados:

  • Límites observados: Son la cantidad real de IDs encontrados para cada partición durante el tiempo de ejecución, según los datos de entrada que se procesan.
  • Si los límites observados superan los límites compilados, se puede producir una pérdida de ID (si allow_id_dropping está habilitado) o errores.

• Cálculo de límites: El proceso para determinar los límites adecuados implica un análisis cuidadoso de la distribución de los datos de entrada. Para cualquier tabla determinada (llamémosla T1, que podría formar parte de una tabla apilada más grande T):

  1. Inicialmente, el lote de entrada (por ejemplo, un SparseTensor 2D con forma [BatchSize, MaxSequenceLength]) se divide entre los SparseCores disponibles. Por ejemplo, si un TensorCore se combina con 2 SparseCores, cada SparseCore podría recibir un sublote con la forma [BatchSize/2, MaxSequenceLength].
  2. Luego, este sublote se convierte al formato COO, lo que genera row_ids y col_ids.
  3. Se quitan los IDs duplicados dentro de la misma muestra (es decir, las entradas con el mismo row_id y col_id).
  4. Para cada col_id único restante (dentro de una muestra), el SparseCore objetivo responsable de este ID se determina con la regla de fragmentación por módulo: target_sc_id = col_id % num_total_sparse_cores.
  5. Se mantiene un recuento de la cantidad total de IDs (ids_per_sparse_core[target_sc_id]++) y la cantidad de IDs únicos (unique_ids_per_sparse_core[target_sc_id]++, después de garantizar la unicidad para ese target_sc_id específico) que se destinan a cada target_sc_id.
  6. Luego, el valor de max_ids_per_partition para la tabla T1 se establece en max(ids_per_sparse_core_array).
  7. Del mismo modo, el max_unique_ids_per_partition para la tabla T1 se establece en max(unique_ids_per_sparse_core_array).
  8. Si la tabla T1 es un componente de una tabla apilada, es posible que se apliquen transformaciones adicionales, como rotaciones o desplazamientos, a las distribuciones de ID antes de sumar las estadísticas de todas las tablas constituyentes. Esto ayuda a equilibrar la carga entre los chips.

Establecer estos límites correctamente es un acto de equilibrio: los límites más bajos pueden generar un mayor rendimiento (ya que se deben procesar menos datos por paso y se reduce la presión de la SPMEM), pero si se establecen demasiado bajos, pueden generar un procesamiento por lotes excesivo o una eliminación no deseada de IDs.

8. Cómo se comunica cada SparseCore

La comunicación de SparseCore, en particular en el contexto del procesamiento de una lista de IDs para las búsquedas de incorporaciones, se basa en varios mecanismos coordinados:

• Fragmentación de mods y enrutamiento implícito:
  • Las tablas de incorporación se fragmentan de forma modular en todos los SparseCores del sistema.
  • Cuando el host proporciona un lote de datos de entrada (que luego se preprocesa en formato COO, incluido col_ids), el valor de col_id se usa para determinar qué SparseCore es responsable de ese ID específico: target_sc_id = col_id % num_total_sparse_cores.
  • Cada SparseCore recibe y procesa de manera eficaz solo el subconjunto de IDs que se asignan a sus particiones de vocabulario asignadas. La etapa de preprocesamiento del host es fundamental para preparar los datos de manera tal que cada SparseCore pueda identificar y operar fácilmente con sus IDs relevantes.
• Distribución de datos por host:
  • La lógica de preprocesamiento del host particiona el lote de entrada general y distribuye las partes pertinentes de row_ids y col_ids (junto con cualquier característica o peso asociado, si corresponde) a la memoria (HBM) a la que puede acceder directamente cada SparseCore o a una HBM compartida desde la que los SparseCores recuperarán los datos que necesiten.
• Procesamiento interno de SparseCore:
  • Una vez que un SparseCore recibe su conjunto designado de IDs para una partición de tabla determinada, realiza operaciones como la deduplicación de estos IDs y la recopilación de los vectores de incorporación correspondientes. Se trata principalmente de cálculos locales que se ejecutan dentro de los propios mosaicos de SparseCore y que utilizan su SPMEM local.
• Comunicación entre SparseCore (todos con todos):
  • Después de la fase de procesamiento inicial (como las búsquedas de incorporaciones), se puede usar un patrón de comunicación "todos con todos" para combinar o redistribuir los resultados en todos los SparseCores (por ejemplo, antes de ingresar las activaciones en una capa de TensorCore que espera una entrada correspondiente a todas las posiciones de la muestra original). Esto es fundamental para reconstruir el conjunto completo de activaciones si el lote de entrada original se distribuyó para el procesamiento paralelo.
• Comunicación con TensorCores:
  • Los SparseCores se comunican con los TensorCores para enviar activaciones de incorporación (durante el paso de avance) y recibir gradientes (durante el paso de retroceso). Esta interacción se coordina con el programa compilado por XLA y, con frecuencia, involucra a la HBM como un búfer intermedio. La estrategia de segmentación (analizada anteriormente) influye en gran medida en los tiempos y la sincronización de esta comunicación entre SC y TC.

En esencia, la "distribución" inicial de los IDs a los SparseCores adecuados se controla en gran medida con el esquema de fragmentación y los pasos de preprocesamiento del host. La comunicación posterior implica que los SparseCores operan en sus datos locales, lo que podría dar lugar a operaciones de comunicación colectiva, como la comunicación de todos con todos, si los datos deben intercambiarse o reordenarse de forma global entre los SparseCores antes de que los TensorCores los procesen aún más.

9. Administración de memoria de SparseCore

Cada SparseCore administra de manera eficiente varios tipos distintos de memoria para realizar sus cálculos:

• Memoria de bloc de notas (SPMEM):
  • Naturaleza: Es una SRAM local relativamente pequeña, pero muy rápida, que está disponible exclusivamente para cada SparseCore. Es importante tener en cuenta que SPMEM no es una caché; su uso lo administra y coordina de forma explícita el compilador de XLA.
  • Propósito: SPMEM se usa para "almacenar datos de forma oportunista". Esto incluye las entradas, las salidas y los resultados intermedios que se requieren para los cálculos continuos de SC. El almacenamiento temporal de datos en la SPMEM reduce significativamente la alta latencia que suele asociarse con el acceso a la HBM.
  • Tamaño: Como se mencionó en la sección sobre "Límites", los búferes de SPMEM tienen un tamaño estático en el momento de la compilación. Este tamaño se basa en parámetros como max_ids_per_partition y max_unique_ids_per_partition. Esta asignación estática garantiza que, para cualquier operación determinada en una partición de tabla (como la reducción de CSR), todos los datos necesarios para los IDs de esa partición (hasta los límites definidos) puedan caber en la SPMEM.
  • Optimizaciones del compilador: El compilador de XLA incorpora optimizaciones sofisticadas para determinar con precisión la cantidad de datos y los elementos de datos específicos que deben almacenarse en SPMEM para ocultar de manera eficaz la latencia de la HBM y maximizar el rendimiento.
  • Restricción de asignación dinámica: Por el momento, el compilador de SparseCore no admite la asignación dinámica de bloc de notas. Esto destaca la importancia fundamental del tamaño estático a través de la configuración cuidadosa de los límites.
• Memoria de alto ancho de banda (HBM):
  • Naturaleza: Es un recurso de memoria grande y compartido al que pueden acceder todos los SparseCores, los TensorCores y el sistema host. Las tablas de incorporación principales se almacenan en HBM.
  • Uso de la pila: Las operaciones de SparseCore a menudo requieren almacenamiento temporal en HBM para los resultados intermedios que no caben en la SPMEM limitada o que deben pasarse entre etapas más grandes de la canalización de procesamiento. El uso de la pila de HBM durante los pases hacia adelante y hacia atrás se puede estimar de la siguiente manera:
    • Pila de HBM de pase hacia adelante (una sola tabla) ≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4 bytes
    • Pila de HBM de pase hacia atrás (tabla única) ≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4 bytes
  • Uso del montón: La HBM también admite el montón, que administra el host. El montón almacena datos como los pesos de las capas densas, las constantes que usa el modelo y los datos de entrada recuperados previamente. El uso de la memoria dinámica tiende a aumentar con la cantidad de pasos para los que el host realiza una recuperación previa de los datos (controlada por la marca maximum_parallel_iterations). Si bien una mayor cantidad de recuperación previa puede mejorar el rendimiento, ya que superpone las transferencias del host al dispositivo con el procesamiento del dispositivo, también consume más HBM.
  • Serialización para la optimización de HBM: La marca xla_sc_num_serialized_tables_to_optimize_hbm proporciona un mecanismo para controlar cuántos datos de tablas se mantienen "activos" en la memoria de la pila de HBM en un momento determinado. Aumentar este número serializa de manera efectiva el procesamiento para más tablas, lo que puede reducir el uso máximo de la pila de HBM, pero puede afectar el rendimiento debido a la reducción del paralelismo.
• Memoria de vectores (VMEM):
  • La VMEM es una memoria local de borrador que usa exclusivamente el TC (TensorCore). Si bien SparseCore no administra directamente la VMEM, es una parte integral del ecosistema de memoria con el que interactúa el SC, principalmente a través de TensorCore.

Estrategia general de administración de memoria:

La estrategia principal de administración de memoria de SparseCore se centra en el uso de la SPMEM pequeña y rápida para los datos "activos" que procesa activamente una unidad de SparseCore, lo que minimiza los accesos a la HBM más lenta. Los límites configurados son el mecanismo principal para garantizar que SPMEM no se desborde. La HBM se utiliza para almacenar tablas de incorporación grandes y datos temporales que superan la capacidad de la SPMEM o que deben compartirse entre diferentes unidades de procesamiento o etapas de la canalización. El compilador de XLA es responsable de coordinar todo el movimiento de datos y la asignación de búferes según estos principios arquitectónicos y los límites configurados por el usuario.

10. Cuellos de botella en el rendimiento y la memoria

Para lograr un rendimiento óptimo con SparseCore, es necesario comprender claramente los posibles embotellamientos y cómo abordarlos. Estos pueden surgir en el host, dentro del propio SparseCore o en su interacción con los TensorCores.

Cuellos de botella habituales en el rendimiento:

• Cuello de botella del host:
  • Problema: Es posible que la CPU del host no pueda realizar el procesamiento previo de los datos y enviarlos a la TPU con la suficiente rapidez, lo que genera una subutilización de los SparseCores y los TensorCores. Este es un factor limitante del rendimiento frecuente.
  • Mitigación: Supervisa el uso de CPU del host y las métricas de la canalización de entrada. Optimiza las rutinas de carga y procesamiento previo de datos del host (consulta las sugerencias de conversión de COO). Ajusta la marca maximum_parallel_iterations para optimizar la recuperación previa de datos.
• Sincronización deficiente de TC/SC (falta de canalización):
  • Problema: Si la canalización entre TensorCore y SparseCore está inhabilitada o no funciona de manera eficiente, una unidad puede pasar mucho tiempo esperando a la otra, lo que reduce el rendimiento general del sistema.
  • Mitigación: Asegúrate de que el procesamiento en canalización esté habilitado (por ejemplo, tf_xla_disable_full_embedding_pipelining = false o su equivalente).
• Cuellos de botella inducidos por límites:
  • Problema:
    • Límites demasiado bajos: Pueden activar el procesamiento por lotes excesivo (dividir los lotes de entrada en numerosos sublotes más pequeños para cumplir con los límites estrictos). Si bien esto mantiene la corrección, cada lote pequeño introduce cierta sobrecarga de procesamiento, lo que podría ralentizar la ejecución general. Si allow_id_dropping es verdadero, los límites demasiado bajos también pueden provocar la pérdida de IDs, lo que afecta la precisión del modelo.
    • Límites demasiado altos (pero que aún se ajustan): Si bien los límites muy altos podrían impedir el procesamiento por lotes pequeños, podrían aumentar la presión de la SPMEM de forma innecesaria si las características reales de los datos rara vez se acercan a estos valores máximos. También podrían generar un uso de la pila de HBM mayor de lo estrictamente necesario.
    • Errores de compilación: Si los límites configurados requieren más SPMEM o pila de HBM que la memoria física disponible, la compilación fallará.
  • Mitigación: Asegúrate de que los límites estén configurados correctamente.
• Sesgo en la distribución de los datos:
  • Problema: Si ciertas particiones de SparseCore reciben de forma constante una cantidad desproporcionadamente mayor de IDs en comparación con otras (lo que indica una distribución deficiente de los IDs), esos SparseCores sobrecargados se convertirán en cuellos de botella de rendimiento.
  • Mitigación: La reorganización de IDs durante el proceso de mini-lotes puede ayudar a mitigar este problema en las tablas apiladas, en especial en aquellas con tablas de usuarios "populares". Analiza las distribuciones de ID con cuidado para establecer límites por tabla adecuados y equilibrados.
• Problemas de apilamiento de tablas:
  • Problema:
    • Demasiadas pocas tablas apiladas: Es posible que no sean suficientes para ocultar de manera eficaz la latencia del ICI ni para reducir adecuadamente los gastos generales de procesamiento.
    • Demasiadas tablas apiladas: Podría generar la creación de tablas lógicas muy grandes que se vuelven difíciles de administrar o que podrían exceder los límites de recursos disponibles.
  • Mitigación:
    • Asegúrate de que haya una cantidad óptima de mesas para apilar. Un lineamiento general sugiere un "punto óptimo" de 5 a 100 tablas para apilar.
• Cuantificación o valores numéricos ineficientes:
  • Problema: Usar la precisión FP32 completa cuando los formatos de menor precisión, como BF16 o los números enteros cuantificados, serían suficientes (y ofrecerían un procesamiento más rápido) puede ser un cuello de botella del rendimiento.
  • Mitigación: Explora opciones de menor precisión. Sin embargo, ten en cuenta que la cuantización en sí tiene cierta sobrecarga y podría requerir un ajuste cuidadoso de los parámetros de cuantización para mantener la precisión del modelo.
• Saturación del ancho de banda de la HBM:
  • Problema: El movimiento excesivo de datos hacia y desde la HBM, que podría deberse a anchos de características muy pequeños (lo que genera una sobrecarga de padding alta), patrones de acceso a la memoria ineficientes o una cantidad extremadamente grande de búsquedas, puede saturar el ancho de banda de la HBM disponible.
  • Mitigación: Ajustar la cantidad de TPU puede ayudar con la saturación del ancho de banda de la HBM.

Cuellos de botella de memoria comunes:

• Desbordamiento de SPMEM (falla de compilación):
  • Problema: Si max_ids_per_partition y max_unique_ids_per_partition se establecen en valores demasiado altos, es posible que el compilador de XLA no pueda asignar suficiente SPMEM, lo que generará errores de compilación como "Fixed size allocations (...) do not fit in TileSpmem (...)". Además, si el término (sample_count * feature_width) / kNumTiles (donde kNumTiles es la cantidad de mosaicos por SC) es demasiado grande para reunir operandos de preparación dentro de la SPMEM del mosaico, pueden producirse errores como "Gather operand too large...".
  • Mitigación: Reduce el tamaño del lote o aumenta la cantidad de chips que se usan para el procesamiento.
• Desbordamiento de pila de HBM (tiempo de ejecución o compilación):
  • Problema: Si la combinación de feature_width, max_unique_nz_per_row y logical_replica_count genera requisitos de memoria de pila de HBM que superan la HBM disponible, esto puede provocar errores de memoria insuficiente (OOM) durante el tiempo de ejecución o la compilación.
  • Mitigación: Ajusta la marca xla_sc_num_serialized_tables_to_optimize_hbm para reducir el uso de la pila de HBM serializando el procesamiento de las tablas (esto suele tener un costo de rendimiento).
• Agotamiento de la pila de HBM:
  • Problema: Se debe principalmente a pesos de capas densas muy grandes, numerosas constantes almacenadas en la memoria o una recuperación previa de datos de entrada demasiado agresiva (maximum_parallel_iterations alto).
  • Mitigación: Supervisa el uso del montón con herramientas como XProf Memory Viewer.
• Sobrecarga de padding:
  • Problema: Las tablas de incorporación se completan con ceros para que se alineen con 32 B (equivalente a 8 números de punto flotante) en la dimensión del atributo. Por lo tanto, los anchos de atributos pequeños (por ejemplo, 1 número de punto flotante) generan una sobrecarga de padding significativa (por ejemplo, 7/8 del espacio de búfer asignado es padding), lo que genera un desperdicio de HBM. La dimensión del vocabulario de las tablas también se completa para que sea un múltiplo de la cantidad de SparseCores en el sistema. Sin embargo, este impacto suele ser insignificante para las tablas con un tamaño de vocabulario suficientemente alto.

Factores generales que afectan el rendimiento y la memoria:

• Topología: Es la cantidad de chips disponibles y su arquitectura de interconexión.
• Tamaño del lote: Afecta directamente el sample_count por SparseCore, lo que, a su vez, influye en el consumo de memoria y la carga de procesamiento.
• Formato de datos: Garantizar un diseño de datos eficiente en el dispositivo es fundamental para un rendimiento óptimo.

11. Cómo analizar un perfil de SparseCore

Analizar un perfil de rendimiento es un paso clave para identificar cuellos de botella y descubrir oportunidades de optimización en tus cargas de trabajo de SparseCore.

1. Obtén un registro:
   • Usa herramientas de generación de perfiles, como XProf, para capturar un registro de ejecución detallado mientras se entrena tu modelo o se ejecuta la inferencia. Este registro proporcionará un cronograma de las operaciones que ocurren en el host, los TensorCores y los SparseCores.
2. Examina el visualizador de seguimiento (por ejemplo, en XProf o TensorBoard):
   • Actividad del organizador: Analiza la actividad del organizador. ¿Hay brechas significativas en la actividad de la TPU? Estas brechas pueden indicar que el host es un cuello de botella y no puede proporcionar datos con la suficiente rapidez. Analiza el rendimiento de tu canalización de entrada.
   • Actividad de TensorCore (TC) y SparseCore (SC):
     • Observa las líneas de tiempo de ejecución de TC y SC. ¿Operan en paralelo, lo que indica una canalización eficaz? ¿O hay períodos prolongados en los que una unidad está inactiva esperando a la otra?
     • Identifica las operaciones que consumen más tiempo (las operaciones de mayor duración) tanto en el SC como en el TC.
     • Los resultados del registro visual (que suelen mostrar bloques de colores que representan diferentes operaciones a lo largo del tiempo, como TPU:0 SparseCore 1 (pid 1005)) son muy valiosos para identificar visualmente las operaciones dominantes y los períodos de inactividad.
   • Análisis del tiempo de pasos: Observa el tiempo de pasos general y comprende cómo se distribuye o desglosa entre el procesamiento del host, el cálculo de SC y el cálculo de TC.
3. Análisis de memoria (visualizador de memoria de XProf):
   • Uso del montón: Usa herramientas como la pestaña "Memory Viewer" de XProf para inspeccionar el uso del montón de HBM. Esto puede ayudar a determinar si los pesos del modelo grandes, las constantes o la recuperación previa de entrada demasiado agresiva consumen una cantidad excesiva de HBM. Habilitar marcas como --vmodule=best_fit_allocator=1 puede proporcionar registros del uso máximo del montón.
   • Uso de la pila (indirecto): Si bien el perfilamiento directo de la pila de HBM puede ser complejo, si encuentras errores de falta de memoria y el uso del montón parece razonable, el agotamiento de la pila de HBM (a menudo debido a límites o anchos de características demasiado grandes) es un fuerte sospechoso. Las fórmulas proporcionadas para el uso de la pila de HBM pueden ayudar a estimar esto.
4. Busca patrones específicos:
   • Minilote: Si los límites se superan con frecuencia, es posible que observes evidencia de minilote en el registro (por ejemplo, una mayor cantidad de operaciones SC más pequeñas de lo esperado para el tamaño del lote global). A menudo, esto se puede inferir a partir de los registros o de la observación de los recuentos de invocaciones de ciertas operaciones.
   • Descarte de ID: Si el descarte de ID está habilitado y se produce, los registros del sistema podrían proporcionar indicaciones de esto. Esto también sería un signo claro de que los límites configurados son demasiado restrictivos para los datos de entrada.
   • Tiempos de compilación: Los tiempos de recompilación extendidos, en especial si la optimización dirigida por comentarios (FDO) está habilitada y se ajustan los límites con frecuencia, pueden agregar una sobrecarga significativa al tiempo de entrenamiento general.
5. Correlaciona con marcas y configuración:
   • Relaciona el comportamiento observado en el perfil con tus configuraciones de SparseCore (parámetros de configuración en archivos de límites, marcas de XLA). Por ejemplo, si xla_sc_num_serialized_tables_to_optimize_hbm se establece en un valor alto, es posible que el rendimiento de la SC sea más lento, pero el consumo de la pila de HBM sea menor.
6. Proceso iterativo:
   • La generación de perfiles suele ser un proceso de perfeccionamiento iterativo. Realiza un cambio específico (ajusta un límite, habilita o inhabilita una función), captura un perfil nuevo y, luego, compáralo con el perfil anterior para ver el impacto de tu modificación.

12. Marcas de depuración generales

Se pueden habilitar varias marcas para ayudar a depurar problemas relacionados con la ejecución de SparseCore. Es importante tener en cuenta que habilitar estas verificaciones a menudo genera una penalización en el rendimiento y, por lo tanto, generalmente deben inhabilitarse para las ejecuciones de producción.

• Verificaciones de ID (fuera de rango):
  • Bandera: xla_sparse_core_enable_id_bound_check = true
  • Propósito: Permite realizar verificaciones en el sistema host para detectar si algún ID de incorporación en los datos de entrada se encuentra fuera del rango de vocabulario válido definido para una tabla de incorporación determinada. Esto ayuda a detectar problemas relacionados con datos de entrada incorrectos o dañados.
• Verificador de NaN:
  • Bandera: xla_sc_detect_nan = true
  • Propósito: Permite la detección de valores NaN (no es un número) dentro de los datos de punto flotante procesados en SparseCore. Si se detecta un NaN en las entradas o salidas de varios pases del compilador, esta marca generará un error. Por lo general, estos errores proporcionan información sobre dónde se encontró el valor NaN.
• Verificador de límites (acceso a la memoria):
  • Bandera: xla_sc_assert_level=bounds
  • Propósito: Esta marca habilita una herramienta de estilo ASAN (AddressSanitizer) que reescribe las instrucciones de acceso a la memoria (como las cargas y los almacenamientos de VMEM, y las operaciones de DMA) para incluir verificaciones dinámicas. Estas verificaciones comprueban si el acceso a la memoria se encuentra dentro de los límites asignados de la región de memoria de destino.
  • Comportamiento: Si se detecta un acceso a la memoria fuera de los límites, la ejecución fallará.
  • Precaución: Es posible que este verificador produzca falsos positivos, por ejemplo, debido a patrones de acceso complejos con pasos que el verificador no comprende por completo. Esta transformación se aplica en una etapa avanzada del proceso de compilación del backend.
• Verificador de búfer (corrupción de memoria):
  • Marcadores:
    • xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'
    • xla_tpu_verify_launch_id_across_cores=true
  • Propósito: Estas marcas ayudan a garantizar que los búferes de memoria no se dañen ni se sobrescriban de forma involuntaria por operaciones no relacionadas. El Buffer Sanitizer verifica el contenido de los búferes para comprobar que no cambien de forma inesperada.

13. Compatibilidad con la cuantización

El SparseDenseMatmulOp de SparseCore está diseñado para admitir operaciones en tablas de incorporación con tipos de datos de punto flotante (FP32) y de números enteros de 32 bits. Si bien el entrenamiento del modelo suele realizarse con precisión FP32 para las tablas de incorporación, se puede aplicar la cuantización posterior al entrenamiento (PTQ). La PTQ permite el uso de tipos de datos de menor precisión (como números enteros de 8 bits) para la inferencia, lo que puede mejorar el rendimiento y reducir la huella de memoria.

Cuantización simulada:

El SparseDenseMatmulOp se puede configurar para realizar una "cuantificación simulada". En este modo operativo, los vectores de incorporación se cuantifican primero a una precisión más baja y, luego, se decuantifican a una precisión más alta (por ejemplo, FP32) antes de usarse en cálculos posteriores. Esta técnica permite entrenar modelos teniendo en cuenta los efectos del ruido de cuantización. El entrenamiento con cuantización simulada puede mejorar la exactitud del modelo final cuando se cuantiza por completo para la inferencia.

Atributos de configuración para SparseDenseMatmulOp (para la cuantificación):

• quantization_config_num_buckets = 256
  • Con este atributo, se especifica la cantidad de discretos buckets o niveles en los que se cuantificará un número de punto flotante de 32 bits. Por ejemplo, cuando se cuantifica en números enteros de 8 bits, se suelen especificar 2^8 =256 buckets.
• quantization_config_low = -X.X
  • Este atributo define el valor de punto flotante mínimo en el rango de cuantificación. Los valores de entrada inferiores a este mínimo especificado se recortarán a este valor mínimo durante la cuantificación.
• quantization_config_high = Y.Y
  • Este atributo define el valor de punto flotante máximo en el rango de cuantificación. Cualquier valor de entrada superior a este máximo especificado se recortará a este valor máximo durante la cuantificación.

Interacción de canalización y datos numéricos:

El comportamiento numérico del modelo puede cambiar según si está habilitado el procesamiento en paralelo entre TensorCore y SparseCore. Si el procesamiento en canalización está activo, es posible que los gradientes procesados por SparseCore estén "obsoletos" (de una iteración anterior). Esto puede interactuar con el proceso de cuantización y afectar potencialmente la dinámica del entrenamiento del modelo o la precisión final.

14. Próximas funciones y mejoras recientes

El ecosistema de SparseCore está sujeto a un desarrollo y una mejora continuos.

Hoja de ruta:

• Mini-lotes de dimensiones de muestra:
  • Esto se planea como una función complementaria a las capacidades existentes de procesamiento por lotes pequeños de la dimensión del vocabulario.
  • Permitiría una mayor partición de las entradas de la incorporación a lo largo de la dimensión de la muestra. Esto se lograría introduciendo bucles en el dispositivo que pueden filtrar y procesar búsquedas de un subconjunto de muestras a la vez. Esta función podría ser beneficiosa para administrar recuentos de IDs por muestra muy grandes o para mejorar el balanceo de cargas en las unidades de procesamiento.
• Se mejoró la compatibilidad con la incorporación de menos de 8 números enteros por fila (ancho de funciones pequeño):
  • El diseño actual suele usar un padding significativo para los anchos de las funciones de incorporación que son inferiores a 8 números de punto flotante (lo que corresponde a 32 bytes). Este padding puede generar un desperdicio de HBM y, posiblemente, un uso insuficiente de los recursos de procesamiento. Las mejoras futuras tienen como objetivo mitigar esta ineficiencia en las tablas con dimensiones de características pequeñas.

Mejoras recientes:

• Almacenamiento temporal de operandos de recopilación en HBM:
  • Esta optimización ayuda a reducir la presión en la memoria de bloc de notas compartida (SPMEM), ya que permite que algunas entradas o salidas de la operación de recopilación se transfieran a la HBM más grande.
• Reducción en el uso de memoria de la pila:
  • Se implementaron mejoras para reducir el consumo de memoria de la pila HBM durante las operaciones de SparseCore, idealmente sin afectar de forma negativa el rendimiento o la capacidad de procesamiento generales.

Estas mejoras se enfocan en optimizar el rendimiento, la eficiencia de la memoria y la flexibilidad operativa de SparseCore para una gama aún más amplia de cargas de trabajo dispersas.