深挖Cerebras：世界上最大AI芯片的架构设计

作者｜Sean Lie

翻译｜胡燕君、程浩源

近年来，神经网络模型规模呈指数级增长，从2018年拥有超1亿参数的Bert到2020年拥有1750亿个参数GPT-3，短短两年模型的参数量增加了3个数量级，而且这种增长还看不到尽头。 

人们刚刚开始发掘神经网络的应用潜力，但传统的训练和推理方式已然无法跟上神经网络规模的飞速增长速度，无法满足大规模机器学习所需的内存和算力需求。为此，国内外诸多创业公司寻求对软硬件等进行实质性的底层技术革新来解决这一挑战。

作为业内备受关注的AI加速器创业公司，成立于2016年的Cerebras希望通过构建全新AI加速器方案解决AI计算问题，以实现数量级计算性能：首先，需要改进计算核心架构，而不只是一味地提升每秒浮点运算次数；其次，需要以超越摩尔定律的速度提高芯片集成度；最后，还要简化集群连接，大幅度提升集群计算效率。

为了实现上述目标，Cerebras设计了一种新的计算核心架构。它让单台设备运行超大规模模型成为可能，此外，它开发出只需简单数据并行的横向扩展和本地非结构化稀疏加速技术，使大模型的应用门槛大幅降低。

图1：近年来各SOTA神经网络模型的内存与算力需求

2021年，Cerebras曾推出全球最大AI芯片Wafer Scale Engine 2（WSE-2），面积是46225平方毫米，采用7nm工艺，拥有2.6万亿个晶体管和85万个AI优化核，还推出了世界上第一个人类大脑规模的AI解决方案CS-2 AI计算机，可支持超过120万亿参数规模的训练。今年6月，它又在基于单个WSE-2芯片的CS-2系统上训练了世界上最大的拥有200亿参数的NLP模型，显著降低了原本需要数千个GPU训练的成本。 

在近期举办的Hot Chips大会上，Cerebras联合创始人&首席硬件架构师Sean Lie深入介绍了Cerebras硬件，展示了他们在核心架构、纵向扩展和横向扩展方面的创新方法。以下是他的演讲内容，由OneFlow社区编译。 

1

Cerebras计算核心架构

计算核心（compute core）是所有计算机架构的“心脏”，而Cerebras针对神经网络的细粒度动态稀疏性重新设计了计算核心。

图2：Cerebras计算核心

图2是一款小型核心，它只有38,000平方微米，其中一半的硅面积用于48 KB内存，另一半是含110,000个标准单元（cell）的计算逻辑。整个计算核心以1.1 GHz的时钟频率高效运行，而峰值功率只有30毫瓦。

先从内存说起。GPU等传统架构使用共享中央DRAM，但DRAM存取速度较慢，位置也较远。即便使用中介层（interposer）和HBM等尖端技术，其内存带宽也远低于核心数据通路带宽。例如，数据通路带宽通常是内存带宽的100倍。

这意味着每一个来自内存的操作数（operand）至少要在数据通路中被使用100次，才能实现高利用率。要做到这一点，传统的方法是通过本地缓存和本地寄存器实现数据复用。

然而，有一种方法可以让数据通路以极致性能利用内存带宽，就是将内存完全分布在要使用内存的单元旁边。这样一来，内存带宽就等于核心数据通路的操作数带宽。

这是一个简单的物理原理：将比特数据从本地内存移动到数据通路，中间只有几十微米的距离，相比将它通过数据包移动到外部设备要容易得多。 

图3：Cerebras计算核心的内存设计：每个核心配有独立内存。

图3展示了Cerebras计算核心的内存设计，每个核心配有48 KB本地SRAM，8个32位宽的单端口bank使其具备高密度，同时可保证充分发挥极致性能，这种级别的bank可提供超出数据通路所需的内存带宽。

因此，我们可以从内存中提供极致数据通路性能，也就是每个循环只需2个64位读取，一个64位写入，因此它可以保证数据通路充分发挥性能。值得注意的是，每个核心的内存相互独立，没有传统意义上的共享内存。

除了高性能的SRAM以外，Cerebras计算核心还具备一个256字节的软件管理缓存，供频繁访问的数据结构使用，如累加器等。该缓存离数据通路非常紧凑，所以消耗的功率极低。上述分布式内存架构造就了惊人的内存带宽，相当于同等面积GPU内存带宽的200倍。

2

所有BLAS级别的极致性能