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

Noticias de ME News, 2 de abril (UTC+8): recientemente, Google publicó una guía de entrenamiento para desarrolladores dirigida a la séptima generación de TPU Ironwood. Esta guía tiene como objetivo ayudar a los desarrolladores a aprovechar al máximo el rendimiento a nivel de sistema de la TPU Ironwood, para entrenar y desplegar de manera eficiente modelos de IA de vanguardia. La TPU Ironwood es una infraestructura de IA personalizada diseñada para satisfacer las demandas de cómputo de modelos de billones de parámetros; mediante tecnologías como la interconexión entre chips (ICI), los conmutadores ópticos de ruta (OCS), la red de centro de datos (DCN) y la memoria de alto ancho de banda agregada (HBM), construye un sistema completo que admite hasta 9,216 chips. El artículo describe detalladamente varias estrategias clave de optimización para este hardware, incluyendo: aprovechar de forma nativa las unidades de multiplicación de matrices (MXU) para el entrenamiento en FP8 y mejorar el rendimiento de procesamiento; utilizar la biblioteca de núcleos JAX optimizada para TPU Tokamax, que procesa tensores irregulares en contextos largos y en modelos de expertos híbridos mediante “atención salpicada” y “multiplicación de matrices por grupos Megablox”; emplear la cuarta generación de núcleos dispersos (SparseCore) para descargar operaciones de comunicación colectiva y ocultar la latencia; ajustar con precisión la asignación de la SRAM rápida a bordo del TPU (VMEM) para reducir la detención de memoria; y seleccionar la mejor estrategia de particionado (como FSDP, TP, EP) según el tamaño del modelo, la arquitectura y la longitud de la secuencia. (Fuente: InFoQ)