设计师如何应对 AI 的内存瓶颈

怀疑论者对当前人工智能技术的批评之一是内存瓶颈——由于无法加速处理器和内存之间的数据移动——阻碍了有用的现实世界应用程序。

用于在数据中心训练 AI 模型的 AI 加速器需要可用的最高内存带宽。在理想情况下，整个模型可以存储在处理器中，这种方法可以消除等式中的片外存储器。这是不可能的，因为最大的模型可以测量数十亿或数万亿个参数。

过去的系统内存受限，而今天的数据中心架构使用各种技术来克服内存瓶颈。

高带宽内存

一种流行的解决方案是使用高带宽内存 (HBM)，它涉及通过硅中介层将 4、8 或 12 个 DRAM 裸片的 3D 堆栈连接到处理器。该技术的最新版本 HBM2E 具有比其前身更快的每针信号速率，每针高达 3.6 Gb/s，从而提高了带宽。三星和 SK 海力士各自提供 8 个芯片 HBM2E 堆栈，总容量为 16 GB，提供 460 GB/s 带宽（相比之下，DDR5 为 2.4 GB/s，GDDR6 为 64 GB/s，SK Hynix 表示）。HBM3 旨在将速度和容量推得更高。

英伟达的 A100 数据中心 GPU，带有六层 HBM2E 内存（出于良率原因，仅使用了五层）（来源：英伟达）

最新版本的 Nvidia 旗舰数据中心 GPU A100提供 80 GB 的 HBM2E 性能和 2 TB/s 的内存带宽。A100 包含五个 16-GB DRAM 堆栈，加入一个使用 HBM2 的 40-GB 版本，总带宽为 1.6 TB/s。两者之间的差异使深度学习推荐模型的 AI 模型训练速度提高了三倍，这是一种已知的内存消耗。

同时，数据中心 CPU 正在利用 HBM 带宽。Intel 的下一代 Xeon 数据中心 CPU Sapphire Rapids 将 HBM 引入 Xeon 系列。它们是英特尔首款使用专为 AI 等矩阵乘法工作负载设计的新 AMX 指令扩展的数据中心 CPU。他们还可以使用片外 DDR5 DRAM 或 DRAM 加 HBM。

“通常，CPU 针对容量进行了优化，而加速器和 GPU 针对带宽进行了优化，”英特尔高级首席工程师 Arijit Biswas 在最近的 Hot Chips 演示中说。“然而，随着模型大小呈指数级增长，我们看到对容量和带宽的持续需求没有权衡取舍。Sapphire Rapids 通过原生支持这两者来做到这一点。” 通过内存分层进一步增强了该方法，“其中包括对软件可见的 HBM 和 DDR 的支持，以及使用 HBM 作为 DDR 支持的缓存的软件透明缓存，”Biswas 补充道。

Sapphire Rapids 的首席工程师 Nevine Nassif 告诉EE Times，HBM 版本是以牺牲芯片面积为代价的。

“[HBM 兼容] 模具略有不同，”Nassif 指出。“还有一个不同于 DDR5 控制器的 HBM 控制器。在没有 HBM 的 Sapphire Rapids 版本中，我们在芯片的一个区域添加了用于加密、压缩等的加速器。所有这些都消失了——除了数据流加速器——而 HBM 控制器取而代之。最重要的是，我们必须对网格进行一些更改，以支持 HBM 的带宽要求，”她补充道。

除了 CPU 和 GPU，HBM 在数据中心 FPGA 中也很受欢迎。例如，英特尔的 Stratix 和赛灵思 Versal FPGA 都有 HBM 版本，一些 AI ASIC 也使用它。腾讯支持的数据中心 AI ASIC 开发商 Enflame 将 HBM 用于其 DTU 1.0 设备，该设备针对云 AI 训练进行了优化。80 Tflops (FP16/BF16) 芯片使用两个 HBM2 堆栈，通过片上网络提供 512 GB/s 带宽。

燧发DTU 1.0数据中心AI加速芯片有两层HBM2内存（来源：燧发）

每美元的性能

虽然 HBM 为数据中心 AI 加速器所需的片外内存提供了极高的带宽，但仍然存在一些值得注意的问题。

Graphcore 就是其中之一。在他的 Hot Chips 演示中，Graphcore 首席技术官 Simon Knowles 指出，在大型 AI 模型中更快的计算需要内存容量和内存带宽。虽然其他人使用 HBM 来提高容量和带宽，但权衡包括 HBM 的成本、功耗和热限制。

Graphcore 对不同内存技术的容量和带宽的比较。当其他人尝试使用 HBM2E 解决这两个问题时，Graphcore 在其 Colossus Mk2 AI 加速器芯片上使用了主机 DDR 内存和片上 SRAM 的组合（来源：Graphcore）

Graphcore 的第二代智能处理单元 (IPU)改为使用其大型片上 896 MiB SRAM 来支持为其 1,472 个处理器内核提供所需的内存带宽。Knowles 说，这足以避免卸载 DRAM 所需的更高带宽。为了支持内存容量，太大而无法在芯片上安装的 AI 模型使用服务器级 DDR 形式的低带宽远程 DRAM。该配置连接到主机处理器，允许中型模型分布在 IPU 集群中的 SRAM 上。

鉴于该公司以每美元性能为基础推广其 IPU ，Graphcore 拒绝 HBM 的主要原因似乎是成本。

“与 AI 处理器集成的 HBM 的净成本是每字节服务器级 DDR 成本的 10 倍以上，”他说。“即使容量适中，HBM 也主导着处理器模块的成本。如果 AI 计算机可以使用 DDR，它可以部署更多 AI 处理器，但总拥有成本相同。”

Graphcore 对 HBM2 与 DDR4 内存的成本分析显示，前者的成本是后者的 10 倍。（来源：Graphcore）

据 Knowles 称，40 GB 的 HBM 有效地将封装的标线大小处理器的成本提高了三倍。Graphcore 的 8 GB HBM2 与 8 GB DDR4 的成本细分估计 HBM 芯片的尺寸是 DDR4 芯片的两倍（将 20-nm HBM 与 Knowles 认为是同时代的 18-nm DDR4 进行比较），从而增加了制造成本. 然后是 TSV 蚀刻、堆叠、组装和封装的成本，以及内存和处理器制造商的利润率。

“DDR DIMM 不会发生这种边距堆叠，因为用户可以直接从内存制造商处采购，”Knowles 说。“事实上，可插拔的计算机组件生态系统出现的一个主要原因是为了避免保证金堆积。”

走得更远

从 Hot Chips 的隐形模式中脱颖而出，Esperanto 提供了另一种解决内存瓶颈问题的方法。该公司的 1000 核 RISC-V AI 加速器针对的是超大规模推荐模型推理，而不是上面提到的 AI 训练工作负载。

世界语的创始人兼执行主席戴夫·迪策尔指出，数据中心推理不需要巨大的片上内存。“我们的客户不想要 250 MB 的芯片，”Ditzel 说。“他们想要 100 MB——他们想用推理做的所有事情都适合 100 MB。任何比这更大的东西都需要更多。”

Ditzel 补充说，客户更喜欢将大量 DRAM 与处理器放在同一张卡上，而不是在芯片上。“他们建议我们：'只需将所有内容都放到卡上，然后使用您的快速接口。然后，只要您能够比通过 PCIe 总线更快地获得 100 GB 的内存，这就是胜利。”

将 Esperanto 的方法与其他数据中心推理加速器进行比较，Ditzel 说其他人专注于消耗整个功率预算的单个巨型处理器。这家初创公司坚持认为，世界语的方法——安装在双 M.2 加速卡上的多个低功耗处理器——可以更好地使用片外内存。单芯片竞争对手“引脚数量非常有限，因此他们必须采用 HBM 之类的产品才能在少量引脚上获得非常高的带宽，但 HBM 确实很昂贵，而且很难获得，而且功率很高，”他说。

Esperanto 声称通过使用六个较小的芯片而不是单个大芯片解决了内存瓶颈，留下可用于连接 LPDDR4x 芯片的引脚（来源：Esperanto）

Esperanto 的多芯片方法使更多引脚可用于与片外 DRAM 通信。除了六个处理器芯片外，该公司还使用了 24 个专为手机设计的廉价 LPDDR4x DRAM 芯片，在低电压下运行，“每比特能量与 HBM 大致相同”，Ditzel 说。

“因为 [LPDDR4x] 的带宽 [比 HBM] 低，我们通过更宽获得更多带宽，”他补充道。“我们在加速器卡上的内存系统上使用 1,500 位宽 [而单芯片竞争对手] 负担不起 1,500 位宽的内存系统，因为对于每个数据引脚，你必须有几个电源和几个接地针，针太多了。

“我们之前处理过这个问题，我们说，让我们分开吧。”

通过 822 GB/s 的内存带宽访问 192 GB 的总内存容量。所有 64 位 DRAM 芯片的总和计算出 1536 位宽的内存系统，分成 96 个 16 位通道以更好地处理内存延迟。这一切都符合 120 W 的功率预算。

流水线权重

晶圆级 AI 加速器公司 Cerebras 设计了一个规模远端的内存瓶颈解决方案。在 Hot Chips 上，该公司发布了 MemoryX，这是一款用于其 CS-2 AI 加速器系统的内存扩展系统，旨在实现高性能计算和科学工作负载。MemoryX 寻求能够训练具有万亿或参数的巨大 AI 模型。

Cerebras 的 MemoryX 系统是其 CS-2 晶圆级引擎系统的片外内存扩展，其行为就像在片上一样（来源：Cerebras）

MemoryX 是 DRAM 和闪存的组合，其行为就像在芯片上一样。该架构被提升为弹性，旨在容纳 4 TB 到 2.4 PB（2000 亿到 120 万亿个参数），足以容纳世界上最大的 AI 模型。

Cerebras 的联合创始人兼首席硬件架构师 Sean Lie 表示，为了使其片外内存表现得像片上一样，Cerebras 优化了 MemoryX，以消除延迟影响的方式将参数和权重数据流式传输到处理器。

“我们将内存与计算分开，从根本上分解它们，”他说。“通过这样做，我们使沟通变得优雅而直接。我们可以这样做的原因是神经网络对模型的不同组件使用不同的内存。因此，我们可以为每种类型的内存和每种类型的计算设计一个专门构建的解决方案……”

结果，这些组件被解开，从而“简化了缩放问题，”Lie 说。

Cerebras 使用流水线在 AI 训练期间消除对延迟敏感的通信。（来源：大脑）

在训练期间，必须立即访问对延迟敏感的激活记忆。因此，Cerebras 将激活保留在芯片上。

Cerebras 将权重存储在 MemoryX 上，然后根据需要将它们流式传输到芯片上。Lie 说，在没有背靠背依赖性的情况下，权重记忆的使用相对较少。这可以用来避免延迟和性能瓶颈。粗粒度流水线还避免了层之间的依赖关系；层的权重在前一层完成之前开始流式传输。

同时，细粒度流水线避免了训练迭代之间的依赖关系；后向传播中的权重更新与同一层的后续前向传播重叠。

“通过使用这些流水线技术，权重流执行模型可以隐藏外部权重的额外延迟，如果权重在晶圆上本地 [访问]，我们可以达到相同的性能，”Lie 说。

审核编辑 黄昊宇