El Wafer-Scale Engine (WSE) es una serie de procesadores a gran escala desarrollados por Cerebras Systems, una empresa fundada en 2015 por Andrew Feldman, Gary Lauterbach y otros. A diferencia de los chips convencionales que se cortan de una oblea de silicio, el WSE utiliza la oblea completa como un solo dado, lo que permite un número sin precedentes de núcleos y memoria en el chip. Este diseño tiene como objetivo reducir los cuellos de botella de comunicación que afectan a los sistemas de múltiples chips, lo que lo hace particularmente adecuado para entrenar modelos de lenguaje grandes y otras cargas de trabajo de aprendizaje profundo.
La primera generación, WSE-1, se anunció en agosto de 2019. Contenía 1,2 billones de transistores y 400.000 núcleos optimizados para IA, fabricados con un proceso de 16 nanómetros por TSMC. El chip medía 46.225 milímetros cuadrados, más de 50 veces más grande que las unidades de procesamiento gráfico más grandes de la época. La segunda generación, WSE-2, siguió en abril de 2021, utilizando un proceso de 7 nanómetros, duplicando el número de núcleos a 850.000 y aumentando la memoria en el chip a 40 gigabytes. En marzo de 2024, Cerebras presentó el WSE-3, construido con un proceso de 5 nanómetros, con 900.000 núcleos y 44 gigabytes de SRAM en el chip, dirigido al entrenamiento de modelos con hasta 24 billones de parámetros.
Arquitectura y Diseño
La innovación central del WSE es su integración a escala de oblea. En lugar de fabricar muchos dados pequeños en una oblea y empaquetarlos por separado, Cerebras mantiene la oblea completa intacta, con circuitos redundantes para enrutar alrededor de los defectos de fabricación. Este enfoque elimina la necesidad de comunicación entre chips a través de una placa de circuito impreso, que a menudo es un cuello de botella de rendimiento en sistemas distribuidos. El chip utiliza una red de núcleos conectados en malla, cada uno con su propia memoria, lo que permite un movimiento de datos de alta banda ancha y baja latencia.
Cada núcleo es un procesador simplificado optimizado para las operaciones de álgebra lineal dispersa comunes en redes neuronales. Los núcleos admiten una variedad de formatos de datos, incluidos FP16, BF16 y FP32, y están diseñados para ejecutar operaciones densas y dispersas de manera eficiente. La memoria en el chip se distribuye entre los núcleos, proporcionando una banda ancha total que supera con creces la de los sistemas de memoria de GPU tradicionales. Esta arquitectura reduce la dependencia de la memoria externa, que a menudo es un factor limitante en el entrenamiento de IA.
El WSE no es un producto independiente; se integra en los sistemas CS-2 y CS-3 de Cerebras, que incluyen un sistema de suministro de energía y una solución de refrigeración. El CS-2, lanzado en 2020, alberga el WSE-2, mientras que el CS-3, anunciado en 2024, alberga el WSE-3. Estos sistemas están diseñados para conectarse a centros de datos existentes y a menudo se utilizan en clústeres para entrenamiento a gran escala.
Rendimiento y Puntos de Referencia
Cerebras ha publicado varios puntos de referencia de rendimiento para el WSE. En 2021, la empresa demostró que un solo sistema CS-2 podía entrenar un modelo de 1.800 millones de parámetros utilizando la arquitectura GPT, logrando una escalabilidad casi lineal cuando se conectaban múltiples sistemas. En 2022, Cerebras y los Laboratorios Nacionales Sandia utilizaron un CS-2 para simular una red neuronal de picos de 200 millones de neuronas en tiempo real, una tarea que requeriría una supercomputadora con miles de GPU.
En 2023, Cerebras informó que sus sistemas podían entrenar un modelo de 7 mil millones de parámetros en menos tiempo que un clúster de 16 NVIDIA A100 GPU, debido a la reducción de la sobrecarga de comunicación. El WSE-3, según la empresa, puede ofrecer hasta 125 petaflops de cómputo de IA, lo que lo convierte en uno de los procesadores de un solo chip más potentes para IA. Sin embargo, los puntos de referencia independientes son limitados, y la mayoría de las afirmaciones de rendimiento provienen de la propia Cerebras.
Ecosistema de Software
Para programar el WSE, Cerebras desarrolló la Plataforma de Software Cerebras (CSoft), que incluye un compilador que mapea modelos de aprendizaje automático a la arquitectura del chip. CSoft admite marcos populares como TensorFlow y PyTorch, lo que permite a los investigadores ejecutar modelos existentes sin cambios significativos en el código. El compilador maneja la programación del flujo de datos, la asignación de memoria y la optimización para el diseño a escala de oblea.
Cerebras también proporciona una biblioteca de implementaciones de modelos preoptimizados, incluidos transformadores, redes convolucionales y redes recurrentes. La pila de software incluye herramientas para el entrenamiento distribuido en múltiples sistemas CS, utilizando una técnica llamada transmisión de pesos, que canaliza los pesos del modelo a través de la memoria del chip. Este enfoque permite entrenar modelos más grandes que la memoria en el chip al transmitir pesos desde el almacenamiento externo.
Aplicaciones y Casos de Uso
El WSE ha sido adoptado por varias organizaciones tanto para fines comerciales como de investigación. En 2021, el grupo G42 en los Emiratos Árabes Unidos se asoció con Cerebras para construir una supercomputadora para la investigación de IA, que luego se utilizó para entrenar los modelos Jais y Falcon. En 2023, Cerebras anunció una colaboración con Qualcomm para desarrollar modelos de IA en dispositivos, aunque los detalles de esta asociación siguen siendo escasos.
En el sector de la salud, Cerebras ha trabajado con empresas farmacéuticas para acelerar las simulaciones de descubrimiento de fármacos. La capacidad del chip para manejar redes neuronales de grafos a gran escala se ha utilizado para la predicción de propiedades moleculares. Además, el WSE se ha empleado en la predicción meteorológica, con experimentos que muestran un entrenamiento más rápido de modelos climáticos de alta resolución en comparación con los clústeres de GPU.
Comparación con Otros Hardware de IA
El WSE compite con otros procesadores de IA especializados, como las unidades de procesamiento de lenguaje de Groq y las unidades de flujo de datos reconfigurables de SambaNova. A diferencia de las GPU, que son procesadores paralelos de propósito general, el WSE es una arquitectura específica de dominio optimizada para cargas de trabajo de IA. Su principal ventaja es la eliminación de la comunicación entre chips, que puede ser una sobrecarga significativa en el entrenamiento a gran escala.
Sin embargo, el tamaño y el consumo de energía del WSE (hasta 15 kilovatios para el CS-2) limitan su despliegue a centros de datos con refrigeración adecuada. En contraste, AWS Trainium y las TPU de Google se ofrecen como servicios en la nube, que pueden ser más accesibles para muchas organizaciones. El WSE también enfrenta competencia de los aceleradores Instinct de AMD y los procesadores Gaudi de Intel, que son más convencionales en diseño pero se benefician de ecosistemas de software maduros.
Desafíos y Limitaciones
Uno de los principales desafíos del WSE es el rendimiento de fabricación. Fabricar una oblea sin defectos es difícil, pero Cerebras utiliza una técnica llamada "redundancia" para tolerar defectos al deshabilitar núcleos defectuosos y redirigir la comunicación. Este enfoque permite a la empresa utilizar obleas que de otro modo se descartarían, pero también significa que el número efectivo de núcleos varía entre chips.
Otra limitación es la capacidad de memoria. Aunque la memoria en el chip es grande, sigue siendo más pequeña que la memoria de alta banda ancha disponible en las GPU. Para modelos que requieren grandes tablas de incrustación o matrices de atención, el WSE debe depender de la memoria externa, lo que reduce su ventaja de rendimiento. Además, el ecosistema de software es menos maduro que el de CUDA, y algunos investigadores han informado dificultades para portar kernels personalizados.
El costo del sistema CS-2 no se divulga públicamente, pero se estima que está en el rango de millones de dólares, lo que lo hace inaccesible para muchos grupos de investigación pequeños. Cerebras ha abordado esto ofreciendo acceso en la nube a través de sus propios centros de datos y asociaciones con Oracle Cloud y Azure.
Direcciones Futuras
Cerebras continúa iterando en la arquitectura WSE. El WSE-3, anunciado en 2024, está diseñado para admitir el entrenamiento de modelos con hasta 24 billones de parámetros, lo que requeriría un clúster de múltiples sistemas CS-3. La empresa también está explorando nuevas tecnologías de empaquetado, como chiplets, para aumentar aún más el rendimiento y reducir los costos.
En 2023, Cerebras presentó una solicitud de OPI, pero los planes se retrasaron debido a las condiciones del mercado. La empresa ha recaudado más de 700 millones de dólares en financiamiento de inversores que incluyen al cofundador de OpenAI, Sam Altman (aunque no en una capacidad oficial) y Benchmark Capital. A partir de 2024, el WSE sigue siendo un producto nicho pero influyente en el panorama del hardware de IA, empujando los límites de lo que es posible en la computación de un solo chip.
Recepción e Impacto
El WSE ha recibido reacciones mixtas de la comunidad de IA. Los defensores destacan su enfoque innovador para resolver el problema de la banda ancha de memoria, mientras que los escépticos cuestionan su practicidad y rentabilidad en comparación con los clústeres de GPU. Estudios independientes han demostrado que para ciertas cargas de trabajo, como la atención dispersa y el procesamiento de grafos, el WSE puede superar a las GPU por un margen significativo, pero para otras, la ventaja es menos clara.
A pesar de los debates, el WSE ha influido en el diseño de otros chips de IA, particularmente en el área de la memoria en el chip y las interconexiones. Su éxito también ha validado el concepto de integración a escala de oblea, que anteriormente se consideraba poco práctico. A medida que los modelos de IA generativa continúan creciendo en tamaño, es probable que aumente la demanda de hardware especializado como el WSE, aunque el panorama competitivo sigue siendo dinámico.
Véase También
Referencias
- Cerebras Systems. "Wafer-Scale Engine: An Introduction." 2019.
- Feldman, A. et al. "The Cerebras Wafer-Scale Engine." IEEE Micro, 2021.
- Cerebras Systems. "CS-3 System Overview." 2024.
- Sandia National Laboratories. "Real-Time Spiking Neural Network Simulation." 2022.
- G42. "Building the World's Largest AI Supercomputer." 2021.