超越掩码稀疏：SNACK 在 GPU 上实现真实稀疏神经网络

摘要

深度神经网络的参数数量持续增长，推高了训练和 GPU 的推理成本。稀疏神经网络和动态稀疏训练 (DST) 有望降低这些成本，但大多数实现依赖于密集张量上的二进制掩码，并且几乎无法恢复理论计算、内存或节能。我们提出 SNACK，一个真正稀疏的 GPU 层，仅存储和计算非零连接。 SNACK 公开了一个简单的 PyTorch API，用于完全在稀疏范式中重组连接和反向传播梯度，并提供 SNACK-COO，这是一个自定义 COO 格式的 SpMM CUDA 内核，具有批处理到流多处理器映射，针对大型模型训练和单流推理的典型小批量、高稀疏性机制进行了调整。在内核级别，SNACK 比屏蔽密集基线 (Dense+Mask) 最高快 7 倍，并且在 95% 的稀疏度下可与 cuSPARSE、Sputnik 和 FlashSparse 竞争。在稀疏度为 90% 的情况下，单个 SNACK 层比 Dense+Mask 层和全密集层分别将训练加速 8 倍和 3.7 倍，推理速度分别提高 4 倍和 2 倍，同时内存的使用量比密集层少 72%，能量也大大减少。端到端，SNACK 在稀疏度为 99% 的情况下，与 Dense+Mask 相比，SNACK 将 GPT-2 峰值训练内存降低了高达 40%，并且计算图式推理延迟降低了 4.8 倍。

超越掩码稀疏：SNACK 在 GPU 上实现真实稀疏神经网络的原论文方法或结果图 — 图 1：SNACK 层概述。两个堆叠的 SNACK 块（$A$ 和 $B$）充当标准 nn.Linear 层的直接替代品。每个块将其权重矩阵 $\mathrm{SpM}(\cdot)$ 存储在 GPU 上，作为仅在其非零值上的 COO 三元组 $(\mathbf{A}_{v},\mathbf{A}_{r},\mathbf{A}_{c})$：突出显示的（蓝色）边缘是消耗记忆和 FLOP 的活动稀疏子网络，而褪色的（灰色）边缘描绘了 SNACK 在训练或推理期间从未实现的底层密集线性层。前向传递由 SNACK-COO 直接在三元组上计算，后向传递使用 $\mathrm{supp}(W)$ 上的稀疏外积，DST 方法可以就地添加或删除活动连接，因此 SNACK 和密集块可以在网络中自由交错，而无需更改任何调用者代码。