跳到正文
Hacker News· mdp2021·· 6 天前AI 评分34

Tensordyne 如何用对数数字格式提升 AI 计算效率

Tensordyne is making AI compute more efficient using logarithms

AI 导读

Tensordyne 采用基于对数的数字格式,将乘法转为加法运算,从而比传统浮点芯片更高效地做矩阵乘法,其软件栈可将 PyTorch 模型转换到该格式运行。但该格式让加法更难,需依赖近似与专有技巧,且客户若要移植为 GPU 优化的自定义 kernel 需大量工作,训练负载也因数值稳定性问题难以支持。文章还讨论了生成式 AI 负载多为内存受限,批处理才能让计算引擎效率真正发挥作用。

正文 · AI 翻译

大多数 AI 芯片公司都在通过改变内存层次结构来加速生成式 AI 工作负载。Groq 通过使用速度极快但昂贵的 SRAM 内存,而非 Nvidia 所依赖的较慢但更具成本效益的高带宽内存(HBM),以高成本实现了超低延迟。Cerebras 正在利用晶圆级处理将这一理念进一步推进,每片晶圆提供 44GB 的 SRAM。d-Matrix 正在利用存内计算,通过减少数据进出内存的移动来提升性能。Taalas 正在使用只读存储器,它比 SRAM 更快、更便宜,但无法更改——因此每块芯片只能运行一个模型。

Tensordyne 对 AI 芯片采取了不同的思路。他们使用一种基于对数的不同数字系统,使其能够比使用传统浮点数学的芯片更高效地进行数字乘法。但将更高效的乘法器转化为实际 AI 工作负载上更好的性能并非易事。今天在博客中,我们将探讨 Tensordyne 的独特技术、他们的芯片相对于市场上其他芯片的定位,并尝试回答一个问题:更高效的计算引擎对 AI 芯片来说真的重要吗?

Tensordyne 架构的核心技术优势来自其对数数字格式。他们不是将值 A 直接存储为 A,而是实际存储 log2(A)。然后,如果他们想将 A 乘以另一个值 B,他们可以利用对数乘积法则,通过加法来完成乘法:

log2(A x B) = log2(A) + log2(B)

这个想法并不新鲜,但此前从未有人成功将其用于商用 AI 加速器。运行大型矩阵乘法需要大量乘法器,但也需要大量加法器。对数乘积法则使乘法更容易,却使加法变得困难得多。Tensordyne 通过一种近似方法绕过了这个问题:

log2(1+x) = x 当 0<x<1

但仅靠这个近似还不够。最初,Tensordyne 使用量化感知训练 来让模型在这种近似下运行,但说服客户仅为 Tensordyne 芯片重新训练模型将是一个巨大的挑战。因此,他们想出了一套额外的专有技巧,使其近似对于大型模型来说足够精确。

Tensordyne 的软件栈将 PyTorch 模型转换为使用其特定数字格式运行。我有点怀疑他们能否在不造成明显精度损失的情况下做到这一点,但即便能做到,这个转换过程还会带来另一个问题。现代优化推理流水线并不只是用 PyTorch 编写的;它们拥有针对 GPU 优化的自定义内核,以及其特定的数值特性。如果客户想将其移植到采用全新数字格式的芯片上,要匹配优化过的 GPU 内核的性能水平需要大量工作。Tensordyne 声称他们将让 AI 智能体帮助客户将代码转换为 Tensordyne 的软件和数字格式……这看起来有点像是敷衍的回答。从根本上说,编写内核很难,而要获得 Tensordyne 芯片(非常真实的)好处,客户可能需要硬着头皮将内核移植到 Tensordyne 的对数数字格式。

这种专用数字格式还使 Tensordyne 支持训练工作负载变得困难得多。推理工作负载对量化等数值变化相对宽容,但训练工作负载存在重大的数值稳定性问题。训练流水线是专门围绕硬件特性设计的,运行成本可能高达数亿美元。为 Tensordyne 的芯片构建新的训练流水线来应对训练这类工作负载根本说不通,毕竟训练主要还是受网络和内存带宽主导。

说到内存带宽……虽然 Tensordyne 在宣传其新的数字系统,但大多数 AI 芯片都专注于利用新的内存层级来实现突破性性能。对于生成式 AI 工作负载,Tensordyne 是不是关注错了瓶颈?

感谢阅读 zach 的科技博客!这篇文章是公开的,欢迎随意分享。

分享

大多数 AI 芯片专注于内存层级的创新,而不是从根本上改变计算方式,这是有原因的:许多生成式 AI 工作负载受内存限制。如果你想用单个 GPU 处理单个 LLM 查询,该 GPU 需要加载和卸载 LLM 中每一层的所有权重来处理该查询。即使你让 GPU 在矩阵乘法上高效得多,你仍受内存带宽的瓶颈制约。

GPU 通过批处理绕过这个问题。它们不是一次处理一个查询,而是处理一批多个查询。如果一个批次有 32 个查询,GPU 对它加载的每一层执行 32 组矩阵乘法。这使工作负载受内存带宽的瓶颈制约减少,并让你的计算引擎效率真正变得重要。这可以用所谓的“roofline 模型”来建模。随着我们增大批次大小,工作负载的算术强度会增加——本质上,每传输一字节内存所做的数学运算更多。算术强度越高,计算引擎效率就越重要。算术强度越低,内存带宽就越重要。

然而,批处理存在一个问题——它会拖慢 LLM 查询速度。批次中的第一个查询需要等待整个批次处理完毕。这就是为什么像 Cerebras 和 Groq 这样批大小为 1 的全 SRAM 架构能提供极低延迟的原因。全 SRAM 架构的支持者认为,延迟对于智能体工作负载极为重要,因为模型可能会运行数小时。如果你生成 token 的速度快一倍,就可能把智能体工作负载从 8 小时缩短到 4 小时,这是有意义的差别。

另一方面,批处理使工作负载更具成本效益。当 Nvidia 将 Groq 的 LPX 芯片集成到他们的 Vera Rubin NVL72 + Groq 3 LPX 混合集群 中时,他们只将 Groq 的芯片用于高交互性、高成本的工作负载。对于专注于最大化每瓦吞吐量(这是成本效益的代理指标)的工作负载,他们只使用 NVL72 机架,采用批处理方式运行。

要让 Tensordyne 更高效的乘法器产生显著的性能差异,它们需要在屋顶线模型右侧的计算受限区域运行。对于成本敏感的工作负载来说,这是一个绝佳的位置,但那些更注重速度而非效率的高成本智能体工作负载可能过于受内存限制,无法真正从 Tensordyne 的硬件中受益。然而,对数计算还带来另一个优势:物理上更小的乘法器。

如果 Tensordyne 的计算引擎更高效,这对内存受限的工作负载来说并不是特别重要。但如果他们的计算引擎更小,那也允许他们对内存层次结构做出有意义的改变。他们通过更小的乘法器节省出的空间,允许他们在每个裸片上容纳更多 SRAM 内存。

由此,性能优势就变得显而易见了。更小的计算引擎带来更多的片上内存,这减少了相同工作负载所需的片外内存访问次数,并使芯片更高效。

不过,这能让 Tensordyne 的芯片高效多少尚不清楚。这肯定不如他们在计算受限的批处理工作负载上获得的优势那么大。要估算实际优势,我需要更多关于芯片本身的技术规格。

归根结底,Tensordyne 的芯片提供了一个明确的权衡。你必须想办法让模型用他们特殊的数学在他们的特殊芯片上运行,这将把它们限制在推理工作负载上。让模型与对数数学良好配合很可能既困难又会降低性能,即使他们有特殊的编译器、出色的支持以及 AI 智能体的帮助。但作为回报,你在计算受限、高延迟的推理工作负载上获得好得多的效率,在内存受限、低延迟的推理工作负载上获得稍好的效率。我们将看看市场是否认为这值得付出努力。无论如何,我认为看到新的、奇特的数字格式真正进入商用硅片是非常令人兴奋和酷的事情。

感谢阅读 zach 的科技博客!这篇文章是公开的,所以请随意分享。

分享

没有文章

来源:Hacker News · zach.be