人工智能驱动的技术正在融入我们的日常生活，有可能增强我们获取知识的渠道并提高我们的整体生产力。这些应用的支柱是大型语言模型 (LLM)。LLM 占用大量内存，通常需要专门的硬件加速器才能有效提供数十百亿亿次浮点运算的计算能力。这篇博文展示了如何通过更有效地利用内存来开始应对计算挑战。
LLM 的大部分内存和计算都被矩阵乘法运算中的权重所消耗。使用较窄的数据类型可以减少内存消耗。例如，将权重存储在 8 位整数（即 U8 或 S8）数据类型中可将内存占用量减少 4 倍（相对于单精度（F32）而言）和 2 倍（相对于半精度（F16）或bfloat16（BF16）。此外，之前的研究表明，使用 S8 中的权重和 F16 中的输入（保留用户输入的更高精度）运行矩阵乘法的 LLM 模型是一种在准确度可接受的权衡下提高效率的有效方法。这种技术称为仅权重量化，需要高效实现混合输入的矩阵乘法，例如半精度输入乘以 8 位整数。硬件加速器（包括 GPU）支持一组固定的数据类型，因此混合输入矩阵乘法需要软件转换才能映射到硬件操作。
为此，在本博客中，我们重点介绍如何将混合输入矩阵乘法映射到NVIDIA Ampere 架构上。我们介绍了解决数据类型转换和布局一致性的软件技术，以便有效地将混合输入矩阵乘法映射到硬件支持的数据类型和布局上。我们的结果表明，软件中额外工作的开销很小，并且性能接近硬件的峰值能力。此处描述的软件技术已在开源NVIDIA/CUTLASS存储库中发布。
具有各种数据类型格式的 175B 参数 LLM 模型的内存占用。
矩阵乘法累加运算
现代 AI 硬件加速器（例如Google 的 TPU和NVIDIA 的 GPU）通过针对 Tensor Core 在硬件中原生地进行矩阵乘法，Tensor Core 是用于加速矩阵运算（尤其是对于 AI 工作负载）的专用处理元素。在本博客中，我们重点介绍 NVIDIA Ampere Tensor Core，它提供矩阵乘法累加（mma）运算。在博客的其余部分，mma对 Ampere Tensor Core 的引用。运算的两个输入矩阵（称为操作数）支持的数据类型、形状和数据布局mma在硬件中是固定的。这意味着，通过将问题平铺到硬件支持的数据类型、形状和布局上，可以在软件中实现具有各种数据类型和更大形状的矩阵乘法。
Tensor Coremma操作的定义是指定两个输入矩阵（例如，如下所示的A和B ）以生成结果矩阵C。该mma操作本身支持混合精度。混合精度 Tensor Core允许将输入（A和B）数据类型与结果（C）数据类型混合。相比之下，混合输入矩阵乘法涉及混合输入数据类型，并且硬件不支持，因此需要在软件中实现。
对 M×K 输入矩阵 A 和 K×N 矩阵 B 进行 M×N×K 的张量核心运算，产生 M×N 输出矩阵 C。
混合输入矩阵乘法的挑战
为了简化讨论，我们限制为混合输入矩阵乘法的一个具体示例：F16 用于用户输入，U8 用于模型权重（写为 F16 * U8）。此处描述的技术适用于各种混合输入数据类型的组合。
GPU 程序员可以访问内存层次结构，包括全局内存、共享内存和寄存器，这些内存按容量减小但速度增加的顺序排列。NVIDIA Ampere Tensor Coremma操作使用来自寄存器的输入矩阵。此外，输入和输出矩阵需要符合一组 32 个线程（称为 warp）内的数据布局。warp内支持的数据类型和布局对于mma操作是固定的，因此为了有效地实现混合输入乘法，必须解决软件中数据类型转换和布局一致性的挑战。
数据类型转换
该mma操作需要两个具有相同数据类型的输入矩阵。因此，混合输入矩阵乘法（其中一个操作数存储在全局内存中的 U8 中，另一个存储在 F16 中）需要将数据类型从 U8 转换为 F16。转换将把两个操作数带到 F16，将混合输入矩阵乘法映射到硬件支持的混合精度Tensor Core。鉴于权重数量众多，存在大量此类操作，我们的技术展示了如何减少它们的延迟并提高性能。
布局一致性
该mma操作还要求两个输入矩阵在 warp 寄存器中的布局符合硬件规范。混合输入矩阵乘法 (F16 * U8) 中 U8 数据类型的输入矩阵B的布局需要符合转换后的 F16 数据类型。这称为布局一致性，需要在软件中实现。
下图显示了一个mma操作，该操作使用来自寄存器的矩阵A和矩阵B来生成分布在一个 warp 中的矩阵C。线程T0被突出显示并放大，以显示权重矩阵B经过数据类型转换，并且需要布局一致性才能映射到硬件支持的 Tensor Core 操作。
软件中的混合输入（F32 = F16 * U8）操作映射到硬件中本机支持的 warp 级 Tensor Core（F32 = F16 * F16）。（原始图片来源：开发 CUDA 内核以将 Tensor Cores 推向 NVIDIA A100 上的绝对极限。）
应对挑战的软件策略
典型的数据类型转换涉及对 32 位寄存器的一系列操作，如下所示。每个矩形块代表一个寄存器，相邻的文本是操作。整个序列显示从 4xU8 到 2x(2xF16) 的转换。该序列涉及大约 10 个操作。
NumericArrayConvertor32 位寄存器中从 4xU8 变为 2x(2xF16)。
实现布局一致性的方法有很多种。现有的两种解决方案是：
位宽较窄的共享内存负载：在这种方法中，线程发出位宽较窄的内存负载，将 U8 数据从共享内存移动到寄存器。这会产生两个32 位寄存器，每个寄存器包含 2xF16 值（如上所示，矩阵B的线程T0）。较窄的共享内存负载可直接在寄存器中实现布局一致性，而无需任何改组；但是，它不会利用完整的共享内存带宽。
全局内存中的预处理：另一种策略涉及重新排列全局内存（内存层次结构中比共享内存高一级）内的数据，从而允许更宽的共享内存加载。这种方法可以最大限度地提高共享内存带宽利用率，并确保数据以一致的布局直接加载到寄存器中。虽然重新排列过程可以在 LLM 部署之前离线执行，从而确保不影响应用程序性能，但它引入了一个额外的、非平凡的硬件特定预处理步骤，需要额外的程序来重新排列数据。NVIDIA /FasterTransformer采用这种方法来有效解决布局一致性挑战。
优化软件策略
为了进一步优化和减少数据类型转换和布局一致性的开销，我们分别实现了FastNumericArrayConvertor和FragmentShuffler。
FastNumericArrayConvertor在 32 位寄存器中对 4xU8 进行操作，而无需解包单个 1xU8 值。此外，它使用成本较低的算术运算，从而减少了指令数量并提高了转换速度。
U8 到 F16 的转换序列如下所示。这些操作使用打包的 32b 寄存器，避免显式解包和打包。FastNumericArrayConvertor使用permute byte将 4xU8 的字节重新排列到两个寄存器中。此外，FastNumericArrayConvertor不使用昂贵的整数到浮点转换指令，而是使用矢量化操作在包含 2x(2xF16) 个值的两个32 位寄存器中获得打包结果。U8FastNumericArrayConvertor到 F16 的大约使用六个操作，与上面显示的方法相比减少了 1.6 倍。
FastNumericArrayConvertor利用permute bytes和打包算法，减少了数据类型转换的指令数量。
FragmentShuffler通过以允许使用更宽位宽的加载操作的方式混洗数据来处理布局一致性，从而增加共享内存带宽利用率并减少总操作数。
NVIDIA Ampere 架构提供了加载矩阵指令（ldmatrix）。 是ldmatrix一个 warp 级别的操作，其中 warp 的 32 个线程将数据从共享内存移动到寄存器，其形状和布局mma与矩阵A和B所用的形状和布局相同。 使用减少了加载指令的数量并提高了内存带宽利用率。 由于指令将 U8 数据移动到寄存器，因此加载后的布局符合 U8*U8操作，而不是 F16*F16操作。 我们实现了使用 shuffle ( 操作重新排列寄存器内的数据，以实现布局一致性。 ldmatrix ldmatrixmmammaFragmentShufflershfl.sync)
这项工作最重要的贡献是通过寄存器改组实现布局一致性，避免在全局内存中进行离线预处理或位宽较窄的共享内存加载。此外，我们还提供了涵盖从U8 到 F16、S8 到 F16、U8 到 BF16和S8 到 BF16 的FastNumericArrayConvertor数据类型转换的实现。
性能结果
我们在 NVIDIA A100 SXM 芯片上测量了我们方法 的八种混合输入变体（如下图蓝色和红色所示；矩阵A和B的数据类型不同）和两种混合精度数据类型（以绿色显示）的性能。性能结果以FLOPS为单位显示（越高越好）。值得注意的是，前八个矩阵乘法相对于后两个需要额外的操作，因为混合精度变体直接针对硬件加速的 Tensor Core 操作，不需要数据类型转换和布局一致性。即便如此，我们的方法也表现出混合输入矩阵乘法性能仅略低于或与混合精度相当。
NVIDIA A100 40GB SMX4 芯片上针对计算受限矩阵问题形状的混合输入矩阵乘法性能m=3456, n=4096, k=2048.
致谢
我们想感谢几位通过技术集思广益和改进博客文章做出贡献的人，包括 Quentin Colombet、Jacques Pienaar、Allie Culp、Calin Cascaval、Ashish Gondimalla、Matt Walsh、Marek Kolodziej 和 Aman Bhatia。我们要感谢我们的 NVIDIA 合作伙伴 Rawn Henry、Pradeep Ramani、Vijay Thakkar、Haicheng Wu、Andrew Kerr、Matthew Nicely 和 Vartika Singh。
评论