摘要：如今，AI 集群是高带宽内存 (HBM) 的主要用途之一。然而，由于多种原因，HBM 并不适合 AI 工作负载。分析表明，HBM 在写入性能上配置过多，但在密度和读取带宽上配置不足，并且每位能量开销也很大。它也很昂贵，由于制造复杂性，良率低于 DRAM。
如今，AI 集群是高带宽内存 (HBM) 的主要用途之一。然而，由于多种原因，HBM 并不适合 AI 工作负载。分析表明，HBM 在写入性能上配置过多，但在密度和读取带宽上配置不足，并且每位能量开销也很大。它也很昂贵，由于制造复杂性，良率低于 DRAM。
于是，我们（微软研究院团队)提出了一种新的内存类：托管保留内存 (MRM：Managed-Retention Memory)，它更适合存储 AI 推理工作负载的关键数据结构。我们相信，MRM 最终可能会为最初提议支持存储类内存 (SCM：Storage Class Memory ) 的技术提供一条可行的途径。这些技术传统上提供长期持久性（10 年以上），但提供较差的 IO 性能和/或耐用性。MRM 做出了不同的权衡，通过了解工作负载 IO 模式，MRM 放弃了长期数据保留和写入性能，以在这些工作负载的重要指标上获得更好的潜在性能。
介绍
迄今为止，存储领域一直呈现二元化：存在非易失性和易失性存储技术。不同形式的 DRAM（GDDR、HBM、LPDDR）是主要的易失性存储器存储技术。一旦切断电源，其存储的数据就会丢失。面向块的 NAND 和 NOR 字节寻址闪存是非易失性存储器最广泛使用的例子。它们不需要持续供电来保存数据。在存储单元级别，数据易失性表示为保留时间，即可靠存储数据而无需刷新的时间。闪存单元的保留时间为 10 年以上，但每个存储单元的读写吞吐量低于 DRAM。这些特性意味着 DRAM 用作处理器的内存，而 Flash 用于二级存储。
其他几种存储器技术，如 RRAM、MRAM 和PCM，都具有提供非易失性的潜力。它们属于通常称为服务器存储级内存 (SCM) 的一类内存。最近停产的英特尔 Optane / 3D XPoint 是 SCM 的标志性代表，旨在克服 Flash 的 IO 限制，同时保持非易失性。其梦想是通过提供可比的 IO 性能和字节寻址能力来取代 DRAM，同时还具有 10 年以上的保留时间。
然而，由于权衡利弊，迄今为止的所有尝试都未能取代 DRAM。由于密度和复杂的制造工艺的挑战，它们未能以低于（或等于）Flash 的成本提供与 DRAM 相当的 IO 性能。对于主存储器，数据的持久性并不像 IO 性能那么重要。对于一般的计算工作负载，没有人愿意用主存储器的 IO 性能来换取 10 年以上的数据保留时间。这些技术还在耐久性方面存在困难，耐久性指的是存储单元在永久退化之前可以支持的写入周期数。因此，SCM 最终对某些用例（例如嵌入式计算)有用，但不能部署在服务器中。
讽刺的是，我们认为闪存的兴起可能是内存创新的祸根。非易失性是存储设备的一个关键属性，但在内存单元层面，它却具有很大的误导性。对于所有技术，内存单元都只提供保留时间，这是一个连续的过程，从 DRAM 的微秒到多年的时间。SCM 所依赖的技术被迫采用非易失性，要求其保留时间为十年或更长。不幸的是，实现这些高保留时间需要权衡其他指标，例如写入和读取延迟、能源效率和耐用性。
或许，历史上将其视为二进制的一个原因是，即使保留时间宽松，SCM 技术也无法在所有对一般工作负载重要的指标上与 DRAM 匹敌。然而，基础模型（大型语言模型或 LLM 是其子集）最近已成为一种具有独特内存 IO 要求的新的主要工作负载。基础模型训练和推理的巨大规模和增长需要新颖的硬件方法。基础模型推理对内存 IO 的要求与历史工作负载不同。例如，很大一部分内存用于存储模型权重，对于顺序读取而言，IO 性能至关重要，但对于写入则不那么重要。内存 IO 是顺序的和可预测的，考虑到 AI 集群的能源挑战，每位读取的能量也是一个问题。
目前唯一可以匹配 IO 性能、能量和密度的技术是 HBM。然而，它不是万能的，基础模型推理中的某些关键阶段受内存而不是计算限制。此外，HBM 价格昂贵，并且在产量方面面临巨大的挑战。
我们认为，有机会重新考虑现有的“非易失性”内存技术，以适应这种新的工作负载。我们提出了一种新的内存类别，我们称之为托管保留内存 (MRM)。MRM 与易失性 DRAM 不同，因为它可以在没有电源的情况下保留数据，并且不会在频繁的单元刷新中浪费能源，但与 SCM 不同，它不以长期保留时间为目标。由于大多数推理数据不需要持久保存，因此保留时间可以放宽到几天或几小时。作为回报，MRM 具有更好的耐用性，并旨在在读取吞吐量、能源效率和容量等关键指标上超越 DRAM（和 HBM）。
基础模型时代的存储
基础模型的工作负载与传统工作负载截然不同。基础模型首先进行训练，通常在大型集群（例如 50,000 多个 AI 加速器）上进行，输出本质上是一组模型权重。然后，这些权重被部署到生产中，用于推理查询。将使用数千甚至数百万个基础模型实例，但每次推理所需的硬件规模要小得多（例如 4 个以上的 AI 加速器）。
据观察，训练和推理工作负载都是内存密集型的。训练规模取决于模型大小，并且是一次性的工作（通常需要数月），而推理工作负载是需求驱动的，并且在相当长的一段时间内提供服务，直到模型权重被淘汰。
训练和推理具有不同的内存访问模式和要求，并且通常部署在不同的集群上。随着需求的增加，我们预计推理基础设施将占据主导地位，因此我们将重点关注推理工作负载。更具体地说，我们考虑执行自回归标记生成的基础模型。推理查询是一系列输入标记，基础模型会响应该查询生成一系列输出标记。上下文由来自用户的所有标记以及模型在交互过程中生成的相应响应组成。上下文尽可能大是可取的，因为它可以通过使用自注意力机制来提高模型的推理能力 。然而，在部署中，上下文的大小有限，范围从 1000 个到几万个 token（取决于模型），并且主要受可用内存量的限制。每个推理查询的计算成本很高，需要跨多个 AI 加速器进行分布式计算。
推理依赖于三个主要的内存数据结构：模型权重、KV 缓存和模型激活。其中，模型权重和 KV 缓存占用了大部分内存容量 （Kwon 等人，2023）。
模型权重（矩阵）是扩展前沿基础模型功能的关键；随着每一代基础模型的出现，模型权重的大小都呈指数级增长。目前，大型模型的权重超过 5000 亿个，根据所使用的权重量化，代表 250 GB 到 1 TB 以上的数据。权重实际上是一种不可变的数据结构。参考模型权重保存在存储中，而副本分布在每个推理集群中的 AI 加速器上。如今有大量的基础模型，但实际上，只有少数最受欢迎的模型被大规模使用。对给定基础模型版本（例如 GPT4）进行的所有推理查询都使用相同权重的副本。
KV 缓存支持模型的自注意力机制。它是一系列自注意力向量，用于编码模型对上下文中所有 token 之间关系的理解。每次在上下文中生成新 token 时，都会将一个向量附加到相应 KV 缓存的末尾。每个向量通常为几 MB，因此 KV 缓存通常会增长到几十 GB，直到达到上下文大小限制。
最后，模型激活是网络前向传递期间创建并在不同层之间传递的瞬态张量。它们通常比权重和 KV 缓存小一个数量级，并且仅在前向传递计算期间存储。
KV 缓存是在预填充阶段创建的，即从用户那里收到第一组token时。随后，在解码阶段，模型会迭代生成响应token。为此，每次迭代时都会按顺序完整读取 KV 缓存，生成新token，并将相应的自注意力向量附加到 KV 缓存。KV 缓存利用内存来减少计算，并且是软状态：它们由模型生成，并且可以在需要时重新计算。但是，每秒的token率通常很低（因此成本很高），因此缓存和使用 KV 缓存通常比重新计算更可取。
在推理过程中，会读取每个生成的 token 的整个自注意力数据和权重，从而在内存和计算之间产生大量带宽需求。在任何给定时间，许多推理请求都会在同一个集群上进行多路复用，但它们都针对同一个模型。因此，每个 AI 加速器的内存都包含模型权重的子集，以及与上下文工作集相对应的多个 KV 缓存和激活。部署新模型时，集群会停止接受新请求，为正在进行的请求提供服务，然后为新模型加载权重。
总而言之，基础模型推理主要由非常大、可预测的内存读取组成，而写入则较小且大多仅附加。需要读取的确切内存范围是预先知道的，并且大部分内存不会长时间被覆盖。然而，尽管以读取为主，但推理仍然需要与存储工作负载相比非常高的写入速率。
1 .
HBM 的诅咒
如今，AI 加速器中使用的大多数数据都存储在 HBM 上，因为所有数据结构都需要以高带宽重复读取。当前的 AI 加速器可以支持非常高的主内存带宽，例如单个 B200 GPU 的带宽为 8 TB/s。此外，由于权重和 KV 缓存很大，AI 加速器需要大量的 HBM 容量。要实现这一目标，需要克服艰难的工程挑战，尤其是在能源使用方面。在所需数据速率下通过铜互连跟踪信号丢失意味着内存必须物理地位于（非常）靠近计算芯片的位置，通常共同封装在同一个中介层上。非常宽的接口和高信号速率意味着更多的能量，而 AI 加速器大约三分之一的能量使用是内存。
使用 HBM 是因为它能够在同一封装上对 DRAM 进行 3D 堆叠，以提高封装内存容量、吞吐量，并最大限度地缩短内存单元与 AI 加速器的距离。目前的 HBM 产品有 8-12 层，B200 封装上总共 192 GB。因此，HBM 被使用，因为它在合理的能源使用下以最高的密度提供最高的吞吐量。然而，即使使用 HBM，每个推理查询的很大一部分也受到内存限制 。
不幸的是，目前还没有可行的 HBM 替代品。非堆叠 DRAM 不具备所需的密度，而 NAND 和 NOR 闪存速度不够快，使用寿命短，尤其是在每个存储单元存储多个位的较高密度下。两者都缺乏封装所需的能效。
值得注意的是，HBM 面临着几个基本挑战。首先，内存供应商正在努力继续扩展密度。每层扩展面临着从 DRAM 继承而来的挑战。因此，与当前的 HBM3e 相比，下一代 HBM（HBM4）预计每层容量仅增加 30%；其次，DRAM 的 3D 堆叠不仅大大降低了制造工艺的良率，还带来了散热挑战，尤其是在与 AI 加速器芯片紧密封装时。
目前，业界预计在可预见的未来它不会扩展到 16 层以上。3D 堆叠极其复杂。这些因素加上云基础设施对基础模型的指数级增长带来的高需求，意味着 HBM 占 AI 集群成本的很大一部分。这种情况在可预见的未来不太可能改变，AI 集群仍将依赖 HBM。
2 .
新的希望？
Foundation 模型推理与 DRAM 设计的通用主内存工作负载有很大不同。首先，它是极其密集的读取操作。例如，解码期间生成的每个 token 都需要读取所有权重和整个 KV 缓存 ，用于一次自注意力向量写入。自注意力向量大小通常最多为几MB, 而权重和KV缓存通常为数十GB，这意味着读写比超过1000:1。
有人努力减少推理过程中读取的数据量。例如，批处理允许在请求之间重用权重。然而，批处理受到延迟要求的限制也被使用，但是它们各自都有其局限性，而且即使结合起来，它们也无法从根本上改变工作负载以读取为主的性质。
其次，内存访问是顺序的且可预测的。权重或 KV 缓存没有就地更新，并且每个基础模型响应都会迭代读取相同的权重和 KV 缓存。已经提出了内存虚拟化机制来解决内存碎片问题 ，但即使在这种情况下，页面的读取顺序也是相同的。每个页面通常超过 10 个向量（通常为几 MB 到几十 MB），并按顺序读取。另外，虚拟页面和物理地址之间的映射通常是静态的。
这些特性表明，大多数 HBM 容量用于存储与 HBM 从 DRAM 继承的通用特性（随机访问、字节寻址、可比的读写性能）几乎无用的数据。从某种意义上说，HBM 超出了此基础模型推理工作负载的要求。这种过度配置会导致成本和能源开销不理想。
这也引发了一个诱人的问题：如果我们正确地为工作负载配置内存，我们是否可以解决推理集群中内存的次优成本和能源挑战？
存储的机会
我们认为，(i) 基础模型基础设施的重要性和规模、(ii) 传统服务器 CPU 和 AI 加速器的工作负载模式之间的巨大差异以及 (iii) HBM 与工作负载的匹配性较差，为该应用开辟了更好内存的计算机架构研究领域。
我们现在认为，这种机会最好通过一种新型内存来解决，而不是 DRAM、HBM 或闪存。闪存无法使用，因为它没有足够的耐久性，即使是使用单层单元 (SLC) 也是如此, 因为它们无法满足高吞吐量和能源效率的要求。Flash 的非易失性也是不必要的：数据要么保存在其他地方（权重），要么是软状态（KV 缓存、激活）。
另一方面，一些工作负载属性接近存储工作负载通常表现出的属性。例如，字节寻址能力不是必需的，因为 IO 很大且是连续的。与存储基础设施类似，存储容量和总拥有成本 (TCO)/TB 是关键指标，而 DRAM 和 HBM 在这方面表现不佳。
MRM 能否满足 AI 集群要求？PCM、RRAM 和 STT-MRAM 的读取性能和能耗与 DRAM 相当甚至更好，并且具有更高的密度和/或更低的 TCO/TB 潜力。
图 1.KV 缓存和模型权重的耐久性要求与内存技术的耐久性。
这些技术的耐久性低于 DRAM，我们现在估计权重和 KV 缓存写入的近似耐久性要求。权重更新不频繁，在更换模型时会批量覆盖。更新频率目前通常较低（数小时以上），但可能会随着模型的多样化而发展。我们估计保守的每小时更新和每秒一次的密集更新需要 5 年以上的耐久性。KV 缓存写入在预填充和解码期间都会发生，每个上下文标记一个自注意向量。预填充的吞吐量通常高于解码，我们使用 Splitwise 中针对 Llama2-70B 模型报告的吞吐量和中位上下文长度。对于预期五年的寿命，我们计算了 KV 缓存写入的次数，并推断出每个单元的平均写入次数。
图 1显示了现有内存/存储技术的耐久性与工作负载耐久性要求之间的比较。在适用的情况下，我们将现有设备中观察到的耐久性与该技术所展示的潜力区分开来。我们观察到 1) HBM 在耐久性方面配置过多，2) 现有的 SCM 设备不满足耐久性要求，但底层技术有潜力做到这一点。我们认为这部分是由于当前的设备是为非易失性设计的，这是通过权衡其他重要指标（例如写入延迟、能源效率或耐久性）来实现的 。我们认为这是一个重新思考现有内存技术的机会，这些技术目前用于 SCM，特别是用于 AI 工作负载，通过权衡非挥发性与其他关键指标。
软件堆栈影响
在本节中，我们将阐述计算机系统社区为何对 MRM 感兴趣。
基础模型正变得无处不在，这导致需求多样化：一些用例具有严格的延迟 SLA（例如，用户在环对话），一些用例对吞吐量要求高且大量使用批处理，另一些用例是后台尽力而为的作业（例如，会议回顾）。工作负载变得越来越复杂，输入：输出token比率差异很大，针对特定用例量身定制的专家模型，以及对高级增强机制的依赖。
此外，工作负载的资源密集型特性和硬件成本需要全面而高效的编排。这可以通过利用关键操作系统机制，有效地构建了机架规模的基础模型推理操作系统。在这种背景下，MRM 的出现带来了一系列令人兴奋的挑战和探索机遇。
（1）保留感知数据放置和调度。
MRM不太可能成为一种万能的解决方案，它将与其他类型的内存共存，例如用于写入密集型数据结构（例如激活）的 HBM 和作为较慢层的 LPDDR。需要对数据的生命周期和访问模式有深入的了解才能对数据进行布局。调度程序需要跟踪数据到期时间，并根据依赖于该数据的请求的状态决定是否刷新它或将其移动到另一层。
（2）轻量级内存控制器。
有可能使 MRM 控制器变得极其简单和节能。缺乏随机访问要求为块级访问内存控制器开辟了独特的前景，这对软件堆栈有影响。许多通常在设备上处理的功能（例如刷新和磨损均衡）可以留给堆栈中更高级别的软件控制平面，该平面最适合做出这些决策，同时满足全局应用程序要求。这种方法类似于 Flash 的分区存储接口。
（3）动态可配置内存 (DCM：Dynamically Configurable Memory)。
由于控制平面对应用程序和用户工作负载都具有集群级可见性，因此它最适合在写入数据时动态决定每个数据所需的保留期，从而有效地将 MRM 正确配置到工作负载。这是 MRM 的完全灵活实例。在硬件级别，内存控制器将支持以不同的持续时间和能量进行写入，从而允许在运行时对保留时间进行编程。然后，基础模型操作系统可以在集群级别协调最佳数据刷新、磨损均衡和垃圾收集。
（4）保留感知错误更正。
MRM放宽的保留要求也提出了一个有趣的问题：我们如何看待数据完整性？存储在 MRM 中的大部分数据要么持久存储在其他地方（例如权重），要么是软状态（例如 KV 缓存）。因此，对持久性的要求并不像传统存储系统那么严格。尽管如此，系统仍然需要强制完整性，以保证涉及数据的计算的正确性，并避免频繁重新计算软状态。利用现有的最先进的内存纠错技术是一个好的开始，然而，基于大型块的 MRM 接口意味着可以考虑对更大代码字进行操作且开销更少的纠错技术。
相关工作
对于 STT-MRAM 和 RRAM，人们已经对保留、耐久性和写入能量效率之间的权衡进行了深入研究 。有人提出利用这种机制来提高混合片上 CPU 缓存的能源效率 。与我们的工作相比，这一系列工作侧重于通用多核 CPU，因此解决了不同的优化问题。AI 集群具有机架级能源和冷却要求，并且具有更复杂的内存层和互连，但工作负载更可预测。
斯坦福大学最近启动了一个为期 5 年的项目，以解决预计即将到来的主存储器分层和异构性增加的问题 。我们有一个共同的观察，那就是记忆墙 是关键工作负载面临的一大挑战，并且由于缺乏一刀切的技术，可能会导致更多的内存异构性。虽然这种趋势在数据中心是新奇的，但在其他应用中却很常见。例如，嵌入式世界历来使用的 ROM（只读存储器） 是一种一次写入多次读取的技术，EPROM (可擦除可编程只读存储器) 写入少量读取大量，用于存储程序，可以使用紫外线擦除，当然还有 RAM。ROM 和 EPROM 提供非易失性存储，必须做出谨慎的设计选择，才能最大限度地利用不同技术的优势。
人们正在努力通过紧密集成内存和计算来克服内存壁垒。这可以通过在计算芯片上添加更多内存来实现 ，或使用内存计算（IMC）或RRAM 为基础。我们的工作是正交的，因为它旨在优化主流内存/计算模型，而不是探索新的范式。
最后，在利用 AI 集群中的异构内存访问模式方面有大量工作要做。例如，有人提出使用 CPU 主内存来卸载空闲的 KV 缓存 。最新的 Nvidia GB200 超级芯片集成了 LPDDR5 控制器，用于更高容量、更慢的内存层。这表明内存异质性将成为 AI 集群中的常见现象。我们的工作建议利用数据访问异质性的更多方面来最大化每美元产生的代币。
结论
AI 工作负载的出现及其对 HBM 内存的依赖凸显了 HBM 的局限性。AI 推理工作负载需要高读取吞吐量、密度和能效，而 HBM 很难以经济高效的方式提供这些。我们提出了一种可以与 HBM 共存的新型内存，即托管保留内存 (MRM)，它允许使用最初为 SCM 提出的内存技术，但会牺牲保留和其他指标（如写入吞吐量）来换取对这些 AI 工作负载至关重要的改进性能指标。通过放宽保留时间要求，MRM 可以潜在地使现有的 SCM 技术提供更好的读取吞吐量、能效和密度。我们希望本文能够真正开启关于内存单元技术和内存芯片设计创新的新思路，专门针对 AI 推理集群的需求。
来源：智慧芯片
