Cerebras Deploys Trillion-Parameter AI Models
Cerebras explains how it deploys trillion-parameter models across wafer-scale chips that each hold only 44 GB of on-chip SRAM.
Extract
One of the most frequently asked questions we get is how we deploy large models on Cerebras wafer-scale chips. A single Cerebras WSE-3T wafer has 44 GB of on-chip memory, and trillion-parameter frontier models need terabytes for their weights alone.In this post, we’ll explain how we split a large frontier model across multiple WSE-3T wafers, to run multi-trillion parameter models at 1000-2000+ tokens per second.We’ll use the open source model Kimi K3, with ~3 trillion parameters, as an illustrative example.By the end, you'll understand:How we fit the full model into on-chip memoryHow we map the model across multiple chipsHow we handle a high # of concurrent requestsBreaking it down with Kimi K3Kimi K3 is an open source, 3T parameter model with 69 linear-attention layers and 24 dense layers.For this illustrative example, we use:A 256k assumed average request size4-bit weight quantization (like the Hugging Face implementation)An 8-bit KV cache#1: Size of Each layerTo calculate the total number of Cerebras WSE-3T wafers we need to support Kimi K3, let’s first calculate the size of each model layer.Kimi K3 has 3 trillion parameters. If all 3T parameters are stored at 4 bits each, the raw parameter payload is:Dividing that payload by 93 layers then gives:#2: Number of WafersNow that we know each layer needs about 16 GB, we can figure out how many layers fit on each 44 GB wafer. The answer depends on the layer type:Dense layers need additional room for a KV cache, because each dense (full-attention) layer stores the keys and values for every token in every active request. So, we place one dense layer per wafer and leave the remaining SRAM for its KV cache.Linear-attention layers do not require additional room for a KV cache. Instead of storing keys and values for every token, a linear-attention layer keeps a small, fixed-size state per request, regardless of context length. That leaves room for two linear layers per wafer (2 × 16 GB = 32 GB).Therefore, with Kimi K3’s 69 linear layers and 24 dense layers, we get:This gives us a total wafer count of 35 + 24 = 59 wafers.#3: Connecting the wafers togetherCerebras executes models in a pipeline. After the Cerebras compiler maps the model layers across regions of wafers, the runtime maps the path for tokens to flow between wafers. For ultrafast latency, only activations need to be passed between the wafers, and this can be done via RDMA over Ethernet. This is roughly 2 microseconds per hop.It’s important to emphasize …