矩阵乘法背景指南
说明：本文翻译自 NVIDIA 官方文档 Matrix Multiplication Background User's Guide，版权归 NVIDIA 所有。
摘要
本指南介绍矩阵乘法及其在众多深度学习操作中的应用。这里描述的规律构成了全连接层、卷积层、循环层等性能趋势的基础。
1. 背景：矩阵-矩阵乘法
GEMM（General Matrix Multiplications，通用矩阵乘法）是神经网络中许多操作的基础构建块，例如全连接层、RNN/LSTM/GRU 等循环层，以及卷积层。本指南描述理解这类层性能所共通的 GEMM 性能基础。
GEMM 定义为运算
2. 数学与内存界限
遵循各种线性代数库（如 BLAS）的惯例，我们称矩阵 A 为 M x K 矩阵，即它有 M 行、K 列。同理，B 和 C 分别为 K x N 和 M x N 矩阵。
A 与 B 的乘积有 M x N 个值，每个值都是 K 个元素向量的点积。因此，计算该乘积总共需要 M * N * K 次融合乘加（Fused Multiply-Add，FMA）。每次 FMA 是 2 次运算——一次乘法、一次加法，所以总共需要
要判断某个矩阵乘法是受数学限制还是受内存限制，我们将它的算术强度与 GPU 的 ops:byte 比值比较，如理解性能所述。假设使用 NVIDIA® V100 GPU 以及 FP16 输入、FP32 累加的 Tensor Core 运算，从 GPU 内存加载数据时 FLOPS:B 比值为 138.9。
举个例子，考虑一个 M x N x K = 8192 x 128 x 8192 的 GEMM。在这个具体情形下，算术强度为 124.1 FLOPS/B，低于 V100 的 138.9 FLOPS:B，因此该操作是内存受限的。如果我们将 GEMM 尺寸增大到 8192 x 8192 x 8192，算术强度提高到 2730，远高于 V100 的 FLOPS:B，因此该操作是数学受限的。特别地，从上述分析可以得出：矩阵-向量乘积（GEMV，即通用矩阵-向量乘法，其中 M=1 或 N=1）永远是内存受限的，它们的算术强度小于 1。
值得牢记的是，把算术强度与 ops:byte 比值比较是一个简化的经验法则，并未考虑实现该计算的许多实际因素（例如指针算术等非算法指令，或 GPU 片上内存层级的贡献）。
2.1 GPU 实现
GPU 通过把输出矩阵划分成分块（Tile），再把这些分块分配给线程块来实现 GEMM。
在本指南中，分块大小通常指这些分块的维度（图 1 中的 Mtile x Ntile）。每个线程块沿 K 维度按分块逐步遍历，从 A 和 B 矩阵加载所需的值，将它们乘积累加到输出中，从而计算出自己的输出分块。
图 1. GEMM 的分块外积方法
2.2 Tensor Core 要求
正如我们在 GPU 架构基础中讨论的，最新的 NVIDIA GPU 引入了 Tensor Core 来最大化张量乘法的速度。使用 Tensor Core 的要求取决于 NVIDIA 库的版本。当矩阵维度 M、N、K 对齐到 16 字节（在 A100 上为 128 字节）的倍数时，性能更好。对于 cuBLAS 11.0 之前的版本或 cuDNN 7.6.3 之前的版本，这是使用 Tensor Core 的前提条件；自 cuBLAS 11.0 和 cuDNN 7.6.3 起，无论维度如何都可以使用 Tensor Core，但当矩阵维度是 16 字节的倍数时效率更高。例如，使用 FP16 数据时，每个 FP16 元素占 2 字节，因此矩阵维度最好是 8 个元素的倍数（在 A100 上为 64 个元素）以获得最佳效率。
表 1. 按 cuBLAS 或 cuDNN 版本划分的、常见数据精度的 Tensor Core 要求。这些要求适用于矩阵维度 M、N 和 K。
| 可以使用 Tensor Core 的精度 | cuBLAS 版本 < 11.0 cuDNN 版本 < 7.6.3 | cuBLAS 版本 ≥ 11.0 cuDNN 版本 ≥ 7.6.3 |
|---|---|---|
| INT8 | 16 的倍数 | 始终可用，但 16 的倍数时最高效；在 A100 上为 128 的倍数 |
| FP16 | 8 的倍数 | 始终可用，但 8 的倍数时最高效；在 A100 上为 64 的倍数 |
| TF32 | 不适用（N/A） | 始终可用，但 4 的倍数时最高效；在 A100 上为 32 的倍数 |
| FP64 | 不适用（N/A） | 始终可用，但 2 的倍数时最高效；在 A100 上为 16 的倍数 |
实际上要求更宽松——只有内存中最快变化的维度需要遵守这一规则——但最简单的方式是把三个维度等同对待。对所有维度遵循这些对齐要求，能确保 Tensor Core 被启用并高效运行。这一效果可见于 图 2：当 K 能被 8 整除时，计算最快（耗时最短）。当 K 不能被 8 整除时，从 cuBLAS 10.2 切换到 cuBLAS 11.0 允许使用 Tensor Core，带来 2-4 倍的加速。同样值得注意的是，使用 cuBLAS 11.0 时，在那些不能被 8 整除的 K 值中，偶数仍比奇数计算得更快。我们建议把矩阵维度选为 16 字节的倍数（如表 1中 FP16 的 8）；如果这不可行，选择较小 2 的幂的倍数（如 8 或 4 字节）在使用 cuBLAS 11.0 及更高版本时往往仍有帮助。在 A100 上，选择高达 128 字节（FP16 为 64）的更大 2 的幂倍数，可以进一步提高效率。
图 2. 使用 FP16 数据时，(a) cuBLAS 10.1 与 (b) cuBLAS 11.0 的 GEMM 执行时间对比。当 K 能被 8 整除时计算最快（耗时最短）。"NN" 表示 A 和 B 矩阵都以非转置方式访问。NVIDIA V100-DGXS-16GB GPU。
2.3 cuBLAS 中的典型分块维度与性能
cuBLAS 库包含 NVIDIA 优化的 GPU GEMM 实现（文档见此处）。
虽然有多种分块策略可用，但较大的分块有更多的数据复用，能使用更少的带宽，因而比小分块更高效。另一方面，对于给定规模的问题，使用较大的分块会产生更少的并行分块，可能导致 GPU 利用不足。当 TensorFlow 或 PyTorch 等框架以特定 GEMM 维度调用 cuBLAS 时，cuBLAS 内部的一个启发式规则会选择一个预期性能最佳的分块选项。另外，一些框架提供"基准测试"（benchmark）模式，在训练之前对所有实现选项计时并选择最快的一个（这构成每次训练会话一次性开销）。
分块效率与分块并行度之间的权衡意味着：GEMM 越大，这种权衡就越不重要——到某个程度，一个 GEMM 有足够的计算量，可以同时使用最大可用的分块并把 GPU 填满。反之，如果 GEMM 太小，分块效率或分块并行度的下降很可能使 GPU 无法达到峰值数学利用率。图 3 和 图 4 展示了这一总体趋势：更大的 GEMM 达到更高的吞吐量。
图 3. 随着 GEMM 的 M-N 规模增大，性能提升。耗时也会增加，但不如 M-N 维度本身增长得快；有时可以只用很小的时间开销就增大 GEMM 规模（使用更多权重）。NVIDIA A100-SXM4-80GB，CUDA 11.2，cuBLAS 11.4。
图 4. 随着 K 维度增大，性能提升，即使 M=N 已经相对较大时也如此，因为点积更长时，计算的建立与拆解开销能被更好地摊还。NVIDIA A100-SXM4-80GB，CUDA 11.2，cuBLAS 11.4。
对于 cuBLAS GEMM，线程块的分块大小通常（但不一定）使用 2 的幂维度。不同用例可能使用不同分块大小，但作为起点，可用的分块如下：
- 256x128 和 128x256（最高效）
- 128x128
- 256x64 和 64x256
- 128x64 和 64x128
- 64x64（最低效）
图 5 展示了其中几种分块大小之间的效率差异示例：
图 5. 更大的分块运行得更高效。基于 256x128 的 GEMM 每个 SM 恰好运行一个分块，其他 GEMM 根据各自的分块大小生成更多分块。NVIDIA A100-SXM4-80GB，CUDA 11.2，cuBLAS 11.4。
图中展示了一个 MxNxK = 6912x2048x4096 的 GEMM 在不同分块大小下的性能。它表明较小分块带来的分块并行度提升（64x64 比 256x128 多 8 倍并行度）是以可观的效率代价换来的。在实践中，对于大到足以用较大分块获得足够并行度的 GEMM，cuBLAS 会避免使用小分块；只有当运行的 GEMM 远小于该示例时，才会求助于更小的分块。顺带一提，NVIDIA 库还具备在 M 和 N 都较小但 K 较大时沿 K 维度"分块"的能力。由于 K 是点积的方向，沿 K 分块最后需要一次归约（Reduction），这可能限制可达性能。为简单起见，本指南大部分内容假定不在 K 方向分块。
3. 维度量化效应
如 GPU 执行模型所述，GPU 函数通过启动若干线程块来执行，每个线程块具有相同数量的线程。这会对执行效率引入两种潜在影响——分块量化（Tile Quantization）与波量化（Wave Quantization）。
3.1 分块量化
当矩阵维度不能被线程块分块大小整除时，就会发生分块量化。
线程块分块的数量足够多，能保证覆盖所有输出元素，但某些分块的实际工作非常少，如图 6 所示（该图假定 128x128 分块和两个矩阵维度）。
图 6. 使用 128x128 线程块分块的示例。(a) 最佳情况——矩阵维度能被分块维度整除；(b) 最坏情况——分块量化导致启动六个线程块，其中两个浪费了大部分工作。
虽然库保证任何分块都不会执行非法内存访问，但所有分块都会执行同样多的数学运算。因此，由于分块量化，图 6(b) 的情形执行的算术运算量是 (a) 的 1.5 倍，尽管算法上只需要多 0.39% 的运算。由此可以看出，当输出矩阵维度能被分块维度整除时，利用率最高。
再举一个体现这种效应的例子：考虑 M = 27648、K = 4096，且库函数使用 256x128 分块的 GEMM，我们考察不同 N 取值的情况。当 N 从 136 以 8 为步长递增到 256 时，Tensor Core 加速的 GEMM 始终运行相同数量的分块，也就是说 N 维度始终被划分为 2 个分块。虽然分块数量保持不变，但这些分块中含有有效数据的比例、从而实际执行的有效 FLOPS 数量会随 N 增大而增加，这反映在下文图 7 的 GFLOPS 中。注意，吞吐量在 N = 128（每行单个分块被有效数据填满）和 N = 136（每行增加第二个分块，但其中只有 8/128 = 6.25% 是有效数据）之间存在显著下降。还要注意，只要分块数量保持不变，耗时就是恒定的。
图 7. 分块量化对 (a) 实际 FLOPS 吞吐量、(b) 运行时间以及 (c) 生成分块数量的影响。通过一个强制在 MxN 输出矩阵上使用 256x128 分块的函数测得。在实践中，cuBLAS 会选择更窄的分块（例如 64 宽）来减小量化效应。NVIDIA A100-SXM4-80GB，CUDA 11.2，cuBLAS 11.4。
3.2 波量化
分块量化意味着问题规模被量化为每个分块的大小；还有第二种量化效应，即分块总数被量化为 GPU 上的多处理器数量——这就是波量化（Wave Quantization）。
我们考虑一个与前面相关的例子，同样改变 N 且 K = 4096，但 M 更小，为 2304。一块 NVIDIA A100 GPU 有 108 个 SM；在 256x128 线程块分块的特殊情况下，每个 SM 可以执行一个线程块，因此波大小（Wave Size）为 108，即最多有 108 个分块可以同时执行。因此，当分块数量是 108 的整数倍或略低于它时，GPU 利用率最高。
M 维度始终被划分为每列 2304/256 = 9 个分块。当 N = 1536 时，N 维度被划分为每行 1536/128 = 12 个分块，总共创建 9*12 = 108 个分块，构成一个完整的波。当 1536 < N <= 1664 时，每行额外增加一个分块，总共 9*13 = 117 个分块，构成一个完整波加上一个仅含 9 个分块的"尾波"（tail wave）。在这个例子中，尾波执行所需的时间与完整的 108 分块波几乎相同，但该期间只使用了 A100 的 9/108 = 8.33% 的 SM。因此，从 N = 1536 到 N = 1544，GFLOPS 大约减半、耗时大约翻倍（图 8）。在 N = 3072、N = 4608 和 N = 6144 之后也能看到类似的跳变，这些值同样对应整数个完整波。
图 8. 波量化对 (a) 实际 FLOPS 吞吐量、(b) 运行时间以及 (c) 生成分块数量的影响。通过一个在 MxN 输出矩阵上使用 256x128 分块的函数测得。注意，当分块数量越过 108 的倍数时就会发生量化效应。NVIDIA A100-SXM4-80GB，CUDA 11.2，cuBLAS 11.4。
值得注意的是，波量化的吞吐量与耗时曲线看起来和分块量化非常相似，只是横轴的比例不同。因为两者都是量化效应，这在意料之中。差别在于量化发生的位置：分块量化意味着计算量被量化为分块的大小，而波量化意味着计算量被量化为 GPU 的大小。图 7(c) 和 图 8(c) 在分块量化与波量化的示意图中都展示了这一差异。