Benchmarking the Performance of Large Language Models on the Cerebras Wafer Scale Engine
Zuoning Zhang, Dhruv Piyush Parikh, Youning Zhang, Viktor K. Prasanna · 2024
Transformer based Large Language Models (LLMs) have recently reached state of the art performance in Natural Language Processing (NLP) and Computer Vision (CV) domains. LLMs use the Multi-Headed Self-Attention (MHSA) mechanism to capture long-range global attention relationships among input words or image patches, drastically improving its performance over prior deep learning approaches. In this paper, we evaluate the performance of LLMs on the Cerebras Wafer Scale Engine (WSE). Cerebras WSE is a high performance computing system with 2.6 trillion transistors, 850,000 cores and 40 GB on-chip memory. Cerebras WSE's Sparse Linear Algebra Compute (SLAC) cores eliminate multiply-by-zero operations and its 40 GB of on-chip memory is uniformly distributed among SLAC cores, enabling fast local access to model parameters. Moreover, Cerebras software configures routing between cores at runtime, optimizing communication overhead among cores. As LLMs are becoming more ubiquitous used, new hardware architectures are required to accelerate LLM training and inference. We benchmark the effectiveness of this hardware architecture at accelerating LLM training and inference. Additionally, we analyze if Cerebras WSE can scale the memory-wall associated with traditionally memory-bound compute tasks using its 20 PB/s high bandwidth memory. Furthermore, we examine the performance scalability of Cerebras WSE through a roofline model. By plotting training throughput against computational intensity, we aim to assess their effectiveness at handling high compute-intensive LLM training and inference tasks.