A principios de este año, Google lanzó el TPU de sexta generación y el más potente hasta la fecha: Trillium. Hoy, Trillium está oficialmente disponible para los clientes de Google Cloud.
Google utilizó el TPU Trillium para entrenar el último Gemini 2.0, el modelo de IA más potente de Google hasta la fecha. Ahora, las empresas y startups pueden aprovechar la misma infraestructura potente, eficiente y sostenible.
El núcleo de la supercomputadora de IA: TPU Trillium
El TPU Trillium es un componente clave de Google Cloud AI Hypercomputer. AI Hypercomputer es una arquitectura de supercomputadora innovadora que integra hardware optimizado para el rendimiento, software abierto, marcos ML líderes y un sistema de modelos de consumo flexible. Con el lanzamiento oficial del TPU Trillium, Google también ha realizado mejoras clave en la capa de software abierto de AI Hypercomputer, incluyendo la optimización del compilador XLA y marcos populares como JAX, PyTorch y TensorFlow, para lograr una relación precio-rendimiento líder en el entrenamiento, ajuste y servicio de IA.
Además, funciones como la descarga del host que utiliza DRAM de host a gran escala (que complementa la memoria de alto ancho de banda o HBM) ofrecen un mayor nivel de eficiencia. AI Hypercomputer le permite obtener el máximo valor de las más de 100.000 implementaciones de chips Trillium de cada arquitectura de red Jupiter, que cuenta con un ancho de banda bidireccional de 13 Petabits/segundo, capaz de escalar un único trabajo de entrenamiento distribuido a cientos de miles de aceleradores.
Clientes como AI21Labs ya están utilizando Trillium para entregar soluciones de IA significativas a sus clientes con mayor rapidez:
Barak Lenz, director de tecnología de AI21Labs, afirma: “En AI21, nos esforzamos continuamente por mejorar el rendimiento y la eficiencia de nuestros modelos de lenguaje Mamba y Jamba. Como usuarios de TPU v4 desde hace tiempo, estamos impresionados con las capacidades de Trillium de Google Cloud. Los avances en escala, velocidad y rentabilidad son muy significativos. Creemos que Trillium desempeñará un papel crucial en la aceleración del desarrollo de nuestros modelos de lenguaje complejos de próxima generación, permitiéndonos ofrecer a nuestros clientes soluciones de IA más potentes y accesibles.”
Mejora significativa del rendimiento de Trillium, varios indicadores baten récords
En comparación con la generación anterior, Trillium ha mejorado significativamente en los siguientes aspectos:
Rendimiento de entrenamiento mejorado en más de 4 veces
Rendimiento de inferencia mejorado en 3 veces
Eficiencia energética mejorada en un 67%
Rendimiento de cálculo máximo por chip mejorado en 4,7 veces
Capacidad de memoria de alto ancho de banda (HBM) duplicada
Ancho de banda de interconexión entre chips (ICI) duplicado
100.000 chips Trillium en una sola arquitectura de red Jupiter
Rendimiento de entrenamiento por dólar mejorado en 2,5 veces, rendimiento de inferencia por dólar mejorado en 1,4 veces
Estas mejoras permiten que Trillium sobresalga en diversas cargas de trabajo de IA, incluyendo:
Escalado de cargas de trabajo de entrenamiento de IA
Entrenamiento de LLM, incluyendo modelos intensivos y modelos de expertos mixtos (MoE)
Rendimiento de inferencia y programación de conjuntos
Modelos intensivos en incrustaciones
Proporcionar rentabilidad en entrenamiento e inferencia
¿Cómo destaca Trillium en diferentes cargas de trabajo?
Escalado de cargas de trabajo de entrenamiento de IA
Entrenar modelos grandes como Gemini 2.0 requiere una gran cantidad de datos y cómputo. La capacidad de escalado casi lineal de Trillium permite acelerar significativamente el entrenamiento de estos modelos mediante la distribución eficaz y eficiente de la carga de trabajo en múltiples hosts Trillium conectados a través de interconexiones entre chips de alta velocidad en pods de 256 chips y nuestra red de centros de datos Jupiter de vanguardia. Esto se logra mediante TPU multichip, una tecnología de pila completa para entrenamiento a gran escala, y se optimiza aún más mediante Titanium, un sistema de descarga dinámico a nivel de centro de datos que abarca desde adaptadores de host hasta la arquitectura de red.
Trillium logró una eficiencia de escalado del 99% en una implementación de 12 pods compuestos por 3072 chips y mostró una eficiencia de escalado del 94% en 24 pods con 6144 chips para el preentrenamiento de gpt3-175b, incluso cuando se ejecuta a través de una red de centros de datos.
Entrenamiento de LLM, incluyendo modelos intensivos y modelos de expertos mixtos (MoE)
Los LLM como Gemini son intrínsecamente potentes y complejos, con miles de millones de parámetros. Entrenar estos LLM intensivos requiere una enorme capacidad de cómputo y optimizaciones de software de diseño colaborativo. Trillium es 4 veces más rápido que la generación anterior de Cloud TPU v5e al entrenar LLM intensivos como Llama-2-70b y gpt3-175b.
Además de los LLM intensivos, el entrenamiento de LLM utilizando la arquitectura de expertos mixtos (MoE) es un enfoque cada vez más popular que combina múltiples redes neuronales "expertas", cada una especializada en diferentes aspectos de una tarea de IA. Gestionar y coordinar estos expertos durante el entrenamiento añade complejidad en comparación con el entrenamiento de un único modelo monolítico. Trillium es 3,8 veces más rápido que la generación anterior de Cloud TPU v5e al entrenar modelos MoE.
Además, en comparación con Cloud TPU v5e, Trillium TPU ofrece 3 veces más memoria dinámica de acceso aleatorio (DRAM) del host. Esto descargará algunos cálculos al host, lo que ayuda a maximizar el rendimiento a gran escala y un buen rendimiento. La función de descarga del host de Trillium proporcionó una mejora del rendimiento de más del 50% en la utilización de FLOP del modelo (MFU) al entrenar el modelo Llama-3.1-405B.
Rendimiento de inferencia y programación de conjuntos
En la inferencia, la importancia de la inferencia de varios pasos está creciendo, lo que requiere que los aceleradores manejen eficazmente las crecientes demandas computacionales. Trillium ofrece avances significativos para las cargas de trabajo de inferencia, permitiendo la implementación de modelos de IA más rápida y eficiente. De hecho, Trillium ofrece nuestro mejor rendimiento de inferencia de TPU para la difusión de imágenes y LLM intensivos. Nuestras pruebas muestran que el rendimiento de inferencia relativa (imágenes por segundo) de Stable Diffusion XL es más de 3 veces mayor en comparación con Cloud TPU v5e, mientras que el rendimiento de inferencia relativa (tokens por segundo) de Llama2-70B es casi el doble.
Trillium es nuestro TPU de mayor rendimiento para casos de uso de inferencia fuera de línea y de servidor. La siguiente figura muestra que el rendimiento de inferencia relativa fuera de línea (imágenes por segundo) de Stable Diffusion XL es 3,1 veces mayor y el rendimiento de inferencia relativa del servidor es 2,9 veces mayor en comparación con Cloud TPU v5e.
Además del mejor rendimiento, Trillium introduce nuevas funciones de programación de conjuntos. Esta función permite que el sistema de programación de Google tome decisiones inteligentes sobre la programación de trabajos para mejorar la disponibilidad y la eficiencia generales de las cargas de trabajo de inferencia cuando hay varias réplicas en un conjunto. Ofrece una forma de gestionar varios fragmentos de TPU que ejecutan cargas de trabajo de inferencia de un solo host o de varios hosts, incluyendo a través de Google Kubernetes Engine (GKE). Agrupar estos fragmentos en un conjunto permite ajustar fácilmente el número de réplicas para que coincida con las necesidades.
Modelos intensivos en incrustaciones
Con la adición de SparseCore de tercera generación, Trillium mejora el rendimiento de los modelos intensivos en incrustaciones en 2 veces y el rendimiento de DLRM DCNv2 en 5 veces.
SparseCore es un procesador de flujo de datos que proporciona una base de arquitectura más adaptable para cargas de trabajo intensivas en incrustaciones. SparseCore de tercera generación de Trillium destaca en la aceleración de operaciones dinámicas y dependientes de datos, como la dispersión-recopilación, la suma de segmentos dispersos y la partición.
Proporcionar rentabilidad en entrenamiento e inferencia
Además del rendimiento y la escala absolutos necesarios para entrenar algunas de las cargas de trabajo de IA más grandes del mundo, Trillium también está diseñado para optimizar el rendimiento por dólar. Hasta la fecha, Trillium ha mejorado el rendimiento por dólar en 2,1 veces en comparación con Cloud TPU v5e y en 2,5 veces en comparación con Cloud TPU v5p al entrenar LLM intensivos como Llama2-70b y Llama3.1-405b.
Trillium destaca en el procesamiento paralelo de modelos grandes de forma rentable. Está diseñado para permitir a los investigadores y desarrolladores ofrecer modelos de imágenes potentes y eficientes a un coste mucho menor que antes. El coste de generar mil imágenes en Trillium es un 27% menor para la inferencia fuera de línea y un 22% menor para la inferencia del servidor en SDXL en comparación con Cloud TPU v5e.
Elevar la innovación en IA a un nuevo nivel
Trillium representa un gran salto adelante en la infraestructura de IA de Google Cloud, ofreciendo un rendimiento, escalabilidad y eficiencia increíbles para una variedad de cargas de trabajo de IA. Con su capacidad para escalar a cientos de miles de chips utilizando software de diseño colaborativo de clase mundial, Trillium le permite lograr avances más rápidos y ofrecer soluciones de IA excepcionales. Además, la excelente relación precio-rendimiento de Trillium lo convierte en una opción rentable para las organizaciones que desean maximizar el valor de sus inversiones en IA. A medida que el panorama de la IA continúa evolucionando, Trillium demuestra el compromiso de Google Cloud de proporcionar infraestructura de vanguardia para ayudar a las empresas a liberar todo el potencial de la IA.
Presentación oficial: https://cloud.google.com/blog/products/compute/trillium-tpu-is-ga
