Google lanza la séptima generación de la guía de entrenamiento para desarrolladores de Ironwood TPU, detallando la optimización del rendimiento a nivel de sistema

Noticias de ME News, 2 de abril (UTC+8), Google publicó recientemente una guía de entrenamiento para desarrolladores dirigida a la séptima generación de TPU Ironwood.
La guía tiene como objetivo ayudar a los desarrolladores a aprovechar al máximo el rendimiento a nivel de sistema del TPU Ironwood, para entrenar y desplegar modelos de IA de vanguardia de manera eficiente.
El TPU Ironwood está diseñado como una infraestructura de IA personalizada para satisfacer las demandas de potencia computacional de modelos con billones de parámetros, construida mediante tecnologías como interconexión entre chips (ICI), conmutadores ópticos (OCS), redes de centros de datos (DCN) y memoria de alta ancho de banda agregada (HBM), formando un sistema completo que soporta hasta 9,216 chips.
El artículo detalla varias estrategias clave de optimización para este hardware, incluyendo: aprovechar la unidad de multiplicación de matrices (MXU) para soportar entrenamiento en FP8 y aumentar el rendimiento; usar la biblioteca de núcleos JAX optimizada para TPU, Tokamax, que procesa tensores irregulares en modelos de contexto largo y modelos de expertos híbridos mediante “atención salpicada” y “multiplicación de matrices en bloques Megablox”; emplear el cuarta generación de núcleos escasos (SparseCore) para descargar operaciones de comunicación colectiva y ocultar latencias; ajustar finamente la asignación de la SRAM rápida en el chip TPU (VMEM) para reducir bloqueos de memoria; y seleccionar la estrategia de partición óptima (como FSDP, TP, EP) según el tamaño del modelo, la arquitectura y la longitud de la secuencia. (Fuente: InFoQ)