AI推理计算的可能终局：存算一体芯片！

生成式AI为提高效率带来了令人难以置信的商机，但我们仍然面临着一个阻碍广泛采用的巨大挑战：运行推理的高昂成本。

我们已经听说过大型语言模型（LLM）训练期间产生的惊人费用：大量的GPU，高昂的电费。分析师估计，Meta 2024年可能会在GPU上花费150亿美元。生成式AI需要大量的内存和带宽来进行权重计算和数据处理，这是大规模运行模型的主要障碍。就连OpenAI的Sam Altman也表示：“没有突破，就不可能实现这一目标。这促使我们加大对核聚变的投资。”

迄今为止，CPU、GPU和许多定制设计的加速器一直是AI的主要选择，但即使是最先进的解决方案也会被传统的冯·诺伊曼架构所拖累。定制设计的ASIC、FPGA和APU提供具有更高内存带宽的专用控制器，并且需要大量RAM来将模型保存在内存中，这使得生成式AI几乎不可能经济地交付。

尽管存在这些成本和计算限制，但很快企业将广泛部署多种AI模型，扩大对推理的需求。这将需要计算能力的另一个显著提高。生成式AI推理需要更多的计算和内存容量，因为今天的模型比过去的非生成式ML模型大得多。

推理的计算不仅取决于模型的大小，还取决于用户输入提示的大小。例如，当提示符长度从8K增加到32K时，OpenAI的GPT-4运行成本增加了一倍。另一方面，模型生成每个输出token的速度是一个重要的用户体验度量。这里内存带宽很重要，但是遗留架构的内存带宽限制加剧了推断成本和功耗的痛点。

AI社区正在积极寻找解决方案来应对这些挑战，包括气候友好型能源（风能、太阳能等）、新的硬件设计方法和算法优化。今天我想谈谈解决这个问题的方法。

存储和能量墙

加州大学伯克利分校BAIR/SkyLab的研究科学家Amir Gholami及其同事在他们的文章“AI和内存墙”中描述了训练和服务transformer模型的挑战。“内存墙”是William Wulf和Sally McKee在1995年创造的一个术语，它既包括有限的容量，也包括存储传输的带宽。分布式推理可以避免单个加速器有限的内存容量和带宽，但这种方法也面临内存墙问题：在加速器之间移动数据的通信瓶颈，这比芯片上的数据移动更慢，效率更低。

从训练过的权重中创建内容需要大量的微小计算。为了进行这些计算，使用了GEMM（General Matrix Multiply，通用矩阵乘法）操作。GEMM不需要大型处理器，但它需要快速高效的小型计算。

这就是记忆墙阻碍的地方。对于每个操作，数据必须在DDR RAM和处理器之间以及处理器之间传输。即使这段旅程只有毫米的距离，由于冯·诺依曼架构中存储和计算处理器的分离，它也需要时间和能量。而且有这么多的计算！

每次数据在存储器总线上移动时，DRAM访问每字节消耗大约60皮焦耳，而每次操作的计算只需要50-60飞焦耳。这意味着来回移动数据所消耗的能量是实际使用数据所消耗能量的一千倍。随着这些能源成本累积到数百万用户提示，每个GPU数百瓦，以及全球数千台服务器和数据中心，你可以看到为什么Altman认为我们需要核能突破。

内存计算

如果在存储和处理之间来回移动数据的效率如此之低，那么我们是否可以将计算移到内存中呢？近年来，内存计算（in-memory computing，简称IMC）通过直接在内存单元附近/内部执行MAC（乘-累加）操作而成为一种很有前途的替代方法。

普渡大学的研究表明，与机器学习推理的既定基线相比，内存计算架构的能耗降低至0.12倍。位于比利时KU Leuven研究型大学的MICAS中心支持IMC，强调其降低访问开销和实现大规模并行化机会的能力，可能导致能源效率和吞吐量的数量级提高。

该领域的早期探索将模拟IMC视为评估权重和在预训练的LLM上运行推理的有效方法。然而，这种方法需要昂贵的数模转换器和额外的错误检查。

数字内存计算（DIMC）提供了一种替代方案，避开了模拟IMC的挑战，提供无噪声计算和更大的空间映射灵活性。与模拟相比，DIMC牺牲了一些面积效率，但为处理未来的AI需求提供了更多的灵活性和能力。KU Leuven的研究还支持SRAM作为IMC的首选解决方案，因为与基于NVM的解决方案相比，SRAM具有鲁棒性和可靠性。

DIMC有望彻底改变AI推理，降低成本并提高性能。鉴于生成式AI的采用速度很快，只有通过将计算引入内存并提高性能来追求降低成本和功耗的新方法才有意义。通过减少不必要的数据移动，我们可以显著提高AI的效率，并改善AI未来的经济状况。

Sree Ganesan是d-Matrix的产品副总裁，d-Matrix是一家为生成式AI推理开发AI芯片的初创公司。