Optimizaciones del rendimiento de TPU7x (Ironwood)

En esta guía, se describen varios métodos para optimizar el rendimiento con TPU7x (Ironwood) administrando de forma eficiente el movimiento de datos entre su sistema de memoria de varios niveles. Esto incluye técnicas como el entrenamiento de baja precisión, la fragmentación, la optimización de la comunicación, la rematerialización de la activación, el ajuste de la memoria vectorial (VMEM) con alcance y los kernels de aceleradores personalizados.

Para optimizar el rendimiento con la TPU7x, primero debes conocer la arquitectura de Ironwood, específicamente la jerarquía de memoria y la topología de interconexión. Para obtener más información, consulta TPU7x (Ironwood).

Entrenamiento de baja precisión con FP8

FP8 (punto flotante de 8 bits) es un formato de datos numéricos eficiente que se usa principalmente para acelerar el entrenamiento y la inferencia de modelos. Cuando se representan números con 8 bits, en lugar de los formatos estándar de 16 bits (FP16 o BF16) y 32 bits (FP32), las TPU pueden procesar datos mucho más rápido y usar menos memoria.

TPU7x admite la aceleración de hardware integrada para los tipos de datos FP8, lo que ofrece un rendimiento teórico máximo de 4,614 TFLOPS por chip. Esta capacidad puede generar tiempos de entrenamiento de extremo a extremo mucho más rápidos. Para las operaciones compatibles, en especial las multiplicaciones de matrices densas que son comunes para las cargas de trabajo de IA, el uso de FP8 puede generar mejoras de rendimiento de 1.3 veces en comparación con el entrenamiento estándar con BF16. En comparación con BF16, FP8 duplica los FLOP máximos y reduce a la mitad el espacio en memoria para los pesos y las activaciones. FP8 debería ser una palanca de ajuste principal para las cargas de trabajo vinculadas al procesamiento y las situaciones que están limitadas por la capacidad de memoria o el ancho de banda.

El uso de FP8 ofrece los siguientes beneficios de rendimiento:

• Presión reducida de la memoria de gran ancho de banda (HBM): Un espacio en memoria más pequeño permite que los modelos más grandes o los modelos con cachés de KV más grandes durante la inferencia quepan por completo en los 192 GiB de HBM. Esto evita la descarga costosa a la memoria del host más lenta.
• Mayor tamaño de lote efectivo: Al reducir la memoria necesaria para las activaciones, FP8 permite usar tamaños de lote más grandes. Esto mejora el paralelismo de datos y puede generar una mayor capacidad de procesamiento y un mejor uso de las unidades de procesamiento.
• Requisitos de ancho de banda de memoria más bajos: Mover la mitad de la cantidad de datos para cada operación reduce la demanda en la ruta de datos de HBM a MXU. En los sistemas en los que el movimiento de datos es un cuello de botella común, esto ayuda a mantener las MXU saturadas de trabajo.

El uso de FP8 con una degradación nula o limitada en el rendimiento requiere la selección cuidadosa de las técnicas de cuantización. Estas son algunas de las prácticas recomendadas que debes tener en cuenta para el entrenamiento con FP8:

• Granularidad del escalamiento: Comienza con el escalamiento por tensor como modelo de referencia. Si hay problemas de calidad o rendimiento, cambia al escalamiento por eje. Es posible que no sea necesario escalar los subcanales.
• Modo de escalamiento: El escalamiento dinámico, que calcula los factores de escala en el entorno de ejecución, es una buena opción predeterminada para mantener la calidad. Si bien el escalamiento estático puede ofrecer un aumento significativo del rendimiento, ya que elimina los cálculos, requiere un perfilado cuidadoso para determinar los factores de escala correctos y puede no ser adecuado para todos los casos de uso, en especial cuando cambian las configuraciones del modelo. Por el contrario, algunos modelos y parámetros de configuración sólidos pueden fijar la escala al límite de FP8 para pesos o activaciones, lo que te permite reducir la sobrecarga de cuantización y, al mismo tiempo, mantener la exactitud y mejorar el rendimiento.
• Formatos FP8 (E4M3 y E5M2): Un enfoque común y eficaz es usar una combinación de formatos FP8. Por ejemplo, usa E4M3 para los pesos y las activaciones en el pase hacia adelante para aprovechar la mayor precisión de E4M3, y usa E5M2 para los gradientes en el pase hacia atrás para adaptarte al rango dinámico más amplio de los gradientes.
• Redondeo: Usar el "redondeo al par más cercano" (RNE) en lugar del redondeo estocástico para los gradientes puede mantener la calidad y, al mismo tiempo, ofrecer un mejor rendimiento y reproducibilidad.
• Habilita FP8 en MaxText: MaxText admite el entrenamiento de FP8 a través de la biblioteca de cuantización de QWIX. Para activar la cuantización, establece la siguiente marca en tu configuración: use_qwix_quantization=true.

Fragmentación y paralelismo

La fragmentación es el proceso de dividir un modelo grande o sus datos de entrenamiento en partes más pequeñas y distribuirlas en varios chips TPU o núcleos. Elegir la estrategia de fragmentación adecuada es importante para lograr un alto rendimiento en TPU7x.

Un enfoque ingenuo que maximice puramente el grado de paralelismo suele generar un rendimiento deficiente, ya que se limita a la comunicación. El mejor enfoque suele ser seleccionar la estrategia de fragmentación más simple que cumpla con las restricciones de memoria, ya que esto minimiza la sobrecarga de comunicación y permite que las unidades de procesamiento se utilicen de manera eficiente.

Antes de seleccionar una estrategia de fragmentación, el primer paso en cualquier esfuerzo de ajuste del rendimiento debe ser un análisis de intensidad aritmética. Este análisis determina si un cálculo determinado está limitado por el procesamiento, el ancho de banda de la memoria o el ancho de banda de la interconexión. Se calcula como la proporción de operaciones de punto flotante con respecto a los bytes de datos que se deben mover.

Una intensidad aritmética alta indica una carga de trabajo vinculada al procesamiento. Una intensidad aritmética baja sugiere una carga de trabajo vinculada a la memoria o a la comunicación, en la que el rendimiento está limitado por la velocidad a la que se pueden mover los datos desde la HBM o a través de la red de ICI. Este análisis fundamenta el tamaño de lote ideal y la estrategia de fragmentación. Por ejemplo, una carga de trabajo vinculada a la comunicación no se beneficiará de una estrategia de fragmentación que introduzca aún más comunicación, como el paralelismo de tensores de alto grado.

Marco de decisiones para la estrategia de fragmentación

MaxText ofrece una variedad de estrategias de fragmentación. La mejor opción depende de la arquitectura del modelo, la longitud de la secuencia y la necesidad de equilibrar la carga computacional con la sobrecarga de comunicación.

• Paralelismo de datos completamente fragmentado (FSDP): Esta es la estrategia predeterminada preferida para el paralelismo de datos. El FSDP fragmenta los pesos, los gradientes y los estados del optimizador del modelo en los dispositivos de paralelismo de datos. Durante el procesamiento, cada dispositivo realiza una operación All-Gather para recuperar los pesos completos necesarios para su microlote local. FSDP es muy eficaz siempre que el tamaño del lote por dispositivo sea lo suficientemente grande como para ocultar la latencia de esta comunicación All-Gather. En el caso de los modelos de mezcla de expertos (MoE), el cálculo de la intensidad aritmética debe tener en cuenta la dispersión.
• Paralelismo de tensores (TP): El TP fragmenta tensores individuales en varios dispositivos. Por lo general, los tensores son matrices de pesos en perceptrones multicapa (MLP) y bloques de atención. La alta intensidad aritmética del hardware (11.5k) impone un requisito muy alto en las dimensiones del modelo para que el TP sea viable sobre ICI, y tratar de usar TP puede hacer que el sistema esté vinculado a la comunicación.
• Paralelismo de expertos (EP): Esta es la estrategia estándar y necesaria para entrenar modelos de MoE. EP fragmenta las capas de "expertos" en un conjunto de dispositivos, y se usa un colectivo de comunicación de todos a todos para enrutar los tokens a su dispositivo experto designado. El EP puede ser eficiente si la dimensión del MLP del modelo es lo suficientemente grande como para acercarse al límite de rendimiento.
• Paralelismo de contexto (CP): El CP es una estrategia especializada que es esencial para entrenar modelos con longitudes de secuencia muy largas. Su función principal es administrar el consumo de memoria de las activaciones, que crece de forma cuadrática con la longitud de la secuencia y puede superar la capacidad de la HBM. CP fragmenta la dimensión de la secuencia de los tensores de activación, lo que permite usar un tamaño de lote fraccionario por dispositivo. Debido a que el CP introduce más comunicación que el FSDP, la regla general es usar el grado mínimo de CP necesario para satisfacer las restricciones de memoria y garantizar que el fragmento del eje del lote siga siendo un número entero.

En la siguiente tabla, se asignan los tipos de cargas de trabajo comunes a la estrategia de fragmentación óptima:

Tipo de carga de trabajo	Fragmentación principal recomendada	Fragmentación secundaria	Cuellos de botella clave	Razones
Modelo denso: secuencia corta	FSDP	N/A	Rematerialización, Matmuls de FF	FSDP proporciona el mejor equilibrio. Con las secuencias cortas, es posible que la memoria de activación no sea una preocupación importante. La clave es un lote global lo suficientemente grande como para ocultar el All-Gather de pesos de FSDP. A medida que aumenta el tamaño del lote, aumenta el tamaño de la activación, y se requiere una política de rematerialización adecuada para garantizar que esta configuración no se quede sin memoria.
Modelo denso: secuencia larga	FSDP	CP	FlashAttention, memoria de activación	La memoria de activación se convierte en la restricción principal. El CP es necesario para habilitar los tamaños de lote fraccionarios por dispositivo y evitar problemas de memoria insuficiente (OOM). FlashAttention es la fuente dominante de procesamiento y tiempo perdido.
Modelo de MoE: secuencia corta	FSDP y EP	N/A	All-to-All (enrutamiento de expertos), rematerialización	Los modelos de MoE requieren EP para fragmentar los expertos. La comunicación de todos a todos para el enrutamiento de tokens es un cuello de botella importante que se debe superponer. La rematerialización también es una fuente importante de desperdicio.
Modelo de MoE: escala muy grande	FSDP + EP + PP	Paralelismo de modelos (MP)	Todos los cuellos de botella mencionados anteriormente, además de las burbujas de canalización	En el caso de los modelos que superan la memoria de un solo pod, se necesita PP para fragmentar las capas en los pods. Esto genera sobrecargas de comunicación de DCN y de burbujas de canalización. Esta es una configuración muy compleja que requiere un ajuste cuidadoso.

Optimización de la comunicación

El mecanismo principal para superponer la comunicación y el procesamiento en TPU7x se llama descarga colectiva de SparseCore. La arquitectura de Ironwood incluye unidades SparseCore dedicadas, que actúan como subprocesos de control independientes capaces de administrar el movimiento de datos en el tejido de ICI. Esto permite que las operaciones de comunicación colectiva (como All-Gather o Reduce-Scatter) se ejecuten en paralelo con los cálculos principales que se realizan en los TensorCores. Este es el método recomendado para los colectivos asíncronos en TPU7x. Usa las marcas recomendadas para habilitar la descarga de los colectivos más comunes.

Rematerialización de la activación

La rematerialización de activación, también conocida como puntos de control de gradiente, es una técnica fundamental para reducir la huella de HBM de un modelo. En lugar de almacenar todas las activaciones intermedias de la propagación hacia adelante en la HBM para usarlas durante la retropropagación, guarda solo algunas activaciones clave (puntos de control) y vuelve a calcular las demás a pedido durante la retropropagación. Esto ahorra una cantidad significativa de memoria a costa de un mayor procesamiento (aproximadamente entre un 25% y un 30% de FLOP adicionales para un bloque de Transformer estándar).

La decisión de qué tan agresivamente se debe aplicar la rematerialización es un parámetro de ajuste fundamental que depende por completo del cuello de botella principal, que suele variar según la longitud de la secuencia.

Para cargas de trabajo de secuencia larga (como 128k): En estos casos, el tamaño de los tensores de activación es el consumidor dominante de HBM. Por lo general, la carga de trabajo está vinculada a la memoria. Por lo tanto, aplicar una política de rematerialización agresiva es muy beneficioso. El ahorro de memoria permite que el entrenamiento continúe sin errores de memoria insuficiente y también permite tamaños de lote más grandes, y la sobrecarga de procesamiento de volver a calcular es una compensación que vale la pena.

Para cargas de trabajo de secuencia corta (como 8k): En estos casos, la memoria de activación es mucho menos preocupante y es más probable que la carga de trabajo esté vinculada al procesamiento. La sobrecarga de procesamiento de la rematerialización puede ser la mayor fuente de ineficiencia.

Ajusta las políticas de rematerialización en MaxText

MaxText proporciona un control detallado sobre la rematerialización a través de un conjunto de políticas predeterminadas y personalizadas, que se configuran con la marca remat_policy.

Políticas predeterminadas

MaxText ofrece las siguientes políticas integradas:

• full: Es la política más agresiva, que rematerializa casi todo. Esto minimiza el uso de HBM, pero maximiza la sobrecarga de recálculo. Es ideal para situaciones de secuencias largas con restricciones de memoria extremas.
• minimal: Es la política menos agresiva, que almacena la mayoría de las activaciones. Esto maximiza el uso de HBM, pero minimiza el recálculo. Es ideal para cargas de trabajo de secuencias cortas y vinculadas al procesamiento en las que la memoria no es un problema.
• Políticas intermedias: Las opciones como save_dot_with_context_except_mlp, save_qkv_proj y save_out_proj proporcionan varias compensaciones, ya que crean puntos de control de forma selectiva para los resultados de las operaciones de producto escalar costosas mientras se vuelven a materializar las operaciones por elementos más económicas.

Políticas personalizadas

Para obtener un mayor nivel de control, puedes establecer remat_policy en custom. Esto te permite especificar el comportamiento de las capas individuales dentro del módulo de decodificación del modelo. A cada capa se le puede asignar uno de tres comportamientos:

• device: La activación se almacena en la HBM en el dispositivo TPU.
• remat: La activación se descarta y se volverá a materializar durante el paso hacia atrás.
• offload: La activación se mueve de la HBM a la memoria del host de la CPU, lo que libera la HBM a costa de la latencia de transferencia de PCIe.

Ajuste de VMEM con alcance

El rendimiento del kernel, como la atención flash, depende de los tamaños de mosaico seleccionados en el kernel, cuyo tamaño está limitado por la memoria vectorial (VMEM) disponible. Cada uno de los dos TensorCores de un chip TPU7x tiene 64 MiB de memoria vectorial (VMEM). Esta capacidad de VMEM se puede dividir entre el alcance actual (VMEM con alcance) y la carga previa de pesos futura. Aumentar la VMEM con alcance permite aumentar los tamaños de los mosaicos en el kernel, lo que puede reducir las detenciones de memoria y aumentar el rendimiento de los kernels. Puedes modificar el tamaño de la VMEM con alcance configurando xla_tpu_scoped_vmem_limit_kib (en LIBTPU_INIT_ARGS), que se puede usar para explorar el rendimiento del kernel y los límites de rendimiento de extremo a extremo. Optimizar el tamaño de la VMEM con alcance puede afectar de forma indirecta el rendimiento del kernel de Pallas personalizado, ya que aumentar la VMEM con alcance desbloquea un espacio de búsqueda de hiperparámetros más grande para los tamaños de mosaicos en el kernel.

Kernels de Tokamax

Tokamax, una biblioteca de kernels de JAX de alto rendimiento con muchos kernels de TPU altamente optimizados, aborda varios cuellos de botella comunes específicos del hardware:

• Splash attention: Splash attention se usa como la implementación de atención principal para eliminar el cuello de botella de HBM de la atención estándar y usa la implementación de atención más eficiente en las TPU.
• Multiplicación de matrices agrupadas (GMM) de Megablox: Para las cargas de trabajo de MoE, Megablox controla de manera eficiente las multiplicaciones de matrices agrupadas a través del procesamiento en la representación de activaciones irregulares. Asigna de manera eficiente la dimensión irregular, calcula las multiplicaciones de matrices entre grupos irregulares de filas en LHS y la matriz de expertos correspondiente, lo que evita la necesidad de rellenar los lotes a un tamaño fijo.
• Ajuste empírico con tune-jax: La biblioteca tune-jax tiene utilidades para realizar búsquedas empíricas de tamaños de bloques óptimos. Los tamaños de kernel predeterminados suelen ser poco óptimos; el ajuste permite elegir tamaños de mosaicos de VMEM compatibles con el hardware para maximizar el uso del hardware.
• Estimación de logits máximos: El kernel de atención de Tokamax Splash se puede optimizar aún más si se establece un valor para max_logit_const. Si se establece, reemplaza el cálculo de reducción del logit máximo durante la operación softmax de atención (softmax(Q * KT)), lo que reduce cierta sobrecarga de procesamiento y sincronización. En MaxText, se implementa con el parámetro de configuración use_max_logits_estimate, que se puede establecer en None (inhabilitado) o en un valor de punto flotante. Verifica que el rango de logits de tu modelo específico siga siendo compatible con la estimación para evitar el desbordamiento numérico. Se recomienda realizar pruebas de convergencia si se establece este valor.