The Wafer-Scale Engine (WSE) is a series of large-scale processors developed by Cerebras Systems, a company founded in 2015 by Andrew Feldman, Gary Lauterbach, and others. Unlike conventional chips that are cut from a silicon wafer, the WSE uses the entire wafer as a single die, enabling an unprecedented number of cores and on-chip memory. This design aims to reduce the communication bottlenecks that plague multi-chip systems, making it particularly suited for training large language models and other deep learning workloads.
The first generation, WSE-1, was announced in August 2019. It contained 1.2 trillion transistors and 400,000 AI-optimized cores, fabricated using a 16-nanometer process by TSMC. The chip measured 46,225 square millimeters, over 50 times larger than the largest graphics processing units of the time. The second generation, WSE-2, followed in April 2021, using a 7-nanometer process, doubling the core count to 850,000 and increasing on-chip memory to 40 gigabytes. In March 2024, Cerebras introduced the WSE-3, built on a 5-nanometer process, with 900,000 cores and 44 gigabytes of on-chip SRAM, targeting training of models with up to 24 trillion parameters.
Architecture and Design
The WSE's core innovation is its wafer-scale integration. Instead of manufacturing many small dies on a wafer and packaging them separately, Cerebras keeps the entire wafer intact, with redundant circuits to route around manufacturing defects. This approach eliminates the need for inter-chip communication over a printed circuit board, which is often a performance bottleneck in distributed systems. The chip uses a mesh-connected network of cores, each with its own memory, allowing for high-bandwidth, low-latency data movement.
Each core is a simplified processor optimized for the sparse linear algebra operations common in neural networks. The cores support a range of data formats, including FP16, BF16, and FP32, and are designed to execute both dense and sparse operations efficiently. The on-chip memory is distributed across the cores, providing a total bandwidth that far exceeds that of traditional GPU memory systems. This architecture reduces the reliance on external memory, which is often a limiting factor in AI training.
The WSE is not a standalone product; it is integrated into Cerebras's CS-2 and CS-3 systems, which include a power delivery system and cooling solution. The CS-2, released in 2020, houses the WSE-2, while the CS-3, announced in 2024, houses the WSE-3. These systems are designed to be plugged into existing data centers and are often used in clusters for large-scale training.
Performance and Benchmarks
Cerebras has published several performance benchmarks for the WSE. In 2021, the company demonstrated that a single CS-2 system could train a 1.8-billion-parameter model using the GPT architecture, achieving near-linear scaling when multiple systems were connected. In 2022, Cerebras and Sandia National Laboratories used a CS-2 to simulate a 200-million-neuron spiking neural network in real time, a task that would require a supercomputer with thousands of GPUs.
In 2023, Cerebras reported that its systems could train a 7-billion-parameter model in less time than a cluster of 16 NVIDIA A100 GPUs, due to the reduced communication overhead. The WSE-3, according to the company, can deliver up to 125 petaflops of AI compute, making it one of the most powerful single-chip processors for AI. However, independent benchmarks are limited, and most performance claims come from Cerebras itself.
Software Ecosystem
To program the WSE, Cerebras developed the Cerebras Software Platform (CSoft), which includes a compiler that maps machine learning models to the chip's architecture. CSoft supports popular frameworks like TensorFlow and PyTorch, allowing researchers to run existing models without significant code changes. The compiler handles dataflow scheduling, memory allocation, and optimization for the wafer-scale design.
Cerebras also provides a library of pre-optimized model implementations, including transformers, convolutional networks, and recurrent networks. The software stack includes tools for distributed training across multiple CS systems, using a technique called weight streaming, which pipelines model weights through the chip's memory. This approach allows for training models larger than the on-chip memory by streaming weights from external storage.
Applications and Use Cases
The WSE has been adopted by various organizations for both commercial and research purposes. In 2021, the G42 group in the United Arab Emirates partnered with Cerebras to build a supercomputer for AI research, which was later used to train the Jais and Falcon models. In 2023, Cerebras announced a collaboration with Qualcomm to develop on-device AI models, though the details of this partnership remain sparse.
In the healthcare sector, Cerebras has worked with pharmaceutical companies to accelerate drug discovery simulations. The chip's ability to handle large-scale graph neural networks has been used for molecular property prediction. Additionally, the WSE has been employed in weather forecasting, with experiments showing faster training of high-resolution climate models compared to GPU clusters.
Comparison with Other AI Hardware
The WSE competes with other specialized AI processors, such as Groq's language processing units and SambaNova's reconfigurable dataflow units. Unlike GPUs, which are general-purpose parallel processors, the WSE is a domain-specific architecture optimized for AI workloads. Its main advantage is the elimination of inter-chip communication, which can be a significant overhead in large-scale training.
However, the WSE's size and power consumption (up to 15 kilowatts for the CS-2) limit its deployment to data centers with adequate cooling. In contrast, AWS Trainium and Google's TPU are offered as cloud services, which may be more accessible for many organizations. The WSE also faces competition from AMD's Instinct accelerators and Intel's Gaudi processors, which are more conventional in design but benefit from mature software ecosystems.
Challenges and Limitations
One of the primary challenges of the WSE is yield. Manufacturing a defect-free wafer is difficult, but Cerebras uses a technique called "redundancy" to tolerate defects by disabling faulty cores and rerouting communication. This approach allows the company to use wafers that would otherwise be discarded, but it also means that the effective core count varies between chips.
Another limitation is the memory capacity. While the on-chip memory is large, it is still smaller than the high-bandwidth memory available on GPUs. For models that require large embedding tables or attention matrices, the WSE must rely on external memory, which reduces its performance advantage. Additionally, the software ecosystem is less mature than that of CUDA, and some researchers have reported difficulties in porting custom kernels.
The cost of the CS-2 system is not publicly disclosed, but it is estimated to be in the millions of dollars, making it inaccessible for many small research groups. Cerebras has addressed this by offering cloud access through its own data centers and partnerships with Oracle Cloud and Azure.
Future Directions
Cerebras continues to iterate on the WSE architecture. The WSE-3, announced in 2024, is designed to support training of models with up to 24 trillion parameters, which would require a cluster of multiple CS-3 systems. The company is also exploring new packaging technologies, such as chiplets, to further increase performance and reduce costs.
In 2023, Cerebras filed for an IPO, but the plans were delayed due to market conditions. The company has raised over $700 million in funding from investors including OpenAI co-founder Sam Altman (though not in an official capacity) and Benchmark Capital. As of 2024, the WSE remains a niche but influential product in the AI hardware landscape, pushing the boundaries of what is possible in single-chip computing.
Reception and Impact
The WSE has received mixed reactions from the AI community. Proponents highlight its innovative approach to solving the memory bandwidth problem, while skeptics question its practicality and cost-effectiveness compared to GPU clusters. Independent studies have shown that for certain workloads, such as sparse attention and graph processing, the WSE can outperform GPUs by a significant margin, but for others, the advantage is less clear.
Despite the debates, the WSE has influenced the design of other AI chips, particularly in the area of on-chip memory and interconnects. Its success has also validated the concept of wafer-scale integration, which was previously considered impractical. As generative AI models continue to grow in size, the demand for specialized hardware like the WSE is likely to increase, though the competitive landscape remains dynamic.
See Also
References
- 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.