/AI 芯片的软件和硬件设计(17)——AI 助手/

翻译自 Bjarke Hammersholt Roune 的博文

/AI 助手的并行化/

本章从整体角度介绍如何将 AI 助手，也就是 LLM，并行部署到多颗芯片上。

作者指出，本节内容已经比较复杂，而真正要以高性能方式实现这些并行方法，其复杂程度还要远高于这里的描述。

如果读者只是希望进行相对轻松的阅读，可以直接跳到后面的**软硬件协同设计（Co-design）**章节，因为相比本节以及下一节关于 AI 芯片网络的内容，那一章更容易理解。

1

什么是并行化？

这里所说的并行化（Parallelization），是指将一个 LLM 的计算分布到：

多颗芯片

和/或：

单颗芯片内部的多个 Core

上执行。

例如，如果使用两颗芯片，理想情况下应该能够使计算速度提高到原来的：

2 倍

这就是并行化最基本的目标。

因此，作者认为，孤立地讨论单颗芯片性能实际上意义有限。

真正重要的是：

整个系统能够以多大的成本提供多高的性能。

也就是说，对于 AI 系统，更应该关注：

系统级性能/成本

而不仅仅是：

单芯片性能。

作者认为，目前很多 AI 芯片宣传材料主要强调单芯片性能，而忽略了系统级并行效率，这并不能完整反映 AI 芯片的实际价值。

2

理解 AI 并行需要首先掌握的三个事实

对于 AI 芯片初创公司而言，作者认为必须准确理解并行化产生的影响。

这是一个非常复杂，而且在很多情况下不符合直觉的问题。

作者首先总结了三个重要事实。

2.1

1. 网络带宽限制层内并行规模

层内并行（Within-layer Parallelization）受到网络带宽限制。

也就是说，如果将同一个 Transformer Layer 中的计算拆分到多颗芯片，那么这些芯片之间通常需要频繁交换中间数据。

因此：

更高层内并行度→更高芯片间通信带宽需求

这也是 AI 推理系统需要高速网络的重要原因。

相比之下，层间并行（Between-layers Parallelization），也就是 Pipeline Parallelism，对网络带宽的需求要低得多。

2.2

2. 更高网络带宽可以降低 Token 延迟

如果网络带宽提高：

2 倍

那么系统可能就能够支持大约：

2 倍层内并行度

而更高的层内并行度又可能将单个 Token 的计算分布到更多芯片，因此 Token 延迟最高可能降低接近：

1/2

这一点对于采用大型脉动阵列的 AI 芯片尤其重要。

大型脉动阵列通常需要较大的 Batch 才能保持高利用率，而较大的 Batch 可能导致较高的单 Token 延迟。

因此，如果芯片存在 Token 延迟过高的问题，可以通过增加网络带宽、提高层内并行度来缓解。

从这个角度看：

更大的脉动阵列→更高并行需求→更高网络带宽需求。

因此，作者认为，大型脉动阵列实际上也是推动 AI 芯片网络不断提高带宽的重要因素之一。

2.3

3. 不同并行方式对 HBM 容量的影响不同

层间并行，也就是 Pipeline Parallelism，可以降低单芯片的权重存储需求，但不能相同比例地降低单芯片的 KV Cache 需求。

相比之下：

层内并行可以降低单芯片 KV Cache 存储需求。

而层内并行规模又受到网络带宽限制。

因此，网络带宽实际上会间接影响每颗芯片需要配置多少 HBM。

例如，如果网络带宽提高 2 倍，使系统能够将层内并行度提高约 2 倍，那么在权重存储已经不再是主要问题的情况下，每颗芯片需要保存的 KV Cache 理论上可能降低到接近：

1/2

因此：

网络带宽不仅影响通信性能，也可能直接影响单芯片 HBM 容量需求。

如果模型进一步采用MoE（Mixture-of-Experts），问题还会更加复杂。

MoE 在某种意义上可以看作一种介于传统层内并行与其他并行方式之间的结构，而且其网络通信模式也有所不同。

此外，还需要记住一个比较直接但非常重要的规律：

如果单颗芯片计算速度提高 2 倍，那么网络速度通常也需要提高约 2 倍，才能继续跟上计算速度。

否则，原本不是瓶颈的网络就可能成为新的系统瓶颈。

3

为什么必须准确理解并行？

作者将在后续内容中进一步解释上述三个结论成立的原因。

作者特别强调，并行问题非常容易计算错误，即使经验丰富的工程师也可能出现误判。

例如，早期 TPU 的网络配置就曾因为对并行需求估算不准确，而出现明显的网络带宽过度配置。

因此，在设计 AI 芯片时，只要计算涉及并行，作者建议对相关推导进行多次检查。

因为如果并行模型理解错误，那么由此推导出的：

HBM 容量、网络带宽、Token 延迟

等指标都会失去意义。

换言之，在 AI 芯片设计中，并行不是一个可以最后再考虑的问题，而是决定大量核心硬件参数的基础因素。

/并行化中的基本概念/

为了进一步理解不同并行方式的优缺点，需要首先明确几个基本概念。

4

吞吐率：Tokens per Second

并行化最直观的收益是提高吞吐率（Throughput）。

对于 LLM，吞吐率通常使用：

Tokens/s

衡量。

例如，如果一颗芯片能够产生：

1000Tokens/s

那么理想情况下，两颗芯片应该能够产生：

2000Tokens/s

这就是最基本的并行扩展。

5

首 Token 延迟：Time to First Token

首 Token 延迟（First Token Latency）表示：

AI 从开始处理用户请求，到生成第一个 Token 所需要的时间。

首 Token 延迟可能比较长。

例如，AI 可能需要：

• 从磁盘读取相关数据；
• 恢复此前对话信息；
• 在内部执行较长时间的“思考”；
• 访问网页或其他外部数据；
• 确定大致应该如何回答。

因此，当用户提交一个问题后，AI 显示“正在思考”，但还没有开始输出内容，这段等待时间就属于首 Token 延迟。

作者特别强调：

First Token Latency≠Token Latency

首 Token 延迟本身并不是并行计算中特有的概念，但非常容易与后面讨论的 Token Latency 混淆，因此需要单独区分。

6

Token 延迟：ms per Token

Token Latency表示系统进入稳定输出阶段之后，生成一个 Token 所需要的时间。

通常使用：

ms/Token

表示。

例如，如果 AI 已经开始回答，但后续文字以很慢的速度逐个出现，比如：

1Token/s

那么就意味着 Token Latency 较高。

需要再次强调：

Token Latency 不是 First Token Latency。

第一个 Token 可能需要等待很长时间，而开始输出之后，后续 Token 可能以完全不同的速度生成。

7

Token Latency 与 Throughput 不是简单的倒数关系

这一点非常重要。

假设某个 AI 系统的吞吐率为：

1000Tokens/s

是否意味着一个用户能够在 1 秒内获得 1000 个 Token 的回答？

并不是。

原因在于，Throughput 统计的是整个系统所有对话产生 Token 的总量。

例如，系统同时服务：

1000 个用户

每个用户每秒只获得：

1 个 Token

那么系统整体吞吐率仍然是：

1000×1=1000Tokens/s

因此：

系统吞吐率很高，并不意味着单个用户的 Token 延迟很低。

进一步假设，其他 999 个用户全部停止使用 AI，只剩下一个用户。

此时是否意味着这个用户可以独占：

1000Tokens/s？

仍然不一定。

因为系统可能只有在同时处理大量相互独立的 Conversation 时，才能达到 1000Tokens/s。

单个 Conversation 中的 Token 具有自回归依赖，因此无法简单地把原来跨 1000 个 Conversation 获得的吞吐率全部集中到一个 Conversation 上。

所以，即使其他用户全部退出：

单用户 Token Latency 也可能基本不变。

8

低 Token Latency 也不代表 Throughput 只能等于其倒数

反过来也是如此。

假设：

Token Latency=1ms/Token

是否意味着：

Throughput=1000Tokens/s？

也不完全正确。

此时系统吞吐率至少可以达到 1000Tokens/s，但可能远高于这一数值。

例如，每 1ms：

Conversation A 生成 1 个 Token

同时：

Conversation B 也生成 1 个 Token

那么 Token Latency 仍然是：

1ms/Token

但系统吞吐率已经达到：

2000Tokens/s

因此：

Token Latency=1ms 并不能唯一确定系统 Throughput。

综合来看：

而：

二者都是简单的指标，但不存在普遍适用的简单对应关系。

因此：

更高 Throughput 不一定意味着更低 Token Latency；

同样：

更低 Token Latency 也不一定意味着更高 Throughput。

在某些系统设计中，两者甚至可能呈现相反的变化趋势。

这也意味着，如果一份 AI 芯片资料只报告：

Throughput

或者只报告：

Token Latency

那么很难完整判断该系统是否真正具有优势。

因为要判断系统是否经济，需要关注 Throughput；而要判断用户实际体验是否足够好，则需要关注 Token Latency。

因此，评价 AI 推理系统时，通常需要同时观察这两个指标。作者也指出，现实中大量资料往往只报告其中一个。

9

Variability：执行时间波动

Throughput 和 Latency 通常都是平均值。

实际执行过程中，不同 Token 的处理时间可能存在明显差异：

某些 Token 很快，某些 Token 很慢。

这种波动称为：

Variability

在并行系统中，Variability 通常是不利的，因为不同计算任务执行时间不一致可能导致其他并行资源等待。

但这种波动并不总能避免。

作者特别指出：

MoE 是造成 Variability 的重要来源之一。

10

Storage：并行方式会改变单芯片存储需求

存储看起来似乎与并行没有直接关系，但对于 Transformer 而言，两者实际上高度相关。

不同并行方式可能显著：

增加或降低单颗芯片的存储容量需求。

例如，模型权重和 KV Cache 能否被有效分布到不同芯片，就直接取决于采用哪种并行策略。

因此，后续内容还会进一步分析不同并行方式对昂贵 HBM 容量需求的影响。

11

Local 与 Global：单芯片容量和系统总容量

分析并行系统时，必须严格区分：

Local Requirement：单颗芯片需求

与：

Global Requirement：整个系统总需求

例如，一个模型大小为：

1TB

而单颗芯片只有：

100GB HBM

是否意味着无法运行这个模型？

并不是。

100GB 只是：

Local Capacity

如果系统拥有：

20 颗芯片

那么整个系统的 Global HBM Capacity 就是：

20×100GB=2TB

因此，虽然任何单颗芯片都无法容纳 1TB 模型，但通过并行将模型权重分散到 20 颗芯片上，整个系统仍然可以容纳该模型。

这也是为什么讨论 AI 芯片存储能力时，不能只观察单颗芯片的 HBM 容量。

12

Minimum Installation Size：最小部署规模

Minimum Installation Size表示：

为了让一个 AI 模型能够正常运行，至少需要部署多少颗芯片以及相应的基础设施。

继续采用前面的例子。

如果一个 1TB 模型必须分布到 20 颗芯片才能运行，那么：

Minimum Installation Size=20 颗芯片

即使 20 颗芯片能够提供远高于实际业务需求的吞吐率，也仍然不能减少芯片数量，因为模型本身无法装入更少的芯片。

因此：

较大的 Minimum Installation Size 通常会降低系统部署的灵活性。

实际最低数量可能位于 10～20 颗之间，因为 HBM 除了存储权重之外，还需要存储 KV Cache 等其他数据。

13

Reliability：并行规模越大，可靠性问题越突出

计算机硬件可能发生故障。

例如：

• Cable 断开；
• Fan 停止工作；
• Chip 发生故障；
• 其他硬件组件失效。

如果一个 AI 任务需要：

1024 颗芯片共同工作

那么只要其中一颗芯片发生故障，就可能导致整个任务无法继续运行。

此时：

1 颗芯片故障→另外 1023 颗芯片也无法有效工作

直到故障芯片被识别并修复或替换。

因此，大规模并行会放大整个系统受到单点硬件故障影响的概率。

14

早期 TPU Pod 的实际案例

作者指出，早期 TPU 曾经实际遇到过类似问题。

TPU 可以组成：

1024 颗芯片的 Pod

但连接这些芯片的 Cable 偶尔会发生故障。

单独一根 Cable 的故障概率很低。

然而，当一个 Pod 中存在数千根 Cable 时：

至少有一根 Cable 发生故障

就不再是非常罕见的事件。

早期情况下，一根关键 Cable 发生故障可能导致整个 Pod 无法正常使用，直到数据中心工作人员找到并更换这根 Cable。

由于无法改变 Cable 本身的故障率，Sameer Kumar 修改了并行软件，使 TPU 系统即使存在少量：

故障、不可用或缺失的 Cable

仍然能够正常运行。

由于 TPU Pod 采用 Torus 网络拓扑，实现这种容错并不容易。

这一解决方案的代价是：

Torus 网络有效带宽降低约一半。

不过，TPUv3 的 Torus 网络原本就存在带宽过度配置，因此仍然有足够余量承担这一损失。

作者认为，这也是硬件资源过度配置有时能够产生额外价值的例子：

虽然最初的 Overprovisioning 并非有意为之，但额外资源为解决设计阶段没有预料到的问题提供了灵活性。

15

Failure Domain：故障域

**Failure Domain（故障域）**是指：

某个组件发生故障后，会因此停止工作的整个系统范围。

例如，如果一个由：

1024 颗 TPU

组成的 Pod，会因为其中一根 Cable 断开而整体停止运行，那么这根 Cable 对应的 Failure Domain 就包含：

全部 1024 颗 TPU。

因此，并行化之所以可能降低系统可靠性，一个重要原因就是：

并行规模扩大了 Failure Domain。

如果两颗芯片必须协同完成一个任务，那么其中一颗芯片故障就可能影响另一颗芯片的有效工作。

Minimum Installation Size 通常也容易成为 Minimum Failure Domain。

例如，如果一个模型必须至少使用：

20 颗芯片

才能运行，那么只要其中任意一颗芯片故障，而系统又没有相应的容错机制，这 20 颗芯片组成的整体就可能无法继续正常执行该模型。