Accelerating Transformer Inference and Training with 2:4 Activation Sparsity

📄 arXiv: 2503.16672v1 📥 PDF

作者: Daniel Haziza, Timothy Chou, Dhruv Choudhary, Luca Wehrstedt, Francisco Massa, Jiecao Yu, Geonhwa Jeong, Supriya Rao, Patrick Labatut, Jesse Cai

分类: cs.LG, cs.AI

发布日期: 2025-03-20

💡 一句话要点

利用2:4激活稀疏性加速Transformer的推理与训练

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Transformer 稀疏性 2:4稀疏 激活函数 模型加速 语言模型 硬件加速

📋 核心要点

  1. 大型语言模型的训练和推理计算成本高昂，现有加速方法仍有提升空间。
  2. 该论文利用Squared-ReLU激活的固有稀疏性，结合硬件加速的2:4稀疏模式，实现加速。
  3. 实验表明，该方法在不损失精度的情况下，前馈网络速度提升高达1.3倍。

📝 摘要（中文）

本文展示了如何利用2:4稀疏性（一种流行的硬件加速GPU稀疏模式）来加速大型语言模型的训练和推理。关键在于，我们利用了Squared-ReLU激活中固有的稀疏性，从而在不损失准确性的前提下实现加速。我们的方法在正向和反向传播中，使前馈网络（FFN）的速度提高了1.3倍。这项工作突出了稀疏性在加速大型语言模型训练和推理方面的潜力。

🔬 方法详解

问题定义：大型语言模型的训练和推理需要大量的计算资源。现有的加速方法，例如模型压缩、量化等，虽然可以降低计算量，但可能会导致精度损失，或者硬件支持不够完善。因此，如何在保证精度的前提下，进一步加速大型语言模型的训练和推理是一个重要的研究问题。

核心思路：该论文的核心思路是利用激活函数中的稀疏性。具体来说，Squared-ReLU激活函数会产生大量的零值激活，这些零值激活可以被有效地利用来减少计算量。同时，论文利用硬件加速的2:4稀疏模式，进一步提升计算效率。2:4稀疏模式是指在每4个元素中，只有2个非零元素，这种模式可以被特定的硬件加速器高效地处理。

技术框架：该方法主要针对Transformer模型中的前馈网络（FFN）进行优化。FFN是Transformer模型中计算量最大的模块之一。该方法通过在FFN中使用Squared-ReLU激活函数，并结合2:4稀疏模式，来加速FFN的计算。整体流程包括：1. 使用Squared-ReLU激活函数；2. 将激活值转换为2:4稀疏格式；3. 利用硬件加速器进行计算。

关键创新：该论文的关键创新在于将Squared-ReLU激活的固有稀疏性与硬件加速的2:4稀疏模式相结合。以往的研究可能只关注激活函数的稀疏性，而忽略了硬件加速的潜力。该论文通过充分利用硬件加速的优势，实现了更高的加速效果。此外，该方法在不损失精度的情况下实现了加速，这对于大型语言模型的应用至关重要。

关键设计：论文的关键设计包括：1. 选择Squared-ReLU作为激活函数，因为它具有较高的稀疏性；2. 使用2:4稀疏模式，以便利用硬件加速器；3. 针对FFN进行优化，因为FFN是Transformer模型中计算量最大的模块。具体的参数设置和网络结构与原始的Transformer模型保持一致，只在激活函数和稀疏模式上进行了修改。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，该方法在不损失精度的情况下，使Transformer模型中的前馈网络（FFN）速度提升高达1.3倍。这一显著的加速效果证明了该方法在加速大型语言模型训练和推理方面的有效性。该方法具有很强的实用价值，可以很容易地集成到现有的Transformer模型中。

🎯 应用场景

该研究成果可广泛应用于各种需要加速大型语言模型训练和推理的场景，例如自然语言处理、机器翻译、文本生成、对话系统等。通过降低计算成本，可以使得大型语言模型更容易部署在资源受限的设备上，并促进其在各个领域的应用。

📄 摘要（原文）

In this paper, we demonstrate how to leverage 2:4 sparsity, a popular hardware-accelerated GPU sparsity pattern, to activations to accelerate large language model training and inference. Crucially we exploit the intrinsic sparsity found in Squared-ReLU activations to provide this acceleration with no accuracy loss. Our approach achieves up to 1.3x faster Feed Forward Network (FFNs) in both the forwards and backwards pass. This work highlights the potential for sparsity to play a key role in accelerating large language model training and inference.