四. TensorRT模型部署优化-pruning(sparse-tensor-core)

前言

自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考

本次课程我们来学习课程第四章—TensorRT 模型部署优化，一起来学习 sparse tensor core

课程大纲可以看下面的思维导图

0. 简述

本小节目标：理解 NVIDIA 是如何使用 sparse tensor core 来处理带有稀疏性的矩阵乘法

这节课我们来学习剪枝的第三个小部分—sparse tensor core，这个小节我们主要是来理解 NVIDIA 的 sparse tensor core 是如何处理稀疏性的矩阵乘法的，这里做一个科普稍微给大家扩展一下

下面我们开始本次课程的学习🤗

1. 自动驾驶中需要关注的电力消耗

下面是 Jetson AGX Orin 的电力消耗图，它的 Peak performance 和对应的电力消耗是：

• GPU：170 TOPS（Sparse INT8），85 TOPS（DENSE INT8）
• DLA0：52.5 TOPS（Sparse INT8），26.25 TOPS（DENSE INT8）
• DLA1：52.5 TOPS（Sparse INT8），26.25 TOPS（DENSE INT8）
• Power consumption ~60W

Jetson AGX Orin

我们可以看到 Orin 中的 Sparse INT8 的 TOPS 比 Dense INT8 高两倍，那我们自然而然回去想 Sparse 和 Dense 分别代表着什么呢？其实 Sparse 就是 NVIDIA Ampere 架构中所支持的一个东西，下面我们一起来看看

2. Ampere架构中的3rd Generation Tensor core

在 Ampere 架构（A100/Jetson AGX Orin）中的第三代 Tensor core 支持带有 sparsity 的 matrix 计算，更准确来说：

• 支持 fine-grained structured sparsity
• fine-grained 细粒度，主要是对权重的某个元素本身进行分析剪枝，是否归零
• structured 表现在 sparsity 的 pattern 是以 1x4 vector 的大小进行 2:4 的归零（vector-wise pruning），具体如下图所示，每一个 1x4 vector 只保留其中的两个元素即图中绿色部分，另外两个归零即图中白色部分
• 50% 粒度的 sparse pruning，理论上可以实现 2x 的吞吐量的提升

结构化稀疏矩阵

tensor core计算示意图

我们下面来看看 sparse 在 tensor core 中是如何做支持的，假设我们有一个 Dense matrix $W_0$，通过 Fine-grained structured pruning 之后变成了一个 structured-sparse matrix $W$，如下图所示：

也就是说一个 dense 的 matrix 会以 2:4 sparsity 的方式进行剪枝（每 4 个连续的权重中最小的两个归零），下图展示了 pruning 的过程：

下面我们以具体的数据来更加形象的说明整个 Fine-grained structured pruning 过程，Original 代表原始的 Dense Matrix，通过剪枝将其中的每四个中的两个进行 pruning 剪枝，得到 2:4 Sparse Matrix：

剪枝完成之后，对于已经 sparse pruning 过的 matrix 可以进一步进行压缩，在 memory 中只保存非零的 weight，至于哪些 weight 是零，哪些 weight 非零，我们可以用一个 2-bits indices 来保存（可以把它理解为一种索引），如下图所示：

这样一来 weight 的大小减半，同时对于 activation values 可以通过 2-bits indices 来决定 activation values 中哪些值是参与计算的，哪些是 skip 掉的（这个过程需要特殊的硬件 unit 来实现），从而实现 2x 的计算吞吐量的提升，整个过程如下图所示：

上图还对比了 Dense Matrix 计算和 Sparse Matrix 计算，可以看到 Sparse Matrix 的计算省去了 Dense Matrix 中很多没必要的计算，它将一个 1x8 与 8x1 的计算通过 indices 变成了一个 1x4 与 4x1 的计算

3. Sparse tensor core做矩阵乘法

我们再来看一个具体的例子，Dence Tensor core（FP16）的计算 A(M,K) * B(K,N) = C(M,N) 的过程可以用下图来表示：

我们知道 A100 中的 tensor core 可以用 1 cycle 完成一个 16x16 * 16x8 = 16 * 8 的矩阵乘法，需要用 2 cycle 完成一个 16x32 * 32x8 = 16 * 8 的矩阵乘法，如下图所示：

上面展示的是 Dense Tensor core（FP16）的计算过程，下面我们来看看 Sparse Tensor core（FP16）的计算 A(M, k) * B(K, N) = C(M, N) 的过程，如下图所示：

如果说 A 中的 matrix 拥有 sparsity，是按照 2:4 的 pattern 进行 pruning，我们可以重构 A，如下图所示，重构后的 A 的 memory 占用空间直接减半达到压缩的效果，我们可以把这里的 A 理解为 conv 或 FC 中的 weight

那么对于 B，可以通过索引对 B 中参与计算的值进行筛选（也可以理解为对 B 的重构），如下图所示，我们可以把这里的 B 理解为 conv 或 FC 中的 activation values。另外这个筛选的过程其实是需要硬件支持的，也就是只能是第三代 tensor core 才能做，其他硬件可能就不行，这也就是为什么 Fine-grained pruning 细粒度剪枝可能会需要特定硬件的支持

这样我们可以得到一个结论：

• 使用 dence tensor core 需要用 2 cycle 完成一个 16x32 * 32x8 = 16 * 8 的矩阵乘法
• 使用 sparse tensor core 只需要 1 cycle 就可以完成一个 16x32 * 32x8 = 16 * 8 的矩阵乘法

然而这里面很容易忽视的一点就是，为了实现 sparse 的计算而添加的额外操作的 overhead

• compress weight 的 overhead
• reconstruct activation values 的 overhead

虽然这些是在硬件上可以完成，从而自动的将 0 参与的计算全部 skip 掉，然而这些多余的操作有时会比较凸显，尤其是当模型并不是很大，参与 sparse 的计算的激活值不是很大时，使用 sparsity 的特性做计算效果不是那么好。目前认为 sparse tensor core 在 NLP 领域的加速可能会比较可观。

上图展示了参与计算的参数量的不同 sparse tensor core 体现出来的加速比也不同，如果参与计算的参数量比较少时，此时 sparse 带来的加速比不是很明显，随着参数量的增加，sparse 带来的加速比也随之增加

总结

本次课程我们主要学习了 NVIDIA 的第三代 tensor core 是如何处理稀疏的矩阵乘法计算的，它主要是通过 Fine-grained pruning 将 Dense Matrix 变成 Sparse Matrix，随后进行 compressed 压缩得到压缩后的 weights 和对应的 indices，通过 indices 可以 skip 掉 weights 和 activation values 之间的一个 0 参与的计算。不过需要注意的是 sparse 的计算会有一些额外的 overhead 导致在参与计算的参数量比较小时加速并不是很可观

OK，以上就是剪枝的第三个小部分 sparse tensor core 的全部内容了，下节我们进入第五章—实战：TensorRT API的基本使用的学习，敬请期待😄

参考

• Delivering Server-Class Performance at the Edge with NVIDIA Jetson Orin
• Accelerating Sparse Deep Neural Networks