共计 8451 个字符，预计需要花费 22 分钟才能阅读完成。
作者|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 级别的极致性能
图 4：稠密 GEMM 即对每个非零权重执行一次 AXPY 操作。
有了极大的内存带宽，就能够实现许多卓越的性能。比方，能够充分发挥所有 BLAS 级别（根底线性代数程序集，BLAS levels）的极致性能。传统的 CPU 和 GPU 架构的片上内存带宽无限，因而只能实现 GEMM（通用矩阵乘法）的极致性能，即矩阵 - 矩阵相乘。
从图 4 可见，在低于矩阵 - 矩阵相乘的任何 BLAS 级别都须要比内存带宽的大幅减少，这一点传统架构无奈满足。
但有了足够的内存带宽后，就能够让 GEMV（矩阵 - 向量相乘）、DOT（向量 - 向量相乘）和 AXPY（向量 - 标量相乘）均实现极致性能。高内存带宽在神经网络计算中尤为重要，因为这能够实现非结构化稠密的充沛减速。一个稠密 GEMM 操作可看作是多个 AXPY 操作的合集（对每个非零元素执行一次操作）。
Cerebras 计算外围的根底是一个齐全可编程的处理器，以适应一直变动的深度学习需要。与通用处理器一样，Cerebras 外围处理器反对算术、逻辑、加载 / 贮存、比拟（compare）、分支等多种指令。这些指令和数据一样贮存在每个外围的 48 KB 本地内存中，这意味着外围之间互相独立，也意味着整个芯片能够进行细粒度动静计算。通用指令在 16 个通用寄存器上运行，其运行在紧凑的 6 级流水线中。
等式 1，FMAC 指令示例
除此之外，Cerebras 外围还在硬件层面反对所有无关数据处理的张量指令。这些张量算子在 64- 位数据通路中执行，数据通路由 4 个 FP16 FMAC（交融乘积累加运算）单元组成。
为了晋升性能与灵活性，Cerebras 的指令集架构（ISA）将张量视为与通用寄存器和内存一样的一等操作数（first-class operand）。上图等式 1 是一个 FMAC 指令的例子，它将 3D 和 2D 张量视为操作数间接运行。
之所以能够做到这一点，是因为 Cerebras 外围应用数据结构寄存器（DSR）作为指令的操作数。Cerebras 外围有 44 个 DSR，每个 DSR 蕴含一个描述符，外面有指针指向张量及其长度、形态、大小等信息。
有了 DSR 后，Cerebras 外围的硬件架构更灵便，即能够在内存中反对内存中的 4D 张量，也可反对织构张量（fabric streaming tensors）、FIFO（先进先出算法）和环形缓冲器。此外，Cerebras 外围还配有硬件状态机来治理整个张量在数据通路中的流动秩序。
3
细粒度数据流调度
图 5：外围数据通路及外围数据流调度。细粒度动静计算外围可晋升计算性能，稠密利用率为 GPU 的 10 倍。
除了改良张量利用，Cerebras 外围还可执行细粒度数据流调度。如图 5 所示，所有计算都由数据触发。Fabric 间接在硬件中传输数据和关联控件，一旦外围接收数据，就开始查找运行指令，查找工作齐全基于接管到的数据。这一数据流机制使整个计算构造变成一个数据流引擎，能够反对稠密减速——因为它只解决非零数据。发送器会过滤所有零值，因而接收器只会接管到非零值，而所有计算都由非零数据触发。
这样做岂但能够节俭功率，还能够省略不必要的计算，放慢运算效率。操作由单个数据元素触发，使得 Cerebras 外围能够反对超细粒度、齐全非结构化的稠密性，同时不会造成性能损失。因为数据流具备动态性，所以 Cerebras 外围还反对 8 个张量操作同时运行，咱们称之为“微线程（micro-threads）”。
微线程之间互相独立，每次循环时硬件可在其间切换。调度器继续为所有待处理张量监控输出和输入是否可用，还具备优先解决机制，保障要害工作失去优先解决。当不同工作间的切换产生大量动静行为时，微线程能够晋升利用率，否则这些动静行为可能会导致流水线呈现气泡。
上述细粒度、动静、小型外围架构等特点使咱们的架构具备前所未有的高性能，其非结构化稠密计算的利用率是 GPU 的至多 10 倍。可见，通过对计算外围架构的改良，Cerebras 可将性能进行数量级晋升。
4
纵向扩大：超过摩尔定律
要纵向扩大芯片，传统的办法都是从芯片制作方面动手，即晋升芯片集成度。过来数十年，芯片行业的倒退都合乎摩尔定律，芯片集成度越来越高。现在，摩尔定律还在连续，但它的增量不够大，每一代制程只能将集成度晋升约两倍，不足以满足神经网络的计算需要。所以，Cerebras 心愿能够超过摩尔定律，实现数量级的性能晋升。
为此，咱们尝试过传统的办法——扩充芯片面积，并在这方面做到了极致，成绩就是 WSE-2（Wafer-Scale Engine，晶圆级引擎）。现在，WSE- 2 的利用已十分宽泛。它是全世界最大的芯片，尺寸超过 46,000 平方毫米，是目前最大的 CPU 的 56 倍。单块 WSE- 2 有 2.6 万亿个晶体管，外围数达 850,000 个。宏大的芯片面积能够实现极大的片上内存和极高的性能。
为了让尺寸惊人的 WSE- 2 也能在规范的数据中心环境中应用，咱们还针对性地设计了 Cerebras CS- 2 零碎，做到了用单块芯片实现集群级计算。
图 6：从小型外围到大型晶圆级引擎
以下是咱们从小型外围构建大型晶圆级引擎的过程：首先，咱们在整片直径约 300 毫米的晶圆上做出一个个传统晶粒（Die），每个晶粒含有约 10,000 个外围；而后，不同于以往的是，咱们不将单个晶粒切割进去做成传统芯片，而是在整片晶圆内切割出一个边长 215 毫米的方块，方块蕴含 84 个晶粒，共有 850,000 个计算外围（图 6）。
图 7：高带宽、低提早的芯片构造
实现这样的超大芯片尺寸，离不开底层架构的配合，底层架构必须能使数据在整片晶圆上高效、高性能地传输（图 7）。Cerebras 的芯片构造应用 2D 网格拓扑，这种构造非常适合扩大，而且只需耗费极低的开销。
网格拓扑将所有外围连接起来，每个外围在网状拓扑中有一个构造路由器（fabric router）。构造路由器有 5 个端口，4 个方向各有 1 个，还有一个端口面向外围本身，各个端口都有 32 位的双向接口。端口数量较少的益处是能够将节点间延时放弃在一个时钟周期以内，从而实现低成本、无损流控和非常低的缓冲。
芯片中的根本数据包是针对神经网络优化后的单个 FP16 数据元素，与之随同的是 16 位的管制信息，它们独特组成 32 位的超细粒度数据包。
为了进一步优化芯片构造，咱们应用了动态路由（static routing），效率高，开销低，而且能够充分利用神经网络的动态连贯。为了让同一物理连贯上能够有多条路由，咱们提供 24 条互相独立的动态路由以供配置，路由之间无阻塞，且都能够通过时分复用（time-multiplexing）技术在同一物理连贯上传输。
最初，因为神经网络传输须要高扇出（fan-out），因而 Cerebras 芯片的每个构造路由器都具备本地播送（native broadcast）和多播（multi-cast）能力。
有了上述根底后，咱们就能够进行扩大。在单个晶粒内进行扩大比较简单，但当初须要将晶粒与晶粒连接起来。为了逾越晶粒间不到一毫米宽的划片槽（scribe line），咱们应用了台积电工艺中的高级金属层。
咱们将计算外围扩大为 2D 网格计算构造，而后又在整个晶圆上造成了齐全同质的计算外围阵列。晶粒 - 晶粒接口是一种高效的源同步并行接口，然而，在如此大的晶圆规模上，总共有超过一百万条线路，所以咱们的底层协定必须采纳冗余度设计。咱们通过训练和主动校对状态机来做到这一点。有了这些接口，即便在制作过程中存在瑕疵，整个晶圆的构造也能做到齐全均质构造（uniform fabric）。
图 8：整个晶圆上的均质构造（uniform fabric）。
芯片上看似简略的短线其实非常重要，它们在硅上的间隔不到一毫米。这种线路设计与传统的 SERDES 办法很不一样。与后面提到的的内存设计雷同，短线设计是出于简略的物理原理：在芯片上将比特数据传输不到 1 毫米的间隔，比通过封装连接器、PCB 或者线缆传输都更容易。
与传统 IO 相比，这种办法带来了数量级的改良。从图 8 数据可看出，WSE- 2 每单位面积的带宽比 GPU 多出约一个数量级，并且每比特的功率效率进步了近两个数量级。这些都表明整个晶圆构造具备了前所未有的高性能。
如果转化为等同的 GPU 面积，WSE- 2 的带宽是 GPU 的 7 倍，而功率仅约 5 瓦。正是这种级别的全局构造性能，使晶圆可能作为单个芯片运行。有了如此弱小的单芯片，咱们就能够解决一些极具挑战性的问题。
5
通过权重流式技术支持超大模型
图 9：通过权重流式（Weight Streaming）技术可在单个芯片上反对所有模型大小。
高性能的芯片构造能够让咱们在单个芯片上运行大型神经网络。WSE- 2 具备足够高的性能和容量来运行现在最大的模型，且无需分区或简单的分布式解决，这是通过合成神经网络模型、权重和计算来实现的。
咱们将所有模型权重存储在名为 MemoryX 的外部设备中，并将这些权重流式传输到 CS- 2 零碎。权重会在神经网络各层的计算中用到，而且一次只计算一层。权重不会存储在 CS- 2 零碎上，哪怕是临时贮存。CS- 2 接管到权重后，应用外围中的底层数据流机制执行计算（图 9）。
每个独自的权重都会作为独自的 AXPY 操作触发计算。实现计算后，该权重就会被抛弃，硬件将持续解决下一个元素。因为芯片不须要贮存权重，所以芯片的内存容量不会影响芯片可解决的模型大小。在反向流传中，梯度以相同的方向流回到 MemoryX 单元，而后 MemoryX 单元进行权重更新。
图 10 残缺的晶圆是 MatMul 阵列，可反对超大矩阵。
以下是芯片中执行计算的具体方法。神经网络各层的计算可归结为矩阵乘法，因为 CS- 2 的规模较大，咱们可能将晶圆的 85 万个外围用作单个巨型矩阵乘法器。
它是这样工作的：对于像 GPT 这样的 Transformer 模型，激活张量具备三个逻辑维度：批次（B）、序列 (S) 和暗藏（H）维度，咱们将这些张量维度拆分到晶圆上的二维外围网格上。暗藏维度在芯片构造的 x 方向上划分（split），而批次和序列维度在 y 方向上划分。这样能够实现高效的权重播送以及序列和暗藏维度的高效归约。
激活函数存储在负责执行计算工作的外围上，下一步是触发这些激活函数的计算，这是通过应用片上播送构造来实现的。咱们应用片上播送构造来向每一列发送权重、数据和命令的办法。
当然，在硬件数据流机制下，权重会间接触发 FMAC 操作。这些是 AXPY 操作。因为播送产生在列上，因而蕴含雷同特色子集的所有外围接管雷同的权重。此外，咱们发送命令来触发其余计算，例如归约或非线性操作。
图 11：数据流调度以低开销实现齐全非结构化的稠密 MatMul 运算。
举个例子，咱们首先在整个晶圆上播送权重行（图 11）。每行的每个元素都是标量，当然，在单行中，有多个权重映射到单个列上，当存在稠密性时，只有非零权重才会被播送到列，触发 FMAC 计算。咱们跳过所有的零权重，并输出下一个非零权重，这就是产生稠密减速的起因。
图 12：稠密输出的 GEMM：乘法和 partial sum 归约。
如果咱们当初放大一个外围，能够看到外围架构是如何进行此操作（图 12）。在数据流机制下，权重到达后，就会触发外围上的 FMAC 计算。权重值与每个激活函数输入相乘，而后累加到软件治理缓存中的本地累加器中。FMAC 计算应用张量指令执行，将激活函数输入视为张量操作数。上述计算都不会对外围造成额定开销。
此外，权重也不会产生内存开销，因为一旦计算实现，外围就会转而计算下一个权重，不须要存储任何权重。若整行外围都接管到权重，每个外围就都会产生一个 partial sum，而后该行外围的所有 partial sum 将进行归约。
归约计算由被播送到每列所有外围的命令包触发。同样，在数据流调度机制下，一旦外围接管到命令包，它就会触发 partial sum 归约计算。理论的归约计算自身是应用外围的张量指令实现，应用的是构造张量操作数。所有列都接管一个 PSUM 命令。然而其中一列会收到一个非凡的 FSUM 命令，它要求内核存储 final sum。这样做是为了应用与输出特色雷同的散布来存储输入特色，从而为下一层计算做好筹备。
收到命令后，外围应用构造上的环形模式进行通信，该模式应用构造动态路由设置。应用微线程，所有归约都与下一个权重行的 FMAC 计算重叠，该种 FMAC 计算并行开始。当所有的权重行都处理完毕，残缺的 GEMM 操作就实现了，同时所有的激活函数输入都已齐备，能够进行下一层计算。
上述设计能让各种规模的神经网络都能够在单个芯片上高性能运行。独特的外围内存和芯片架构使芯片能够无需分块或分区即可反对超大矩阵，即便是具备多达 100,000 x 100,000 MatMul 层的超大模型也能够在不拆分矩阵的状况下运行。
若应用单个 WSE- 2 芯片运行此模型，FP16 稠密性能可达 75 PetaFLOPS（若稠密性更高，性能还可更高），FP16 密集性能可达 7.5 PetaFLOPS。这就是咱们应答机器学习硬件挑战的第二个方面，通过扩大进一步带来一个数量级的性能晋升。
6
横向扩大：为什么这么难
最初一个方面：集群横向扩大。现在曾经存在集群解决方案，但为什么横向扩大依然如此艰难？
图 13：散布复杂性随集群规模显著减少。
让咱们看看现有的横向扩大技术（图 13）。最常见的是数据并行，这也是最简略的办法，但它不适用于大型模型，因为它要求每个设施都有足够的容量包容整个模型。
为了解决这个问题，常见的办法是采纳模型并行，即划分模型，以流水线形式用不同的设施运行模型的不同层。但随着流水线变长，激活值内存（activation memory）以二次方的速度增长。
为了防止这种状况，另一种常见的模型并行办法是跨设施划分层，但这会造成很大的通信开销，而且划分单个层非常复杂。
因为上述种种限度，明天仍没有一种万能的形式来实现横向扩大。在大多数状况下，训练海量模型须要数据并行和模型并行混合的办法。现存的横向扩大解决方案仍有许多有余，根本原因很简略：在传统的横向扩大中，内存和计算是紧密联系的，如果在数千台设施上运行单个模型，扩大内存和计算就变成相互依赖的分布式束缚问题。
图 14：GPU 集群在实践中的复杂性。
这种复杂性导致的后果是：图 14 显示了过来几年在 GPU 上训练的最大模型及其应用的不同并行办法。从中可见，越大的模型须要的并行类型也越多，减少了复杂性。
例如，张量模型的并行级别始终限度为 8，因为在单个服务器中通常只有 8 个 GPU。因而，大型模型大多采纳流水式模型并行，这是最简单的办法，起因就是之前提到的内存问题。在 GPU 集群上训练模型须要解决这些分布式系统问题。这种复杂性导致须要更长的开发工夫，并且往往无奈实现最佳扩大。
7
Cerebras 架构使扩大变得容易
Cerebras 架构可能在单个芯片上运行所有模型，无需模型宰割，因而扩大变得简略而天然，能够仅通过数据并行进行扩大，不须要任何简单的模型并行宰割。
图 15：应用 MemoryX 和 SwarmX 进行扩大，只需近线性的数据并行。
咱们为数据并行专门设计了 SwarmX（图 15）互联技术。它位于贮存权重的 MemoryX 单元和用于计算的 CS- 2 零碎之间，但又独立于两者。
SwarmX 向所有 CS- 2 零碎播送权重，并缩小所有 CS- 2 的梯度，它不仅仅是一个互联，更是训练过程中的一个沉闷组件，专为数据并行横向扩大而构建。
在外部，SwarmX 应用树形拓扑来实现模块化和低开销扩大，因为它是模块化和可分解的，所以扩大到任意数量的具备与单个零碎雷同的执行模型的 CS- 2 零碎。要扩大到更多计算，只需在 SwarmX 拓扑中增加更多节点和更多 CS- 2 零碎。这就是咱们应答机器学习硬件需要的最初一个方面：改良并大大简化横向扩大。
8
总结
在过来的几年里，机器学习工作负载的需要减少了三个数量级以上，而且没有放缓的迹象。预计几年后将增长到图 16 的箭头地位，咱们问本人，能够满足这种需要吗？
图 16：各种最先进的神经网络的内存和计算要求。横、纵坐标每一格代表一个数量级的晋升。
Cerebras 置信，咱们能够，但不是通过传统技术做到这一点，而是通过非结构化稠密减速、晶圆级芯片和集群横向扩大的联合将性能晋升三个数量级。神经网络模型规模仍然呈指数级增长，能够应用这些大模型的公司很少，而且将来只会更少。
然而，Cerebras 架构反对用单个设施运行超大模型，同时反对只需数据并行的横向扩大，以及本地非结构化稠密减速，将让更多人都能应用大模型。
（原文：https://www.cerebras.net/blog…；视频：https://www.youtube.com/watch…）
欢送下载体验 OneFlow v0.8.0 最新版本：https://github.com/Oneflow-In…