Accelerating Transformer Pre-training with 2:4 Sparsity
机器学习
2024-10-29 v3
摘要
训练大型 transformer 模型速度较慢，但最近的 GPU 架构创新为我们提供了优势。NVIDIA Ampere GPU 能够以细粒度 2:4 稀疏矩阵乘法的速度翻倍于其稠密等效操作。在此基础上，我们全面探讨了加速 transformer 前馈网络 (FFN) 在预训练中的可行性。首先，我们定义一种“翻转率”来监控 2:4 训练过程的稳定性。利用这一指标，我们提出了三种技术来保持准确性：通过在梯度上应用掩码衰减项来修改稀疏优化的直向估计器，以确定预热阶段中可行的衰减因子，并通过在预训练末期进行密集细化来提高模型质量。此外，我们设计了两种技术来实际加速训练：通过卷积计算可转置的 2:4 掩码，并通过减少 GPU L2 缓存丢失来加速门控激活函数。实验表明，我们的 2:4 稀疏训练算法在几个 transformer 预训练任务上实现与密集训练算法相似的收敛效果，而在不同形状的 transformer 块上实际加速显而易见。我们的工具包已在 https://github.com/huyz2023/2by4-pretrain 上提供。
引用
@article{arxiv.2404.01847,
title = {Accelerating Transformer Pre-training with 2:4 Sparsity},
author = {Yuezhou Hu and Kang Zhao and Weiyu Huang and Jianfei Chen and Jun Zhu},
journal= {arXiv preprint arXiv:2404.01847},
year = {2024}
}