Google detalla TPUv4 y su loca red de inteligencia artificial reconfigurable ópticamente

Aug 17, 2023

En Hot Chips 2023, Google mostró su loca red de inteligencia artificial reconfigurable ópticamente. La empresa está realizando conmutación de circuitos ópticos para lograr un mejor rendimiento, menor potencia y más flexibilidad para su grupo de entrenamiento de IA. Lo más sorprendente es que lo han tenido en producción durante años.

Esto se está haciendo en vivo, así que disculpen los errores tipográficos.

El gran objetivo de esto es unir los chips TPU de Google.

Aquí está el Google TPUv4 de 7 nm. Esperamos que esta semana comencemos a escuchar más sobre TPUv5. Google normalmente puede realizar artículos y presentaciones sobre hardware antiguo de una generación. La TPU v4i era la versión de inferencia, pero esta es más la charla centrada en TPUv4.

Google dice que aprovisiona energía en exceso en comparación con la energía típica para poder cumplir con un SLA de tiempo de servicio de 5 ms. Entonces, el TDP en los chips es mucho mayor, pero eso es para permitir que las ráfagas cumplan con esas ráfagas de SLA.

Aquí está el diagrama de arquitectura TPUv4. Google construye estos chips TPU no solo para que sean un acelerador único, sino para escalarlos y ejecutarlos como parte de una infraestructura a gran escala.

Aquí están las estadísticas de Google TPUv4 versus TPUv3 en una de las tablas más claras que jamás hayamos visto sobre esto.

Google ha más que duplicado el pico de FLOPS, pero ha reducido la potencia entre TPUv3 y TPUv4.

Google tiene un acelerador SparseCore integrado en TPUv4.

Aquí está el rendimiento de TPUv4 SparseCore de Google.

La placa en sí tiene cuatro chips TPUv4 y está refrigerada por líquido. Google dijo que tuvieron que reelaborar los centros de datos y las operaciones para cambiar a refrigeración líquida, pero el ahorro de energía vale la pena. La válvula de la derecha controla el flujo a través de los tubos de refrigeración líquida. Google dice que es como un controlador de velocidad de un ventilador, pero para líquido.

Google también dice que está utilizando PCIe Gen3 x16 en el host, ya que este era un diseño de 2020.

Google tiene energía que ingresa desde la parte superior del bastidor como muchos centros de datos, pero luego tiene varias interconexiones. Dentro de un bastidor, Google puede utilizar DAC eléctricos, pero fuera de un bastidor, Google necesita utilizar cables ópticos.

Cada sistema tiene 64 racks con 4096 chips interconectados. Hasta cierto punto, los clústeres de IA de NVIDIA en 256 nodos tienen la mitad de GPU.

También al final de los bastidores vemos un bastidor para CDU. Si desea obtener más información sobre la refrigeración líquida, puede ver Cómo funcionan los servidores de refrigeración líquida con Gigabyte y CoolIT. Pronto tendremos más contenido de refrigeración líquida. Google dice que el caudal de líquido es mayor que el del agua en la manguera de un camión de bomberos con gancho y escalera.

Cada bastidor es un cubo de 4x4x4 (64 nodos) con conmutación de circuito óptico (OCS) entre las TPU. Dentro del rack, las conexiones son DAC. Las caras del cubo son todas ópticas.

Aquí hay un vistazo al OCS. En lugar de utilizar un interruptor eléctrico, el uso del OCS proporciona una conexión directa entre chips. Google tiene matrices MEMS 2D internas, lentes, cámaras y más. Evitar toda la sobrecarga de la red permite compartir datos de manera más eficiente. En pocas palabras, esto en cierto modo se parece a los televisores DLP.

Google dijo que tiene más de 16.000 conexiones y suficiente distancia de fibra en el superpod como para rodear el estado de Rhode Island.

Debido a que hay tanta comunicación punto a punto, se requieren muchos hilos de fibra.

Más allá de eso, cada grupo se puede conectar a grupos más grandes.

El OCS, al ser reconfigurable, puede generar una mayor utilización de los nodos.

Luego, Google puede cambiar las topologías ajustando el enrutamiento óptico.

Aquí Google muestra los beneficios de diferentes topologías.

Esto es importante ya que Google dice que los cambios en las necesidades del modelo pueden impulsar cambios en el sistema.

Aquí está el escalado de Google en una escala logarítmica con aceleraciones lineales en hasta 3072 chips.

Google también aumentó la memoria del chip a 128 MB para mantener el acceso a los datos local.

Aquí está la comparación de Google con la NVIDIA A100 en términos de rendimiento por vatio.

Aquí está el modelo PaLM entrenando sobre 6144 TPU en dos pods.

¡Ese es un número enorme!

Ya es hora de que Google empiece a hablar del TPUv5, y resulta que Google NEXT es esta semana. Aún así, esta interconexión óptica es una tecnología realmente innovadora.

Algo que está bastante claro ahora es que Google está resolviendo grandes problemas con una enorme infraestructura. Tiene la oportunidad de avanzar más en el espacio de la IA. Es solo una cuestión de qué tan rápido Google comenzará a impulsar su hardware de inteligencia artificial y servicios en la nube contra NVIDIA y, al mismo tiempo, necesitará comprar GPU NVIDIA para sus clientes que las utilizan en lugar de TPU.