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

ME News Noticias, 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 de Ironwood TPU, para entrenar y desplegar de forma eficiente modelos de IA de vanguardia. Ironwood TPU es una infraestructura de IA personalizada diseñada para satisfacer las necesidades de cómputo de modelos de billones de parámetros. Mediante tecnologías como la interconexión entre chips (ICI), el conmutador de rutas ópticas (OCS), las redes del centro de datos (DCN) y la memoria de alto ancho de banda agregada (HBM), construye un sistema completo compatible con hasta 9.216 chips.

En el artículo se detallan varias estrategias clave de optimización para este hardware, incluyendo: aprovechar de manera nativa el soporte de entrenamiento FP8 en su unidad de multiplicación de matrices (MXU) para aumentar el rendimiento; usar la biblioteca de núcleos JAX optimizada para TPU, Tokamax, para procesar tensores irregulares en modelos de atención de “atención relámpago” y “multiplicación de matrices agrupadas Megablox” a fin de manejar contextos largos y modelos de expertos mixtos; utilizar la cuarta generación de núcleos dispersos (SparseCore) para descargar operaciones de comunicaciones colectivas y así ocultar la latencia; ajustar con precisión la asignación de la SRAM en el chip rápido de TPU (VMEM) para reducir la detención de memoria; y seleccionar la estrategia de particionado óptima (como FSDP, TP, EP) según el tamaño del modelo, la arquitectura y la longitud de la secuencia. (Fuente: InFoQ)