arXiv:2505.23949cs.LGcs.AI2025-05NeurIPS

提出高效算法生成可转置的N:M稀疏掩码,支持大模型训练加速。

TSENOR: Highly-Efficient Algorithm for Finding Transposable N:M Sparse Masks

  • 将掩码生成建模为带熵正则化的最优传输问题,用狄克斯特拉算法求解。
  • 在百亿参数模型上实现100倍加速,误差仅1-10%,支持任意N:M组合。
  • 可集成至Wanda、SparseGPT等框架,适用于大模型训练与推理加速。

网络剪枝能降低大模型的计算开销,其中N:M稀疏性(每M个连续权重保留N个)在压缩效果与硬件加速间取得良好平衡。然而,传统N:M稀疏性在矩阵转置时无法保持结构,限制了训练阶段的效率。虽已有可转置N:M稀疏方法,但现有方案难以扩展至大模型,或仅支持M=4,导致压缩-精度权衡不佳。本文提出一种高效的可转置N:M掩码求解器,可扩展至百亿参数模型。通过将掩码生成建模为最优传输问题,结合熵正则化与狄克斯特拉算法求解,并辅以取整步骤。基于张量的实现充分利用GPU并行性,相比现有方法最高提速100倍,误差仅1%-10%。该方法可与层级剪枝框架如Wanda、SparseGPT和ALPS集成,生成任意N:M的可转置稀疏模型。实验表明,采用可转置16:32稀疏的LLaMA3.2-8B模型性能接近标准N:M版本,且优于标准2:4稀疏模型,验证了其实际价值。

原文摘要 · Abstract (English)

Network pruning reduces the computational requirements of large neural networks, with N:M sparsity -- retaining only N out of every M consecutive weights -- offering a compelling balance between compressed model quality and hardware acceleration. However, N:M sparsity only accelerates forward-pass computations, as N:M patterns are not preserved during matrix transposition, limiting efficiency during training where both passes are computationally intensive. While transposable N:M sparsity has been proposed to address this limitation, existing methods for finding transposable N:M sparse masks either fail to scale to large models or are restricted to M=4 which results in suboptimal compression-accuracy trade-off. We introduce an efficient solver for transposable N:M masks that scales to billion-parameter models. We formulate mask generation as optimal transport problems and solve through entropy regularization and Dykstra's algorithm, followed by a rounding procedure. Our tensor-based implementation exploits GPU parallelism, achieving up to 100x speedup with only 1-10% error compared to existing methods. Our approach can be integrated with layer-wise N:M pruning frameworks including Wanda, SparseGPT and ALPS to produce transposable N:M sparse models with arbitrary N:M values. Experiments show that LLaMA3.2-8B with transposable 16:32 sparsity maintains performance close to its standard N:M counterpart and outperforms standard 2:4 sparse model, showing the practical value of our approach.

模型剪枝稀疏性高效算法大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。