作者：Meta FAIR：Mostafa Elhoushi，Meta Reality Labs Research 传感器与系统团队：Syed Shakib Sarwar, Aaryan Kothapalli, Mia Kasperek, Barbara De Salvo，Meta PyTorch 团队：Christian Puhrsch, Jesse Cai, Joe Isaacson，Quantsight：Andrew James, Pearu Peterson, Nikita Vedeneev

TLDR (太长不看): 通过对 MLP 模块的权重应用块稀疏性，我们在 A100 GPU 上使用 float32 Vision Transformer (ViT) 取得了喜人的结果，实现了高达 1.46 倍的加速，且精度下降幅度小于 2%。这种方法有可能应用于包括大型语言模型在内的其他类型的 Transformer。我们的实现和用于复现结果的基准测试代码可在 https://github.com/pytorch-labs/superblock 获取。

引言

PyTorch 在实现块稀疏矩阵乘法的 CUDA 内核方面进行了大量改进。PyTorch 的最新更新可以在具有高稀疏度的大型矩阵乘法形状上，相比密集基线实现高达 4.8 倍的加速。

在本篇博客中，我们将展示在 Vision Transformer (ViT) 的 MLP（多层感知机）层中的线性层权重上应用块稀疏性所获得的喜人结果，并展示在 A100 Nvidia GPU 上的端到端模型加速效果。

回顾一下，块稀疏性是以预定大小的块为单位对权重进行稀疏化，而不是对单个元素进行稀疏化。这种特定的稀疏模式很有意义，因为它可以通过快速稀疏内核进行 GPU 加速。有关不同稀疏模式之间的差异或稀疏性整体的更多信息，请查阅 torchao。

Illustrations of different types of sparsity.

不同类型稀疏性的图示。

方法

我们的方法可以分解为两个不同的步骤

  1. 使用块稀疏掩码子网络从头开始训练模型。
  2. 将这些掩码折叠到我们的权重中以加速推理。

我们将在下面解释训练和推理步骤

训练

我们从一个未初始化的 Vision Transformer 开始，在注意力块的输出投影线性层、MLP（即 FFN，前馈网络）内部的两个线性层以及最终的线性分类层的权重上，应用具有指定块大小和稀疏度的随机可训练掩码。训练期间的前向传播遵循 supermask 方法，因为每个掩码都会根据稀疏度要求使用调整后的阈值转换为二值图，例如，如果我们需要 80% 的稀疏度，阈值将自动调整以保留前 20% 的权重。掩码是边长为 <block size>x<block size> 的正方形元素，其中 <block size> 是一个超参数。权重的优先级取决于训练得到的掩码值或分数。我们 将每一层的二值掩码与权重相乘 来稀疏化模型。

Illustration of the Supermask sparsification approach

Supermask 稀疏化方法的图示。

推理

训练后，可以通过 将密集权重与掩码相乘 转换为稀疏权重并存储用于推理。在此阶段，尽管权重具有很高比例的零值，但它们仍然以密集格式存储。我们使用 PyTorch 的 to_sparse_bsr() API 将权重转换为 块稀疏表示 (BSR) 格式，该格式仅存储非零值及其块的索引。此步骤只需执行一次，结果可以在运行时缓存。

在运行时，无需更改代码。我们只需将任何输入张量传递给模型，当稀疏化线性层的 forward() 函数被调用时，PyTorch 会负责调用针对块稀疏权重的优化矩阵乘法。这应该适用于 A100 和 H100 NVIDIA GPU。

结果：微基准测试

为了从性能角度验证块稀疏性的可行性，我们首先使用此 简单脚本 运行了一系列微基准测试。我们使用 ViT-b 的线性层形状，在单个线性层上比较了我们块稀疏内核的加速效果，同时改变了权重矩阵的稀疏度和块大小。

我们在 NVIDIA A100 上使用 PyTorch 2.3.0.dev20240305+cu121 nightly 版本运行，并报告了每种稀疏配置相对于密集基线的加速效果。我们观察到，对于 float32，当块大小 >=32 或稀疏度 >= 0.8 时，会获得正向加速；而对于 bfloat16，我们观察到较小的加速效果，通常出现在块大小为 64 及更高的稀疏度下。因此，在本篇博客中，我们将重点关注 float32 的端到端模型加速，bfloat16 将留待未来工作。

Micro benchmarking results on linear layers of ViT-b-16.

Micro benchmarking results on linear layers of ViT-b-16.

ViT-b-16 线性层上的微基准测试结果。

结果：Vision Transformer

一旦我们确认能够在线性层上实现加速，我们便着手展示在 ViT_B_16 上的端到端加速效果。

我们使用标准的 ViT_B_16 方案 在 ImageNet 数据集上从头训练了该模型。我们展示了对 MLP 模块进行稀疏化的加速效果，对注意力输入和输出投影权重进行稀疏化将留待未来工作。

我们考察了实际推理加速（wall-clock speedup），重点关注批处理大小为 256 的情况。我们发现：

  • 对于 90% 的稀疏度，块大小分别为 16、32 和 64 时，我们可以获得 1.24 倍、1.37 倍和 1.65 倍的加速。
  • 为了获得加速，块大小为 16、32 和 64 时所需的最小稀疏度分别为 0.86、0.82 和 0.7。因此，正如预期，块大小越大，获得加速所需的稀疏度越小。

我们注意到 sparse_bsr() API 的一个限制：层维度需要是块大小的倍数。由于 ViT 中最后一个 FC 分类层的维度不是块大小的倍数，因此在我们的实验中它们未被转换为 BSR 表示。

Speedup on ViT-b-16 with batch size 256 on MLP modules across different batch sparsities and block sizes.

批处理大小为 256 时，ViT-b-16 在 MLP 模块上不同稀疏度等级和块大小下的加速效果。

我们还探索了 90% 稀疏度下不同批处理大小的加速效果。我们观察到，从批处理大小 16 开始及以上的批处理大小均实现了相对于基线的加速。虽然较大的块大小在最大批处理大小时能获得更大的加速，但对于较小的块大小，获得 >1 加速的最小批处理大小则更小。

我们认为设备端硬件在批处理大小为 1 时也能获得加速，因为它们——与服务器 GPU 不同——在这种小批处理大小时可以得到充分利用。

Speedup on ViT-b-16 with 90% sparsity on MLP modules across different batch sizes and block sizes.

90% 稀疏度下，ViT-b-16 在 MLP 模块上不同批处理大小和块大小下的加速效果。

考察稀疏化模型在 ImageNet-blurred 测试集上不同块大小和稀疏度下的 Top-1 精度，我们看到了一些预期结果：

  • 低稀疏度（<=70%）对精度没有显著影响
  • 中等稀疏度（>=80% 到 <90%）对精度的影响有限
  • 高稀疏度（>=90%）移除了大量权重，导致精度受到显著影响

可以进行更多研究来提高更高稀疏度和更大块大小下的精度。我们希望 PyTorch 对块稀疏性的支持以及本博客中展示的加速效果能够鼓励研究人员探索更精确的稀疏化方法。

Accuracies on training ViT-b-16 on ImageNet-blurred using the SuperMask approach.

使用 SuperMask 方法在 ImageNet-blurred 数据集上训练 ViT-b-16 的精度结果。

下一步

我们展示了在 float32 精度下对 ViT 的 MLP 模块进行块稀疏化的喜人加速效果。要在 bfloat16 上实现加速还有更多工作要做，我们希望能尽快取得进展。进一步优化 Vision Transformer 以及更广泛的 Transformer 上块稀疏性的可能下一步包括：

  • 对注意力的输入和输出投影执行块稀疏化。
  • 在微调期间而非从头训练期间执行块稀疏化。
  • 对适用于 ViT 线性算子特定形状的 matmul 内核进行进一步优化（特别是对于 80% 及更低的稀疏度）。
  • 与 int8 量化和其他优化（如 torch.compile()）结合使用
  • 探索其他权重稀疏化算法，例如 Spartan，以提高精度
  • 探索选择要稀疏化的权重（例如，特定的 Transformer 层）

如果您有问题或有兴趣为块稀疏化做出贡献，请联系 melhoushi@meta.com！

此外，如果您对稀疏性广泛感兴趣，请随时联系 @jcaip / jessecai@meta.com，并请查看 torchao，这是一个我们正在构建的社区，专注于量化和稀疏性等架构优化技术。