Cerebras Architecture Deep Dive: First Look Inside the HW/SW Co-Design for Deep Learning [Updated]

Core Architecture

Scale-Up: Amplifying Moore’s Law

Now, let’s look at how we scale that up. Traditionally, scaling up within a chip has been the domain of the fabs. Moore’s Law has carried our industry for decades, enabling denser and denser chips. Today, Moore’s Law is still alive and well. But it’s only giving incremental gains, maybe a two times improvement per process generation, and that’s simply not enough. So, we ask ourselves, can we amplify Moore’s Law, and get an order of magnitude or more improvement?

The traditional way to amplify Moore’s Law is to make larger chips. We did that, and we took it to the extreme. The result is our second-generation Wafer-Scale Engine, the WSE-2. It’s currently generally available and being used by our customers every day. It’s the largest chip ever built, at 56 times larger than the largest CPU today. It’s over 46,000 square millimeters in size, with 2.6 trillion transistors on a single chip, and we can fit 850,000 cores. With all those cores integrated on a single piece of silicon, we get some truly mind-boggling numbers for memory and performance because everything is on-chip.

We built a specially designed system around it, called the Cerebras CS-2. This was co-designed around the WSE-2, enabling the wafer-scale chip to be used in a standard datacenter environment. It’s truly cluster-level compute in a single box.

Here’s how we build up to that massive wafer from all those small cores. First, we create a traditional die with 10,000 cores each. Instead of cutting up those die to make traditional chips, we keep them intact, but we carve out a larger square within the round 300-millimeter wafer. That’s a total of 84 die, with 850,000 cores, all on a single chip (Figure 6.). All of this is only possible if the underlying architecture can scale to that extreme size.

The fundamental enabler is the fabric. It needs to enable efficient and high-performance communication across the entire wafer (Figure 7). Our fabric does this by using a 2D mesh topology that is well-suited to scale on silicon with extremely low overheads. This fabric ties together all the cores, where each core has a fabric router within the mesh topology. The fabric routers have a simple 5-port design with 32-bit bi-directional interfaces in each of the four cardinal directions and one port facing the core itself. This small port count enables single clock cycle latency between nodes, enabling a low-cost, lossless flow control with very low buffering.

The fundamental data packet is just a single FP16 data element optimized for neural networks. Along with that FP16 data, there are 16 bits of control information making up a 32-bit ultra-fine-grained packet. To optimize the fabric further, it uses entirely static routing, which is extremely efficient and low overhead, while perfectly exploiting the static connections of neural networks. To enable multiple routes on same physical link, we have 24 independent static routes that can be configured. We call those colors. All of them are non-blocking between one another and they’re all time-multiplexed onto

the same physical links. Lastly, neural network communication inherently has a high degree of fan-out, so our fabric is assigned with native broadcast and multi-cast capabilities within each fabric router.

Now that we have a scalable foundation, we need to scale it. Scaling within a single die is simple. To scale beyond the die, we extend the fabric across those die boundaries. We cross less than a millimeter of scribe line, and we do this using high-level metal layers within the TSMC process. This extends that to the 2D mesh compute fabric to a fully homogeneous array of cores across the entire wafer. The die-to- die interface is a highly efficient, source-synchronous, parallel interface. But, at this wafer scale, it adds up to over a million wires, so we must have redundancy built directly into the underlying protocol. We do that with training and auto-correction state machines. With these interfaces, even with defects in the fab process, we get a fully uniform fabric across the entire wafer.

These seemingly simple short wires are a big deal because they span less than a millimeter of distance on silicon. When you compare that to traditional SERDES approaches, the difference is massive. Just like the memory, it’s just physics. Driving bits less than a millimeter on chip is much easier than across package connectors, PCBs, and sometimes even cables. This results in orders of magnitude improvement compared to traditional IO. As you can see from the table in Figure 8, we’re able to achieve about an order of magnitude more bandwidth per unit area, and almost two orders of magnitude better power efficiency per bit. All of this translates into unprecedented fabric performance across the entire wafer. If you normalize to GPU equivalent area, that’s seven times more bandwidth than GPU die-to-die bandwidth in the same GPU area, at only five watts of power. It’s this level of global fabric performance that enables the wafer to operate as a single chip. That’s important because with such a powerful single chip, we can solve some really, really large problems.

Weight Streaming Enables the Largest Models

The fabric enables us to run extremely large neural networks all on a single chip. The WSE-2 has more than enough performance and capacity to run even the largest models without partitioning or complex distribution. This is done by disaggregating the neural network model, memory weights, and the compute. We store all the model weights externally in a device called MemoryX, and we stream all those weights onto the CS-2 system as they’re needed to compute each layer of the network, one layer at a time. The weights are never stored on the system, not even temporarily. As the weights stream through, the CS-2 performs the computation using the underlying dataflow mechanisms in the cores (Figure 9).

Each individual weight triggers the computation as an individual AXPY operation. Once each weight is complete, it’s discarded, and the hardware moves on to the next element. Now, since the weights are never stored on the chip, the size of the model is not constrained by the memory capacity on chip. On the backward pass, the gradients are streamed out in the reverse direction back to the MemoryX unit, where the weight updates happen.

Let’s take a deeper dive into how the computation is performed, so we can see how the architecture properties uniquely enable this capability. Neural network layers boil down to matrix-multiplication. Because of the scale of the CS-2, we’re able to use all 850,000 cores of the wafer as a single, giant matrix multiplier (Error! Reference source not found.). Here’s how that works: For transformer models, like GPT, activation tensors have three logical dimensions: Batch, Sequence, and Hidden dimension (B, S and H). We split these tensor dimensions over the 2D grid of cores on the wafer. The Hidden dimension is split over the fabric in the x-direction and the Batch and Sequence dimensions are split over the fabric’s y-direction. This arrangement allows for efficient weight broadcast and reductions over Sequence and Hidden dimensions.

With activations stored on the cores where the work will be performed, the next step is to trigger the computation on those activations. This is done by using the on-chip broadcast fabric. That’s what we use to send the weights, the data, and the commands to each column. Of course, using the hardware dataflow mechanisms, the weights then trigger the FMAC operations directly. These are the AXPY operations. Since the broadcast occurs over columns, all cores containing the same subset of features receive the same weights. Additionally, we send commands to trigger other computations such as reductions or non-linear operations.

Let’s walk through an example. We start by broadcasting the row of weights across the wafer (Figure 11). Each element of the row is our scalar. Within that row, there are multiple weights that map onto a single column, of course. When there’s sparsity, only those non-zero weights are broadcasted to the column, triggering those FMAC computations. We skip all the zero weights, and we stream in the next non-zero weight. This is what produces a sparsity acceleration.

If we now zoom into a core, we can see how the core architecture is being used for this operation (Figure 12). When a weight arrives, using the dataflow mechanisms, it triggers an FMAC operation on the core. The weight value is multiplied with each of these activations and added into a local accumulator which resides in the SW-managed cache. The FMAC compute is performed using a tensor instruction with the activations as a tensor operand. All of this is done with zero additional overhead on the core. Additionally, there is no memory capacity overhead for the weights since once the compute done, the core moves onto the next weight. We never store any of the weights. Once all weights for the row have been received, each core contains a partial sum that needs to be reduced across the row of cores.

That reduction is then triggered by a command packet broadcasted to all the cores of each column. Again, using the dataflow scheduling mechanisms, once the core receives the command packet, it triggers the partial sum reduction. The actual reduction compute itself is done using the core’s tensor instruction, this time with fabric tensor operands. All the columns receive a PSUM (partial sum) command. But one column receives a special FSUM (final sum) command. The FSUM command indicates to the core that it should store the final sum. We do this so the output features are stored using the same kind of distribution as was used for input features, setting us up for the next layer. Once it receives the commands, the cores communicate using a ring pattern on the fabric which is set up using the fabric static routing colors. Using micro-threads, all this reduction is overlapped with the FMAC compute for the next weight row which started in parallel. Once all rows of weights are processed, the full GEMM operation is complete, and all our activations are in place for the next layer.

This enables neural networks of all sizes to run with high performance, all on a single chip. That’s made possible by the unique core memory and fabric architecture. This means extremely large matrices are supported without blocking or partitioning, even the largest models with up to 100,000 by 100,000 MatMul layers can run without splitting up the matrix. When you bring this together across the single WSE-2 chip, this results in 75 petaflops of FP16 sparse performance (and potentially even more with higher sparsity), or 7.5 petaflops of FP16 dense performance, all on a single chip. Now, bringing this back to the Grand ML Demand challenge in front of us, this is how we can achieve an order of magnitude level improvement in scaling up.

Scale-Out: Why Is it So Hard Today?

Let’s talk about the last component: Cluster scale-out. Clustering solutions already exist today. So why is it still so hard to scale?

Let’s look at existing scale out techniques (Figure 13). The most common is data parallel. This is the simplest approach, but it doesn’t work well for large models because the entire model needs to fit in each device to solve that. To solve that problem, the common approach is to run model parallel. This splits the model and runs different layers on different devices as a pipeline. But as the pipeline grows, the activation memory increases quadratically to keep the pipeline full. To avoid that, it’s also common to run another form of model parallel by splitting layers across devices. This has significant communication overhead, and splitting individual layers is extremely complicated. Because of all these constraints, there is no single one-size-fits-all way to scale out today. In fact, in most cases, training massive models requires a hybrid approach with both data parallel and model parallel. Although scale out solutions technically exist, they have many limitations. And the fundamental reason is simple: In traditional scale out, our memory and compute are tied to each other. Trying to run a single model on thousands of devices turns the scaling of both memory and compute into distributed constraint problems that are interdependent.

Here’s the result of that complexity: Figure 14 shows the largest models trained on GPUs over the last few years and the different types of parallelism used. As you can see, as the models get larger, the more types of parallelism are needed, and this results in a tremendous amount of complexity. For example, you can see that the level of tensor model parallelism is always limited to 8 because that’s the number of GPUs that are typically in a single server. As a result, most parallelism for large models is pipelined model parallelism, but that’s the most complex because of the memory tradeoffs. Training these models on GPU clusters today requires navigating all these bespoke distributed system problems. This complexity results in longer development times and often suboptimal scaling.

The Cerebras Architecture Makes Scaling Easy

On the other hand, because the Cerebras architecture enables running all models on a single chip without portioning, scaling becomes easy and natural. We can scale with only data parallel replication, no need for any complicated model parallel partitioning.

We do this with a specially designed interconnect for data parallel (Figure 15). This is called SwarmX. It sits between the MemoryX units that hold the weights and the CS-2 systems for compute but is independent from both. SwarmX broadcasts weights to all CS-2 systems and it reduces gradients from all CS-2s. It’s more than just an interconnect – it’s an active component in the training process, purpose- built for data parallel scale-out. Internally, SwarmX uses a tree topology to enable modular and low overhead scaling. Because it’s modular and disaggregated, you can scale to any number of CS-2 systems with the same execution model as a single system. Scaling to more compute is as simple as adding more nodes to the SwarmX topology and adding more CS-2 systems. That is how we address the last component of the Grand ML Demand Challenge; to improve and drastically simplify scale-out.

Conclusion

Let’s go back to where we started today. In the past couple of years, we have seen over three orders of magnitude greater demand from ML workloads, and there’s no sign of slowing down. In the next couple of years, this is where we will be (Figure 16). We ask ourselves, is this even possible?

At Cerebras, we know it is. But not by using traditional techniques. Only by improving the core architecture by an order of magnitude with unstructured sparsity acceleration. Only by scaling up by an order of magnitude with wafer scale chips, and only by improving cluster scale out by an order of magnitude with truly scalable clustering. With all of these, that future is achievable. Neural network models are continuing to grow exponentially. Few companies today have access to them, and that list is only getting smaller.

With the Cerebras architecture, by enabling the largest models to run on a single device, data parallel only scale out and native unstructured sparsity acceleration, we’re making these large models available to everyone.

Sean Lie, Co-Founder and Chief Hardware Architect | August 29, 2022