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

ME Noticias: El 2 de abril (UTC+8), Google publicó recientemente una guía de entrenamiento para desarrolladores dirigida a la séptima generación de Ironwood TPU. La guía busca ayudar a los desarrolladores a aprovechar al máximo el rendimiento a nivel de sistema de Ironwood TPU para entrenar y desplegar eficientemente modelos de IA de vanguardia. Ironwood TPU es una infraestructura de IA personalizada diseñada para satisfacer las necesidades de cómputo de modelos con billones de parámetros; mediante tecnologías como interconexiones entre chips (ICI), switches de intercambio de rutas ópticas (OCS), redes de centros de datos (DCN) y memoria de ancho de banda alto agregada (HBM), construye un sistema completo que admite hasta 9.216 chips. El artículo detalla varias estrategias clave de optimización para este hardware, incluyendo: aprovechar su unidad nativa de multiplicación de matrices (MXU) para el entrenamiento FP8 a fin de mejorar el rendimiento; utilizar la librería de núcleos de JAX optimizada para TPU, Tokamax, mediante “atención por salpicadura” y “multiplicación de matrices agrupadas Megablox” para manejar tensores irregulares en contextos largos y modelos de expertos mixtos; usar el cuarto núcleo disperso (SparseCore) para descargar operaciones de comunicaciones colectivas y así ocultar la latencia; ajustar con precisión la asignación de la SRAM rápida en el chip de TPU (VMEM) para reducir la detención de memoria; y, según el tamaño del modelo, la arquitectura y la longitud de la secuencia, seleccionar la mejor estrategia de particionado (como FSDP, TP, EP). (Fuente: InFoQ)