El Cerebras Wafer Scale Engine (WSE) es un procesador masivo de inteligencia artificial desarrollado por Cerebras Systems, una empresa fundada en 2015 por Andrew Feldman y otros. A diferencia de los chips convencionales que se cortan de una oblea de silicio, el WSE utiliza una oblea completa como un solo chip, integrando cientos de miles de núcleos y una enorme memoria en el chip. Este diseño tiene como objetivo reducir los cuellos de botella de comunicación y las limitaciones de memoria que afectan a los sistemas tradicionales basados en GPU para cargas de trabajo de aprendizaje automático y aprendizaje profundo.
El WSE-1 de primera generación, anunciado en 2019, contenía 1,2 billones de transistores y 400.000 núcleos optimizados para IA, lo que lo convertía en el chip más grande jamás construido en ese momento. El WSE-2 de segunda generación, lanzado en 2021, duplicó el número de transistores a 2,6 billones y aumentó el número de núcleos a 850.000, además de incrementar la memoria en el chip a 40 gigabytes. En 2024, Cerebras presentó el WSE-3, que mejoró aún más el rendimiento y la eficiencia, dirigido al entrenamiento de modelos de lenguaje grandes y otras aplicaciones de IA generativa.
Arquitectura y diseño
La arquitectura del WSE es fundamentalmente diferente a la de los procesadores convencionales. En lugar de depender de un pequeño número de núcleos potentes, el WSE empaqueta cientos de miles de núcleos simples y eficientes energéticamente en una sola oblea. Estos núcleos están interconectados por una red de malla que proporciona un ancho de banda masivo, permitiendo que los datos se muevan rápidamente entre núcleos. El WSE también integra una gran cantidad de memoria estática de acceso aleatorio (SRAM) directamente en el chip, reduciendo la necesidad de acceder a memoria externa, que a menudo es un cuello de botella de rendimiento en cargas de trabajo de IA.
El diseño a escala de oblea elimina la necesidad de empaquetado de chips y comunicación entre chips, que son comunes en sistemas de múltiples chips. Esto resulta en menor latencia y mayor rendimiento para tareas intensivas en datos. El WSE se fabrica utilizando procesos semiconductores avanzados, con el WSE-2 y el WSE-3 construidos en un proceso de 7 nanómetros por TSMC.
Software y ecosistema
Cerebras desarrolló una pila de software para programar el WSE, incluida la Plataforma de Software Cerebras (CSoft). CSoft es compatible con marcos populares de aprendizaje profundo como TensorFlow y PyTorch, permitiendo a los desarrolladores ejecutar modelos existentes con cambios mínimos en el código. La plataforma incluye un compilador que mapea los grafos de redes neuronales en los núcleos de la oblea, optimizando el rendimiento y el uso de memoria.
La empresa también ofrece el Cerebras Wafer-Scale Cluster, un sistema que combina múltiples WSE para escalar el entrenamiento de modelos muy grandes. Este clúster está diseñado para manejar modelos con billones de parámetros, que son comunes en arquitecturas modernas basadas en transformadores.
Rendimiento y aplicaciones
El WSE ha demostrado ventajas significativas de rendimiento tanto en entrenamiento como en inferencia. Por ejemplo, en 2023, Cerebras anunció que su sistema CS-2, impulsado por el WSE-2, podía entrenar un modelo de 175 mil millones de parámetros (similar en escala a GPT-3) utilizando un enfoque simple de paralelismo de datos, sin necesidad de paralelismo de modelos complejo. Esto se logró aprovechando la gran memoria en el chip y el alto ancho de banda del WSE.
El rendimiento de inferencia también es fuerte, con el WSE-3 capaz de generar más de 1.000 tokens por segundo para modelos de lenguaje grandes, según la empresa. Esto lo hace adecuado para aplicaciones en tiempo real como IA conversacional y generación de código.
El WSE ha sido adoptado por diversas organizaciones, incluidas agencias gubernamentales e instituciones de investigación. Por ejemplo, el centro de investigación atómica Bhabha en India ha utilizado sistemas Cerebras para computación científica. Además, Cerebras se ha asociado con g42 (una empresa holding de tecnología) para construir supercomputadoras de IA en Oriente Medio.
Comparación con otros chips de IA
El WSE compite con otros procesadores de IA especializados, como AWS Trainium de Amazon Web Services, los procesadores de flujo tensorial de Groq y las unidades de flujo de datos reconfigurables de SambaNova. A diferencia de las GPU de NVIDIA (que no están en la lista proporcionada pero son un competidor importante), el WSE se centra en maximizar la memoria en el chip y minimizar el movimiento de datos. Este enfoque puede ser más eficiente para ciertas cargas de trabajo, pero también presenta desafíos en el rendimiento de fabricación y la integración del sistema.
En comparación con la IPU (Unidad de Procesamiento de Inteligencia) de Graphcore, el WSE ofrece un equilibrio diferente: más núcleos pero menos flexibilidad en la programación. El diseño a escala de oblea del WSE también lo distingue de las ofertas de Intel y AMD, que utilizan empaquetado de chips tradicional con múltiples matrices.
Desafíos y limitaciones
Uno de los principales desafíos del WSE es el rendimiento de fabricación. Debido a que la oblea completa se utiliza como un solo chip, cualquier defecto en la oblea puede hacer que todo el chip sea inutilizable. Cerebras ha abordado esto implementando redundancia y técnicas de evitación de defectos, como mapear alrededor de núcleos defectuosos. Sin embargo, esto aún limita el tamaño máximo y aumenta el costo.
Otra limitación es el consumo de energía y la refrigeración. El WSE-2 consume hasta 15 kilovatios de potencia, lo que requiere soluciones de refrigeración especializadas. Cerebras utiliza un sistema personalizado de refrigeración líquida para disipar el calor de manera efectiva.
La compatibilidad del software también es una preocupación. Si bien CSoft es compatible con marcos populares, algunas características avanzadas u operaciones personalizadas pueden no estar completamente optimizadas. La empresa continúa expandiendo su ecosistema de software para mejorar la usabilidad.
Direcciones futuras
Cerebras continúa evolucionando la arquitectura del WSE. El WSE-3, anunciado en 2024, está diseñado para soportar modelos aún más grandes y tiempos de entrenamiento más rápidos. La empresa también está explorando aplicaciones más allá de la IA, como simulaciones científicas y computación de alto rendimiento.
En 2024, Cerebras presentó una solicitud de oferta pública inicial (IPO), lo que indica planes de expansión. La empresa tiene como objetivo competir más directamente con actores establecidos en el mercado de hardware de IA.
Conclusión
El Cerebras Wafer Scale Engine representa una desviación audaz del diseño de chips convencional, ofreciendo una escala y un rendimiento sin precedentes para cargas de trabajo de IA. Si bien enfrenta desafíos en fabricación y adopción, ha creado un nicho en el mercado de computación de IA de alta gama. A medida que los modelos de IA continúan creciendo, el enfoque del WSE de maximizar los recursos en el chip puede volverse cada vez más relevante.