Cerebras 介绍在 WSE-3T 上部署万亿参数 AI 模型的方法

我们收到的最常见问题之一，是如何在 Cerebras 晶圆级芯片上部署大型模型。单块 Cerebras WSE-3T 晶圆拥有 44 GB 片上内存，而万亿参数级前沿模型仅权重就需要 terabytes 的存储空间。

在本文中，我们将解释如何将一个大型前沿模型拆分到多个 WSE-3T 晶圆上，从而以 1000-2000+ tokens per second 的速度运行数万亿参数模型。

我们将使用开源模型 Kimi K3 作为示例，该模型约有 ~3 万亿参数。

读完本文后，你将了解：

• 我们如何将完整模型装入片上内存
• 我们如何将模型映射到多块芯片
• 我们如何处理大量并发请求

使用 Kimi K3 拆解分析

Kimi K3 是一个开源的 3T 参数模型，包含 69 个线性注意力层和 24 个稠密层。

在这个示例中，我们使用：

• 平均请求大小假定为 256k
• 4-bit 权重量化（类似于 Hugging Face 实现）
• 8-bit KV 缓存

#1：每一层的大小

为计算支持 Kimi K3 所需的 Cerebras WSE-3T 晶圆总数，我们先计算每个模型层的大小。

Kimi K3 有 3 万亿参数。如果全部 3T 参数都按每个 4 bits 存储，原始参数数据量为：

将该数据量除以 93 层，得到：

#2：晶圆数量

现在我们知道每层约需 16 GB，就可以确定每块 44 GB 晶圆能容纳多少层。答案取决于层的类型：

• 稠密层需要为 KV 缓存留出额外空间，因为每个稠密（全注意力）层都会为每个活跃请求中的每个 token 存储键和值。因此，我们每块晶圆放置一个稠密层，并将剩余的 SRAM 用于其 KV 缓存。
• 线性注意力层不需要为 KV 缓存预留额外空间。线性注意力层不会为每个 token 存储键和值，而是为每个请求保留一个小型固定大小的状态，无论上下文长度如何。因此，每块晶圆可以容纳两个线性层（2 × 16 GB = 32 GB）。

因此，对于 Kimi K3 的 69 个线性层和 24 个稠密层，我们得到：

这样一来，晶圆总数为 35 + 24 = 59 块晶圆。

#3：连接各块晶圆

Cerebras 以流水线方式执行模型。Cerebras 编译器将模型层映射到晶圆上的各个区域后，运行时再映射 token 在晶圆之间流动的路径。为了实现超快延迟，晶圆之间只需传递激活值，并且可以通过以太网上的 RDMA 完成。每跳大约为 2 microseconds。

需要强调的是，由于每块 Cerebras 晶圆都有 44 GB 片上 SRAM，而每层仅占 16 GB，因此我们通常能够将模型的完整层直接放在芯片上。

这意味着在大多数情况下，模型权重无需移动，各层也不必拆分到多块晶圆上。只有激活值会从一块晶圆传递到下一块。

#4：我们如何处理 1 个请求

现在，各块晶圆已经连接起来，模型布局也已确定，我们可以把每块晶圆看作装配线上的一个阶段：每个阶段的模型权重保持不动，请求的中间结果则依次经过各层阶段。

例如，晶圆 1 运行模型第 1–2 层，并将其激活值发送到晶圆 2。晶圆 2 运行第 3–4 层，并将结果传递给晶圆 3，依此类推，直到最后一块晶圆生成下一个 token。

#5：我们如何处理大量请求

这引出了一个常见误解。如果一个 token 必须依次通过每一片晶圆，那么在等待轮到自己时，每片晶圆不都会闲置吗？实际上，这样一来，整个系统不就像是在以批大小为 1 运行吗？

不！如果我们每次只沿着流水线发送一个请求，并等待它完成后才开始下一个请求，就会浪费流水线的效率。相反，我们会让不同请求始终填满整条流水线。由于模型已经拆分为多个流水线阶段，因此每个阶段都可以同时处理不同的请求。

当阶段 3 正在为请求 A 处理 token 时，阶段 2 正在为请求 B 处理 token，而阶段 1 正在为请求 C 处理 token。在我们的示例中，每个阶段是一层。在生产环境中，Cerebras 使用粒度细得多的阶段，而且阶段数量也更多。

正在加载播放器...

每个阶段还可以处理一个微批，即一小组请求，而不必只处理单个请求。因此，同时处理中的请求数量大致等于阶段数乘以微批大小。

正在加载播放器...

要理解这种权衡，较小的微批会让每个请求处理得更快，而较大的微批则能提高总吞吐量。

#6: 我们能处理多少个请求？

分配给 KV 缓存的片上内存量决定了并发请求的最大数量。每个同时处理中的请求都需要在每片稠密层晶圆上拥有自己的 KV 缓存，因此片上内存大小决定了并发请求数量的上限。如上文所述，我们只需要计算稠密层的 KV 缓存需求。

对于每片已经包含一个稠密层的 44 GB 晶圆，我们可以按以下方式计算 KV 缓存的剩余空间：

晶圆总存储空间 - 一个稠密层的大小。在本例中：

计算并发请求数量

要找出任何特定模型映射所支持的并发请求数量，我们首先计算单个 256k 请求会占用多少缓存：

然后，要计算这一配置所支持的并发请求总数，我们将两者相除：

并发请求总数表示任意时刻同时流经流水线的请求数量。与 GPU 批处理不同，这些请求不会等待彼此完成，而是会持续不断地流过流水线。

当然，并发数和批大小都无法准确反映推理系统的运行效率。在评估推理系统吞吐量时，更有意义的指标是单位功率下的 token 总吞吐量（例如，每兆瓦每秒的 token 数），并且需要在相同的输出生成速度下测量。

如何扩展到超过 3T 参数

同样的方法也适用于更大的模型：计算权重和请求状态所需的内存，将模型分布到各片晶圆上，并将它们连接成流水线。每增加一片晶圆，都会增加计算能力和片上内存，使我们能够容纳容量超过任何单颗芯片的模型。

感谢您与我们一起学习！有意在 Cerebras 上部署模型或使用我们的推理云进行开发吗？

• 开始使用我们的免费 API：cloud.cerebras.ai
• 联系我们的销售团队：sales@cerebras.ai

Halley Chang 设计。感谢 Atharva Talpade、Zhenwei Gao、Angela Yeung、Jessica Liu、Joyce Er、Emma Call 和 Rohan Gupta 的审阅与反馈。