NVIDIA Ampere 架构的布局化稀疏功能及其在搜刮引擎中的应用
NVIDIA Ampere 架构的布局化稀疏功能及其在搜刮引擎中的应用深度学习彻底改变了我们分析、理解和处理数据的方式，而且在各个领域的应用中都取得了巨大的乐成，其在盘算机视觉、自然语言处理、医疗诊断和医疗保健、自动驾驶汽车、推荐系统以及气候和气候建模方面有许多乐成案例。
在神经网络模子不绝变大的时代，对盘算速率的高需求对硬件和软件都形成了巨大的挑战。模子剪枝和低精度推理是非常有效的办理方案。
自 NVIDIA Ampere 架构开始， 随着 A100 Tensor Core GPU 的推出，NVIDIA GPU 提供了可用于加速推理的细粒度布局化稀疏功能。在本文中，我们将先容此类稀疏模子的训练方法以保持模子精度，包罗基本训练方法、渐进式训练方法以及与 int8 量化的结合。我们还将先容如何利用 Ampere 架构的布局化稀疏功能进行推理。
腾讯呆板学习平台部门 (MLPD) 利用了渐进式训练方法，简化了稀疏模子训练并实现了更高的模子精度。借助稀疏功能和量化技能，他们在腾讯的离线服务中实现了 1.3 倍～1.8 倍的加速。
NVIDIA Ampere 架构的布局化稀疏功能
NVIDIA Ampere 和 NVIDIA Hopper 架构 GPU 增长了新的细粒度布局化稀疏功能，该功能主要用于加速推理。此功能是由稀疏 Tensor Core 提供，这些稀疏 Tensor Core 需要 2:4 的稀疏模式。也就是说，以 4 个相邻权重为一组，其中至少有 2 个权重必须为 0，即 50% 的稀疏率。
这种稀疏模式可实现高效的内存访问能力，有效的模子推理加速，并可轻松规复模子精度。在模子压缩后，存储格式只存储非零值和相应的索引元数据（图 1）。稀疏 Tensor Core 在实行矩阵乘法时仅处理非零值，理论上，盘算吞吐量是划一稠密矩阵乘法的 2 倍。
https://i-blog.csdnimg.cn/blog_migrate/faec3c0765d4ef716d0a31ae17e9b8db.jpeg
图 1. 2:4 布局化稀疏模式及其压缩格式
（布局化稀疏矩阵具有 2:4 的稀疏模式。在 4 个相邻权重当中，至少有 2 个值为零。在模子压缩后，仅存储非零值和相应的索引元数据。）
布局化稀疏功能主要应用于可以或许提供 2:4 稀疏权重的全连接层和卷积层。如果提前对这些层的权重做剪枝，则这些层可以使用布局化稀疏功能来进行加速。
训练方法
由于直接对权重做剪枝会降低模子精度，因此在使用布局化稀疏功能的时间，您需要进行训练来规复模子精度。下面，我们将先容一些基本训练方法和新的渐进式训练方法。
基本训练方法
基本训练方法可保持模子精度，并且无需任何超参数调优。相识更多技能细节，请参阅论文 Accelerating Sparse Deep Neural Networks（https://arxiv.org/abs/2104.08378）。
基本训练方法易于使用，步调如下：
训练一个常规稠密模子，不需要稀疏化的特殊处理。
对全连接层和卷积层上的权重以 2:4 的稀疏模式进行剪枝。
按照以下规则重新训练经过剪枝的模子：
a. 将所有权重初始化为第 2 步中的值。
b. 使用与第 1 步相同的优化器和超参数（学习率、调理方法、训练次数等）进行稀疏调优训练。
c. 保持第 2 步中剪枝后的稀疏模式。
https://i-blog.csdnimg.cn/blog_migrate/cd278e643757952fef4bdfee95cc2832.jpeg
图 2. 基本训练方法
（基本训练方法就是使用剪枝后的权重和掩码后的优化器重复原始稠密模子的训练过程。）
对于复杂环境，还有一些进阶的方法。
例如，把稀疏训练应用在多阶段式的稠密模子训练当中。好比对于一些目标检测模子，如果下游任务的数据集足够大，您只需做稀疏调优训练。对于像 BERT-SQuAD 等模子，调优阶段使用的数据集相对较小，您则需要在预训练阶段进行稀疏训练以得到更好的模子精度。
别的，通过在稀疏调优之前插入量化节点，您可以轻松将稀疏调优与 int8 量化调优结合起来。所有这些训练以及调优方法都是一次性的，即最终得到的模子只需要经过一次稀疏训练处理。
渐进式稀疏训练方法
一次性稀疏调优（fine-tuning）可以覆盖大多数任务，并在不损失精度的环境下实现加速。然而，就一些对权重数值变革敏感的困难任务而言，对所有权重做一次性稀疏训练会导致大量信息损失。在小型数据集上只做稀疏化调优大概会很难规复精度，对于这些任务而言，就需要稀疏预训练（pretraining）。
然而稀疏预训练需要更多数据，而且更加耗时。因此，受到卷积神经网络剪枝方法的启发，我们引入了渐进式稀疏训练方法，在此类任务上仅应用稀疏化调优便可以实现模子的稀疏化，同时不会造成明显的精度损失。相识更多细节，请参阅论文 Learning both Weights and Connections for Efficient Neural Networks （https://arxiv.org/pdf/1506.02626.pdf）。
https://i-blog.csdnimg.cn/blog_migrate/5a834daf437171dc71a41cf0325546ba.jpeg
图 3. 渐进式稀疏训练的概念
（渐进式稀疏训练方法将稀疏率分为几个步调，以轻松规复精度。渐进式稀疏训练方法的核心思想是将目标稀疏率进行多少次切分。）
https://i-blog.csdnimg.cn/blog_migrate/be7719409e75c5435a17a8fde508aea1.jpeg
如上述公式和图 4 所示，对于目标稀疏率 S，我们将其分为 N 份，这将有助于在稀疏调优过程中快速规复信息。根据我们的实验，在相同的调优迭代次数内，使用渐进式稀疏训练相比一次性稀疏训练，可以得到更高的模子精度。
https://i-blog.csdnimg.cn/blog_migrate/cc7596f8f446a1fd31b4e92c05f7498b.jpeg
图 4. 渐进式稀疏训练方法 （以 50% 稀疏率的 2:4 布局化稀疏模式为例）
（渐进式稀疏训练方法的示例：盘算权重掩码以到达 25% 稀疏率，再进行稀疏调优规复性能，末了重新盘算掩码使之到达 50% 稀疏率并对网络进行调优。）
我们以 50% 稀疏率的 2:4 布局化稀疏为例，将稀疏率分为两份，然后逐步稀疏和调优模子中的权重参数。如图 4 所示，首先盘算权重掩码以实现 25% 的稀疏率，然后实行稀疏调优以规复模子精度。接下来，重新对剩余权重盘算权重掩码以到达 50% 的稀疏率，并对网络进行调优，以得到一个精度无损的稀疏模子。
Sparse-QAT：稀疏化与量化、蒸馏相结合
为了得到更轻量的模子，我们进一步将稀疏与量化、蒸馏相结合，即 Sparse-QAT。
量化（PTQ 和 QAT）
下方的公式表现一个通用的量化过程。对于 32 位浮点数值 x，我们使用 Q 表现其具有 K-bits 表现的量化值。
https://i-blog.csdnimg.cn/blog_migrate/bf77d56b158b90a839823101fb7ef3db.jpeg
通常环境下，我们首先将原始参数量化到特定范围，并将其近似为整数。然后，可以使用这个量化比例 scale (s) 来规复原始值。这样就得到了第一种量化方法，即校准，也称为训练后量化（post-training quantization, PTQ）。在校准中，一个关键的因素是要设置一个得当的量化比例（scale）。如果这个比例值过大，量化范围内的数字将不太准确。相反，如果这个比例值过小，会导致大量的数字落在 lmin 到 lmax 的范围之外。因此，为了均衡这两个方面，我们首先得到张量中数值的统计分布，然后设置量化比例以覆盖 99.99% 的数值。许多工作已经证明，这种方法对于在校准过程中找到合适的量化比例非常有帮助。
然而，只管我们已经为校准设置了一个合理的量化比例，但是对于 8 bit 量化来说，模子精度仍然会明显降落。因此，我们引入量化感知训练（quantization-aware training, QAT），以进一步进步校准后的精度。QAT 的核心思想是以模拟量化的方法来训练模子。
在前向流传过程中，我们将权重量化为 int8，然后将其反量化为浮点数来模拟真实量化。在反向流传过程中，引入 straight through estimation (STE) 的方法来更新模子权重。STE 的核心思想可以用如下公式表现：
https://i-blog.csdnimg.cn/blog_migrate/2ccd9aebfa99b40e7645a201b00e2e81.jpeg
由上述公式可知，阈值范围内的值对应的梯度直接反向流传，超出阈值范围的值对应的梯度被裁剪为 0。
知识蒸馏
除了上述方法外，我们还引入了知识蒸馏（knowledge distilation, KD），以进一步确保 Sparse-QAT 模子的精度。我们以原始稠密模子作为教师模子，以量化稀疏模子作为学生模子。在调优过程中，我们采用了 Mini-distillation，这是一种层级别的蒸馏方法。使用 MiniLM，我们只需要使用 Transformer 模子末了一层的输出。引入蒸馏作为辅助工具乃至可以得到比教师模子精度更高的稀疏量化学生模子。相识更多信息，请参阅 MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers（https://arxiv.org/abs/2002.10957）。
Sparse-QAT 训练流水线
图 5 显示了 Sparse-QAT 的训练流水线。稀疏化、量化、蒸馏以并行的方式实行，最终得到一个稀疏的 int8 量化模子。整个流水线包罗如下三条路径：
在稀疏路径中，应用渐进式稀疏化来获取一个稀疏权重张量。
在量化路径中，使用 PTQ 和 QAT 来获取 int8 类型的权重张量。
在知识蒸馏路径中，使用 MiniLM 来进一步保障最终稀疏 int8 模子的精度。
https://i-blog.csdnimg.cn/blog_migrate/dbc4a263b792a0d6ae5b0e3c9794d1ba.jpeg
图 5. Sparse-QAT 流水线
（将稀疏、量化和知识蒸馏相结合，以得到最终的稀疏 int8 模子。）
使用 NVIDIA Ampere 架构的
布局化稀疏功能进行推理
在训练好稀疏模子后，您可以使用 NVIDIA TensorRT 和 cuSPARSELt 库来加速基于 NVIDIA Ampere 架构布局化稀疏功能的推理。
使用 NVIDIA TensorRT 进行推理
自 8.0 版本开始，TensorRT 可以支持稀疏卷积，矩阵乘法 (GEMM) 需要用 1x1 的卷积替代来进行稀疏化推理。在 TensorRT 中启用稀疏化推理非常简单。在导入 TensorRT 之前，模子的权重应具有 2:4 的稀疏模式。如果使用 trtexec 构建引擎，只需设置 -sparity=enable 标志即可。如果您正在编写代码或脚本来构建引擎，只需按如下所示设置构建立置：
对于 C++：
config->setFlag(BuilderFlag::kSPARSE_WEIGHTS)
对于 Python：
config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS)
使用 NVIDIA cuSPARSELt 库加强 TensorRT
在某些用例中，TensorRT 大概因为输入尺寸不同而无法提供最佳性能。您可以使用 cuSPARSELt 进一步加速这些用例。
办理方案是使用 cuSPARSELt 编写 TensoRT 插件，我们可以为不同的输入尺寸初始化多个描述符以及多个 cuSPARSELt 稀疏矩阵乘法 plan，并根据输入尺寸选择合适的 plan。
假设您在实现 SpmmPluginDynamic 插件，该插件继续自 nvinfer1:: IPluginV2DynamicExt，您可以使用一个私有布局来存储这些 plan。
struct cusparseLtContext { cusparseLtHandle_t handle; std::vector plans; std::vector matAs, matBs, matCs; std::vector matmuls; std::vector alg_sels;}
TensorRT 插件应实现 configurePlugin方法，该方法会根据输入和输出类型及尺寸设置插件。您需要在这个函数当中初始化 cuSPARSELt 的相关布局。
cuSPARSELt 的输入尺寸有一些限制，应为 4、8 或 16 的倍数，具体取决于数据类型。在本文中，我们将其设置为 16 的倍数。相识该限制条件的相关信息，请查看此文档（https://docs.nvidia.com/cuda/cusparselt/functions.html#cusparseltdensedescriptorinit）。
for (int i = 0; i < size_num; ++i) { m = 16 * (i + 1); int alignment = 16; CHECK_CUSPARSE(cusparseLtStructuredDescriptorInit( &handle, &matBs, n, k, k, alignment, type, CUSPARSE_ORDER_ROW, CUSPARSELT_SPARSITY_50_PERCENT)) CHECK_CUSPARSE(cusparseLtDenseDescriptorInit( &handle, &matAs, m, k, k, alignment, type, CUSPARSE_ORDER_ROW)) CHECK_CUSPARSE(cusparseLtDenseDescriptorInit( &handle, &matCs, m, n, n, alignment, type, CUSPARSE_ORDER_ROW)) CHECK_CUSPARSE(cusparseLtMatmulDescriptorInit( &handle, &matmuls, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_TRANSPOSE, &matAs, &matBs, &matCs, &matCs, compute_type)) CHECK_CUSPARSE(cusparseLtMatmulAlgSelectionInit( &handle, &alg_sels, &matmuls, CUSPARSELT_MATMUL_ALG_DEFAULT)) int split_k = 1; CHECK_CUSPARSE(cusparseLtMatmulAlgSetAttribute( &handle, &alg_sels, CUSPARSELT_MATMUL_SPLIT_K, &split_k, sizeof(split_k))) int alg_id = 0; CHECK_CUSPARSE(cusparseLtMatmulAlgSetAttribute( &handle, &alg_sels, CUSPARSELT_MATMUL_ALG_CONFIG_ID, &alg_id, sizeof(alg_id))) size_t ws{0}; CHECK_CUSPARSE(cusparseLtMatmulPlanInit(&handle, &plans, &matmuls, &alg_sels, ws)) CHECK_CUSPARSE( cusparseLtMatmulGetWorkspace(&handle, &plans, &ws)) workspace_size = std::max(workspace_size, ws);}
在 enqueue 函数中，您可以检索得当的 plan 来实行矩阵乘法。
int m = inputDesc->dims.d;int idx = (m + 15) / 16 - 1;float alpha = 1.0f;float beta = 0.0f;auto input = static_cast(inputs);auto output = static_cast(outputs);cusparseStatus_t status = cusparseLtMatmul( &handle, &plans, &alpha, input, weight_compressed, &beta, output, output, workSpace, &stream, 1);
搜刮引擎中的应用
在本部分中，我们将展示在搜刮引擎中应用了稀疏化加速的四个应用案例：
第一是搜刮中的相关性猜测，旨在评估输入文本和数据库中视频之间的相关性。
第二是查询性能猜测，用于文档召回交付策略。
第三是用于召回最相关文本的召回任务。
第四是文生图任务，该任务根据输入的提示词自动天生相应的图片。
搜刮相关性案例
我们使用 PNR (Positive Negative Rate，正负率) 或 ACC (accuracy，精度) 标准来评估稀疏化加速在这些应用案例的效果。在相关性案例 1 中，我们运行 Sparse-QAT 得到了一个稀疏 int8 模子，该模子在两个紧张的 PNR 评估指标均优于在线 int8 模子。
https://i-blog.csdnimg.cn/blog_migrate/6c464abab17a9a0a1bc922fb9882df59.jpeg
在相关性案例 2 中，稀疏 int8 模子可以得到与 float32 模子靠近的 Acc 分数，相比稠密 int8 模子，其得到了 1.4 倍的推理加速。
https://i-blog.csdnimg.cn/blog_migrate/eb39d56a3f41879a2a79eab28dd974ed.jpeg
查询性能猜测案例
在这部分，我们展示了查询性能猜测 (query performance prediction, QPP) 的四个用例，其效果使用 NDCG（normalized discounted cumulative gain, 标准化折扣累积增益）评估。如表 3 所示，这些稀疏 float16 模子乃至可以得到比原始 float32 模子更高的 NDCG 分数，同时推理速率相比于 float32 模子进步了 4 倍。
https://i-blog.csdnimg.cn/blog_migrate/2078e4508bdc08278afc0273a1916f75.jpeg
文档查询案例
表 4 显示了搜刮引擎中文档查询案例的结果，与稠密 int8 模子相比，使用我们推荐的 Sparse-QAT 训练流水线，稀疏 int8 模子可以实现 1.4 倍的推理加速，准确度损失可忽略不计。
https://i-blog.csdnimg.cn/blog_migrate/e1a14f9fad0211c6b160db8d38dc2e96.jpeg
文生图案例
图 6 展示了文生图模子的结果，上面四张图片是用稠密 float32 模子输出，下面四张图片是用稀疏 float16 模子输出。
从结果中您会发现，输入相同的提示，稀疏模子可以输出与稠密模子相称的结果。而且引入模子稀疏化和额外的渐进式稀疏调优使得模子从数据中学习到了更多内容，因此部分稀疏模子的输出结果看起来更为合理。
https://i-blog.csdnimg.cn/blog_migrate/ac49feee6cb7843f5195114792a85e9e.jpeg
图 6. 搜刮引擎中的文生图案例
（在文生图的案例中，稀疏模子大概会产生比密集模子更合理的结果。）
总结
NVIDIA Ampere 架构中的布局化稀疏功能可以加速许多深度学习工作负载，并且易于结合 TensorRT 和 cuSPARSELt 稀疏加速库一起使用。
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。
页:
[1]