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

在这篇文章中，我们将解释如何把一个大型前沿模型拆分到多块WSE-3T晶圆上，从而以每秒 1000-2000+ token 的速度运行数万亿参数的模型。

我们以开源模型 Kimi K3（约 3 万亿参数）作为示例来说明。

读完后，你将了解：

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

以 Kimi K3 为例拆解

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

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

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

#1：每一层的大小

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

Kimi K3 有 3 万亿参数。如果全部 3T 参数都以 4 位存储，原始参数负载为：

把这个负载除以 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 微秒。

需要强调的是，由于每块 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 处理一个，阶段 1 则为请求 C 处理一个。在我们的示例中，每个阶段是一层。在生产环境中，Cerebras 使用的阶段粒度要细得多，数量也多得多。

Video thumbnail

正在加载播放器...

每个阶段还可以处理一个微批次，一小批请求，而不是单个请求。因此，在途请求的数量大致等于阶段数乘以微批大小。

Video thumbnail

正在加载播放器...

要理解这一权衡，更小的微批能让每个请求更快，而更大的微批则能提升总吞吐量。

#6：我们能服务多少个请求？

分配给 KV 缓存的片上内存量决定了最大并发请求数。每个在途请求在每个 dense 层晶圆上都需要自己的 KV 缓存，因此片上内存的大小设定了并发请求的上限。如上所述，我们只需要计算 dense 层的 KV 缓存需求。

对于每块已经包含一个 dense 层的 44 GB 晶圆，我们可以这样计算 KV 缓存的剩余空间：

晶圆总存储 - 一个 dense 层的大小。在本例中：

计算并发请求数

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

然后要计算该配置支持的总并发请求数，我们做如下除法：

总并发请求数代表在任意时刻流经流水线的同时请求数量。与 GPU 批处理不同，这些请求不会互相等待完成，而是持续流经流水线。

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

这如何扩展到 3T 参数以上

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

感谢与我们一起学习！有兴趣在 Cerebras 上部署模型，或用我们的推理云进行构建吗？

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