论文

RED-PIM:用存内计算减少Transformer数据搬移

RED-PIM: Reducing Data Movement for Transformers using Processing-in-Memory

摘要

Transformer 广泛应用于许多领域,包括自然语言处理、计算机视觉、网络搜索和 DNA 序列分析。鉴于其广泛的适用性,提高 Transformer 模型的性能至关重要。然而,注意力操作期间处理单元和内存之间的大量数据移动极大地限制了它们的效率。内存处理 (PIM) 通过直接在内存中执行计算来缓解此问题。虽然之前的工作提出了基于 PIM 的 Transformer 实现,但它们面临着昂贵的存储体间通信成本,并且由于存储体容量有限而难以扩展。因此,与注意力相关的数据必须在各个存储体之间进行分割,从而削弱了 PIM 的潜在优势。在这项工作中,我们提出了 RED-PIM,一种算法架构协同设计,通过最小化存储体间数据移动从 O(N^2) 到 O(N) 并将中间注意力矩阵从 N x N 缩小到 d x d 来减少注意力延迟。通过重新组织矩阵运算、在本地执行计算以及采用优化的数据传输策略,RED-PIM 显着降低了计算成本和互连流量。与基线 PIM 实现相比,RED-PIM 的推理时间减少了 16.05% 至 99.99%(几何平均值为 66.42%),其中较长序列的增益最大。在真实数据集上,RED-PIM 将长文档的性能提高了 99.60%,将较短文档的性能提高了 13.44%,同时保持或提高了准确性。这些结果证明了 RED-PIM 在可扩展且高效的 Transformer 推理方面的有效性。