IBM 全新低精度 AI 芯片实现 25.6 TFLOPS 训练与 102.4 TOPS 推理

在 2024 年 ISSCC 上，IBM Research 展示了一款 7 nm 原型芯片，凝聚多年在低精度 AI 训练和推理方面的研究成果。该芯片支持 16 位和 8 位训练，以及 4 位和 2 位推理——远低于业界常规的 32 位训练 / 8 位推理。

降低数值精度可减少计算和功耗需求，但 IBM 还引入多项架构优化进一步提升效率。核心挑战是在不损害精度的前提下降低精度——IBM 多年来从算法层面攻克此难题。

IBM AI 硬件中心成立于 2019 年，目标是每年将 AI 计算性能提升 2.5 倍，到 2029 年实现 FLOPS/W 提升 1000 倍。这一雄心源于 AI 模型的爆发式增长——尤其是万亿参数 NLP 系统——及其训练带来的碳足迹日益加重。

在原型芯片上，4 核设计在 1.6 GHz、0.75 V 下实现 8 位训练 25.6 TFLOPS，4 位整数推理达 102.4 TOPS。将频率降至 1 GHz、电压降至 0.55 V，能效分别提升至 3.5 TFLOPS/W (FP8) 和 16.5 TOPS/W (INT4)。

IBM Research 测试芯片性能（图片来源：IBM Research）**报告性能基于 0% 稀疏度。(1) FP8. (4) INT4.

低精度训练

这一性能建立在多年低精度训练和推理算法研究基础上。该芯片首创支持 IBM 在 NeurIPS 2019 推出的混合 FP8 格式，使 8 位训练在不妥协结果的前提下，将 16 位训练的计算量减半。

IBM Research 致力于解决降低精度同时保持精度的问题（图片来源：IBM）

“低精度训练极具挑战，但采用正确的数值格式，8 位训练是可行的，”IBM Fellow、加速器架构与机器学习高级经理 Kailash Gopalakrishnan 表示。“针对正确张量应用正确的数值格式至关重要。”

混合 FP8 结合两种格式：前向传播中权重和激活用一种，反向传播中梯度用另一种。训练需要双向传播，推理仅用前向格式。

“我们发现前向张量需要更高保真度——尾数精度更高；而梯度需要更宽动态范围——指数位更大。这种权衡导致前向采用 1-4-3（符号-指数-尾数）划分，反向采用 5 位指数，动态范围达 2^32。我们的硬件现支持两种变体。”

分层累加

IBM 的“分层累加”技术在累加过程中同步降低精度，配合权重和激活。传统 FP16 训练用 32 位算术累加以保精度，但 8 位方案在 FP16 中累加。通过将长累加分组为更小块并求和，芯片在保持低精度优势的同时最小化舍入误差。

“按顺序累加数千个数会降低精度，”Gopalakrishnan 解释。“分块累加既能保持求和准确，又能保持硬件精简。”

低精度推理

虽然大多数推理引擎用 INT8，IBM 证明 4 位整数运算可匹配最先进精度。量化后，量化感知训练将精度损失控制在 0.5% 以内，对许多应用可接受。

片上环网

芯片集成针对深度学习优化的片上网络，使核心能高效多播数据，减少片外内存流量和带宽需求。

“我们将环网工作频率与核心解耦，允许独立优化，提高核心速度而不受环网长线延迟影响，”机器学习与加速器架构研究员 Ankur Agrawal 表示。

功耗管理

IBM 引入频率缩放方案，使芯片在整个负载期间保持接近功耗包络。通过激进设定初始电压和频率，芯片在计算密集阶段维持高性能，在低需求时适应而无需昂贵的电压变更。

“深度学习负载有明显的功耗峰谷，”Agrawal 指出。“我们的策略让功耗保持稳定，将低谷转化为性能增益。”

测试芯片

4 核设计平衡互连复杂度与核心规模，提供可扩展架构。IBM 预见 1-2 核芯片用于边缘设备，32-64 核用于数据中心部署，均支持 FP16、混合 FP8、INT4 和 INT2。

“不同应用领域对精度和能效权衡需求不同，”Gopalakrishnan 说。“我们的‘瑞士军刀’式格式组合，让我们能针对各用例而不妥协效率。”

配合硬件，IBM Research 发布“Deep Tools”编译器栈，在芯片上实现 60-90% 利用率。

据早先采访，基于此架构的低精度 AI 芯片将在两年内面市。

>> 本文原载于姊妹站点 EE Times。

相关内容：

• AI 芯片通过模型缩减保持精度
• 边缘训练 AI 模型
• 边缘 AI 竞赛
• 边缘 AI 挑战存储技术
• 工程团体寻求将 1mW AI 推向边缘
• 小规模任务的神经网络应用
• AI IC 研究探索替代架构

更多 Embedded 内容，请订阅 Embedded 每周电子邮件通讯。