S-STE：用于高效2:4稀疏预训练的连续剪枝函数
机器学习
2024-12-30 v3
摘要
训练深度神经网络（DNN）成本高昂。幸运的是，Nvidia Ampere和Hopper GPU通过实现2:4稀疏性，可以比稠密等效运算快一倍地加速矩阵乘法。然而，先前基于STE的2:4预训练方法（例如，带硬阈值的STE、SR-STE）因剪枝函数不连续而遭遇优化困难。在本研究中，我们全面分析了传统N:M稀疏训练的瓶颈，并认识到不连续性带来的三个缺点：错误的下降方向、无法预测下降量以及稀疏掩码振荡。鉴于此，我们提出了S-STE，这是一种简单而强大的2:4训练方法，包含两个部分：连续地将权重投影为2:4稀疏，并使用一个逐张量的固定缩放因子重新缩放稀疏权重。此外，我们对激活梯度采用最小方差无偏估计，并对整个过程采用FP8量化。结果表明，我们的方法超越了先前的2:4预训练方案，甚至可与全参数模型相媲美。我们的工具包可在https://github.com/huyz2023/2by4-pretrain获取。
引用
@article{arxiv.2409.09099,
title = {S-STE: Continuous Pruning Function for Efficient 2:4 Sparse Pre-training},
author = {Yuezhou Hu and Jun Zhu and Jianfei Chen},
journal= {arXiv preprint arXiv:2409.09099},
year = {2024}
}