博客社区

超越量化:为 PyTorch 带来稀疏推理

作为开发者,我们都很清楚:大语言模型(LLM)具有革命性,但其成本极其高昂。运行前沿模型需要配备大量能源消耗的专用 GPU 集群。多年来,社区一直依赖低精度量化和定制的混合精度算子来使这些模型具备实用性。但对于我们这些专注于边缘计算和端侧推理的人来说,仅仅做到这些还不够。我们需要下一个优化前沿,即稀疏性(Sparsity)。我们坚信,未来的方向是建立一个统一的稀疏推理框架,而我们正在 PyTorch 中实现它。

解锁大语言模型中的稀疏性

像 Meta 的 OPT 等使用 ReLU 激活函数的早期模型曾是“金矿”。研究表明,对于平均输入,其 MLP 模块中超过 95% 到 99% 的权重甚至根本没有被激活。

图 1: OPT 模型在注意力机制和 MLP 模块中观察到的稀疏性
(图片来自 Liu 等人 [2023])

因此,我们可以更新 MLP 模块中的 FFN 计算,跳过这些未激活的神经元,从而同时节省内存和计算资源。

为了提前找到稀疏索引,我们可以使用低秩预测器(low-rank predictors)来计算门控矩阵(gate matrix)的近似分解。通过将维度投影到隐藏层大小的 4-10%,可以实现极快的门控近似计算,同时保持精度。通过观察到 Transformer 的残差结构导致不同层之间的嵌入(embeddings)非常相似,我们可以进一步加速这一过程——这意味着我们可以使用第 i-1 层的输入异步计算第 i 层的预测器,从而最大程度地降低延迟。

当结合各种硬件优化和稀疏预测器层的异步前瞻执行(asynchronous look-ahead execution)时,这种方法(被称为“Deja Vu”)使推理速度提高了 2-6 倍,且几乎没有精度损失 [图 1]。该方法很快被其他研究人员采用,并经由《LLM in a Flash》[Alizadeh 等人, 2024]、《PowerInfer》[Song 等人, 2023] 和《PowerInfer2》[Xue 等人, 2024] 等论文进行了迭代。

图 2: 使用 Deja Vu 稀疏预测器的稀疏化 OPT 的下游精度。即使在极高的稀疏水平下,也观察到极小的性能下降。
(图片来自 Liu 等人 [2023])

现代大语言模型的“长尾”挑战

图 3: ReLU、SiLU 和 GeLU 的激活曲线图。

自 OPT 以来,更新的先进模型(包括 Llama、Mistral 和 Gemma)已将 ReLU 替换为 SiLU 和 GeLU 等更平滑的激活函数。这些函数不会将负输入硬截断为零;相反,它们具有平滑延伸过零点的“长尾”。如图 2 所示,这种变化显著降低了激活稀疏性。天真地对激活值进行阈值处理会导致严重的精度损失,且稀疏水平远低于 OPT 中所见的 >95%。

在现代大语言模型中寻找稀疏性

那么,激活稀疏性在现代时代已经终结了吗?完全没有。目前出现了两大主要学派来应对这一挑战。

1. Relufication:微调回 ReLU

最简单的想法是将 SiLU/GeLU 激活函数替换为 ReLU,然后通过微调模型来恢复性能。Mirzadeh 等人 [2023] 表明,在 150 亿个标记(tokens)上微调 7B 参数的 Llama 模型,可在 9 项零样本(zero-shot)任务中恢复 ≈60% 的稀疏性,同时仅牺牲约 1% 的精度。后续研究进一步提高了稀疏性:Song 等人提出了辅助损失函数和多阶段训练计划,在 Llama-2(7B 和 13B)上实现了 80–90% 的稀疏性,且精度损失可忽略不计。

Relufication 依然强大,且预训练的“ReluLlama”、“ProSparse”和“TurboSparse”检查点已经可用。然而,为每一个新模型构建“Relufied”版本成本高昂,且需要大规模的微调运行,限制了其实际应用价值。

2. 无训练的“误差预算”阈值处理:CATS 和 CETT

在无法重新训练的情况下,我们可以通过对激活值进行阈值处理来近似稀疏性。一种简单的方法称为上下文感知阈值稀疏性(Contextually Aware Thresholding Sparsity, CATS),它预先计算校准集上的激活范数,并丢弃范数低于第 p 百分位的神经元;这种方法在不进行微调的情况下,在 Llama-2 7B 和 Mistral 7B 上实现了约 50% 的稀疏性。然而,CATS 仅考虑了门控激活;乘以大权重的较小数值仍可能对输出做出重大贡献。

为了解决这个问题,Zhang 等人提出了尾部截断累积误差(Cumulative Errors of Tail Truncation, CETT)。对于每个神经元 i,CETT 计算该神经元的完整贡献 ni ——包括门控、上投影和下投影——然后选择一个阈值 τ,使得被忽略贡献的 L2 范数保持在目标误差预算内。固定一个“误差预算”(例如允许丢弃 20% 的输出范数)即可通过二分搜索获得与数据相关的 τ。

在我们的实验中,使用 0.2 的标准 CETT 预算,在多个现代模型上恢复了 >60% 的稀疏性

图 4: 使用 0.2 的 CETT 阈值时各层的稀疏度

在 PyTorch 中分页大模型权重

使用 CETT 找到稀疏掩码只是挑战的一半:我们还需要高效地执行稀疏操作。一种简单的稀疏实现方案在每次前向传递时执行完整的 index_select (gather) 以加载激活的权重;这种操作受限于内存带宽瓶颈。我们观察到“神经元集中度”(neuron concentration),即一小部分神经元在连续的标记中保持“活跃”的趋势,这意味着许多权重在不同步骤间是被复用的。为什么每次都要重新加载它们?

图 5: 左图显示了简单的稀疏 MLP 实现,在每个步骤都有完整的索引选择(index select)操作。

右图显示了我们带有权重缓存的改进版本。

我们在 PyTorch 中构建了一个自定义的权重缓存算子,将先前活跃的权重保留在缓存中,仅将连续掩码之间的差异作为独立的索引交换(index swaps)进行加载。该算子以行优先格式存储上投影权重,以列优先格式存储下投影权重,以实现高效的内存复制。初步实验很有前景:该缓存实现将独立的 index_select 操作加速了 6.7 倍(29.89 毫秒 → 4.46 毫秒),并在结合融合稀疏核(fused sparse kernels)和 OpenMP 并行化后,在 CPU 上实现了高达 5 倍的 MLP 推理加速(30.1 毫秒 → 6.02 毫秒)。

图 6: MLP 块操作的整体加速比与每次缓存更新的索引交换次数的关系。右图显示了使用图 3 稀疏度得出的常见模型的平均索引交换次数。

在 LLM 训练中强制实施稀疏性

稀疏性正迅速从一种优化技巧转变为核心架构特征。DeepSeek v3.2 和 Google 的 Gemma 3-n 在训练时就已经在强制实施它。

DeepSeek 的闪电索引器(Lightning Indexer)

图 7: DeepSeek 稀疏注意力架构(图片来自 DeepSeek-AI [2025])

DeepSeek 的 v3.2 模型引入了 DeepSeek 稀疏注意力 (DSA),它使用一种称为闪电索引器(lightning indexer)的轻量级预测器。该索引器本质上是一个小型、ReLU 激活的网络,作用于低精度向量以估计注意力能量。通过从这些近似能量中选择 Top-k 条目,DSA 将 10 万个标记的上下文窗口缩小到固定的 2048 个标记切片进行完整的注意力计算。结果是在长上下文场景下实现了显著更快的推理。尽管 DeepSeek 在庞大的语料库上训练了该机制,但使用小型预测器修剪注意力窗口的理念具有广泛的适用性。

Google 的 Spark Transformer

Spark Transformer 为注意力层和前馈层提出了不同的稀疏预测器。它将输入向量划分为两个部分:仅用于计算 Top-k 掩码的一小段前缀,以及用于完整计算的其余部分。他们的方法包括定义稀疏算子“SparkAttention”和“SparkFFN”,其功能与我们之前讨论的轻量级稀疏预测器有些类似。

虽然完整的 SparkTransformer 需要大规模的预训练,但其第二项贡献——统计 Top-k(statistical top-k),具有即时应用价值。它不通过排序激活值来寻找阈值,而是将激活值建模为近似正态分布,并使用均值和标准差计算阈值。

θ(x,k)=μ(x)+n(k)·σ(x)

这消除了昂贵的排序需求,使得 Top-k 选择在 GPU 上运行快得多。该技术的版本已经出现在 Google 的 Gemma 3-n 模型中。

迈向统一的稀疏推理框架

稀疏性正在从一种研究层面的好奇心转变为生产环境的必需品。上下文稀疏性和权重缓存所展现的 2–6 倍加速并非增量式的调整;它们决定了端侧大语言模型部署的可行性。然而,现有的实现大多孤立地探索单个技巧。真正的挑战在于系统性的集成。

在 NimbleEdge,我们致力于结合各种优秀理念——预测性掩码、权重缓存、统计 Top-k 和硬件感知算子——构建一个统一的稀疏推理框架。我们的存储库是开放的,我们会在组件和基准测试成熟时陆续发布。你可以在此处找到它:

https://github.com/NimbleEdge/sparse_transformers

我们目前的研究重点包括:

权重缓存验证

确认各模型间的神经元集中模式,并将我们的 6.7 倍索引选择加速转化为端到端的性能增益。

CETT 集成

结合“Relufied”模型和基于 CETT 的阈值处理,以恢复长尾激活中的稀疏性。

融合稀疏核(Fused sparse kernels)

开发能够同时执行稀疏预测和缓存权重访问的核函数。

轻量级注意力索引器

探索类似于 DeepSeek 闪电索引器的长上下文注意力索引器,用于边缘部署。

大语言模型将持续增长,节能推理仍将是核心挑战。稀疏性提供了一条前进道路。通过重温激活函数(Relufication)、采用更智能的阈值方法(CETT)、利用神经元集中特性(权重缓存)以及借用前沿模型的创新(闪电索引器和统计 Top-k),我们可以实现端侧部署所需的多倍速度提升。但构建生产级的稀疏推理栈需要开放合作。我们邀请研究人员和实践者为这项努力做出贡献:边缘 AI 的未来取决于让稀疏推理成为标准,而不是例外。

如需深入了解稀疏性和最新进展,请参阅 NimbleEdge 团队的白皮书:使用稀疏性加速 LLM 推理