作者 | Manish Gupta
OneFlow编译
翻译｜宛子琳、杨婷
AI驱动的技能正逐渐融入人们日常日子的各个旮旯，有望提高人们获取知识的能力，并提高全体生产功率。言语大模型（LLM）正是这些应用的中心。LLM对内存的需求很高，一般需求专用的硬件加快器，以高效地供给数百亿亿次浮点运算（Exaflops级别）的核算能力。本文将展现怎么经过更有效地运用内存来处理核算方面的应战。
LLM中的大部分内存和核算资源都耗费在了矩阵乘法操作中的权重上。运用范围更小的数据类型能够下降内存耗费，例如，将权重存储为8位整数（即U8或S8）的数据类型，相关于单精度（F32）能够削减4倍的内存占用，相关于半精度（F16）或bfloat16（BF16）能够削减2倍的内存占用。
此外，先前的研究标明，LLM模型选用S8格局的权重和F16格局的输入进行矩阵乘法运算，能够在保持可接受的准确性的同时提高功率。这一技能被称为仅权重量化（weight-only quantization），需求对带有混合输入的矩阵乘法进行高效完成，例如半精度输入与8位整数相乘。因为硬件加快器（包括GPU）支撑一组固定的数据类型，因而，混合输入矩阵乘法需求经过软件转化来映射到硬件操作。
为此，本文重点重视将混合输入的矩阵乘法映射到NVIDIA Ampere架构上。咱们提出了处理数据类型转化和布局一致性的软件技能，以有效地将混合输入矩阵乘法映射到硬件支撑的数据类型和布局上。成果显示，在软件中进行额定工作的核算开销很小，而且能够完成挨近硬件峰值的功能。本文所介绍的软件技能已在开源的NVIDIA/CUTLASS库（github.com/NVIDIA/cutlass/pull/1084）中发布。
175亿参数的LLM模型在不同数据类型格局下的内存占用。
（本文作者为谷歌研究院高档软件工程师Manish Gupta。以下内容由OneFlow编译发布，转载请联系授权。原文：blog.research.google/2024/01/mix…
1
矩阵乘累加（matrix-multiply-accumulate）运算
当时的AI硬件加快器，如Google的TPU和NVIDIA的GPU，经过针对张量中心（Tensor Core）在硬件中本地履行矩阵乘运算（这些张量中心是专门加快矩阵运算的处理单元），特别适用于AI工作负载。本文咱们重点重视NVIDIA Ampere张量中心，它供给矩阵乘累加（mma）运算。在本文其余部分，mma指的是Ampere张量中心。在mma运算中，两个输入矩阵（称为操作数）所支撑的数据类型、维度和数据布局在硬件中是固定的。这意味着，软件中不同的数据类型和更大维度的矩阵乘法是经过将问题划分为硬件所支撑的数据类型、形状和布局完成的。
张量中心的mma运算经过指定两个输入矩阵（如下图所示的A和B）来核算生成成果矩阵C。mma运算本身支撑混合精度。混合精度张量中心答应混合输入（A和B）数据类型与成果（C）数据类型。相比之下，混合输入矩阵乘法涉及混合输入数据类型，这在硬件上不受支撑，因而需求经过软件完成。
对M乘K的输入矩阵A和K乘N的输入矩阵B进行的M乘N乘K的张量中心操作， 得到M乘N的输出矩阵C。
2
混合输入矩阵乘面对的应战
为简化评论，咱们挑选了混合输入矩阵乘法的一个详细示例：用户输入选用F16，模型权重选用U8（表明为F16 * U8）。本文评论的技能适用于各种混合输入数据类型组合。
GPU程序员能够访问一系列内存，包括大局内存、同享内存和寄存器，这些内存按容量递减但速度递增的顺序排列。NVIDIA Ampere Tensor Core的mma操作从寄存器中获取输入矩阵。此外，输入和输出矩阵需求契合在一个名为warp的32个线程组内的数据布局。关于mma操作，warp内支撑的数据类型和布局是固定的，因而要高效完成混合输入乘法，就需求在软件中处理数据类型转化和布局一致性问题。
数据类型转化
mma操作要求两个输入矩阵具有相同的数据类型。因而，在混合输入矩阵乘法中，当一个操作数以U8存储在大局内存中，而另一个以F16存储时，就需求进行从U8到F16的数据类型转化。这种转化将两个操作数转化为F16，从而将混合输入矩阵乘法映射到硬件支撑的混合精度张量中心。鉴于权重的数量巨大，因而需求很多的转化操作，咱们的技能展现了怎么下降当时延并提高功能。
布局一致性
mma操作还要求两个输入矩阵的布局（即在一个warp的寄存器中的布局）契合硬件标准。在混合输入矩阵乘法（F16 * U8）中，U8数据类型的输入矩阵B的布局需求契合转化后的F16数据类型。这被称为布局一致性（layout conformance），需求经过软件完成。
下图展现了一个mma操作，它从寄存器中提取矩阵A和矩阵B，然后在寄存器中生成矩阵C，这个过程分布在一个warp中。其间，线程T0被突出显示，并对其进行了放大，以展现权重矩阵B经过数据类型转化，需求契合布局一致性才干映射到硬件支撑的张量中心操作。
将软件中的混合输入（F32=F16U8）操作映射到硬件中原生支撑的warp级张量中心（F32=F16F16）。原图来源：《在NVIDIA A100上开发CUDA中心以充分发挥张量中心的功能极限》。
2
应对核算应战的软件战略
典型的数据类型转化涉及对32位寄存器的一系列操作，如下图所示。每个矩形块代表一个寄存器，相邻文本则表明相应的操作。整个序列展现了从4个U8转化为2个(2个F16)的过程。该序列大约包括10个操作。
在32位寄存器中，将4个U8转化为2x(2个F16)的NumericArrayConvertor。
完成布局一致性的办法有很多，两种现有处理方案如下：
1.较窄位宽的同享内存加载：在这种办法中，线程宣布较窄位宽的内存加载操作，将U8数据从同享内存移动到寄存器。这会导致两个32位寄存器，每个寄存器包括2个F16值（如上所示，关于矩阵B的线程T0）。较窄的同享内存加载直接完成了布局一致性，使其存入寄存器，而无需任何移动（shuffles）操作；然而，这种办法未充分运用同享内存带宽。
2.大局内存中的预处理：另一种战略是，在大局内存中重新排列数据（在内存层次结构中位于同享内存的上一级），答应更宽的同享内存加载。这种办法最大程度地运用了同享内存带宽，并保证数据以一致的布局直接加载到寄存器中。虽然重新排列过程能够在LLM部署之前离线履行，保证不影响应用程序的功能，但它引入了一个额定的、有意义的硬件特定的预处理步骤，需求额定的程序来重新排列数据。
NVIDIA/FasterTransformer选用这种办法有效地处理了布局一致性的应战。
3
优化的软件战略
为进一步优化并削减数据类型转化和布局一致性的核算开销，咱们别离完成了FastNumericArrayConvertor和FragmentShuffler。
FastNumericArrayConvertor在32位寄存器中直接处理4xU8，而无需拆解单个1xU8值。此外，它运用的算术操作本钱较低，削减了指令数量，提高了转化速度。
U8到F16的转化序列如下图所示。这些运算运用打包的32位寄存器，避免了显式的解包和打包。FastNumericArrayConvertor运用置换字节来重新排列4xU8的字节，将其放入两个寄存器中。此外，FastNumericArrayConvertor不运用开销较大的整数到浮点数转化指令，并选用矢量化操作，在两个32位寄存器中获取包括2x(2xF16)值的打包成果。相关于上述办法，U8到F16的FastNumericArrayConvertor大约运用了六个操作，相对上文说到的方式，其功能有约1.6倍的提高。
FastNumericArrayConvertor运用permute字节和packed核算，削减了数据类型转化中的指令数量。
FragmentShuffler经过对数据进行重新排列，能够运用更宽的位宽加载操作，完成了布局一致性，增加了同享内存带宽运用率，并削减了总操作数。
NVIDIA Ampere架构供给了一个加载矩阵指令（ldmatrix）。ldmatrix是一种warp级操作，其间一个warp的32个线程将数据从同享内存移动到寄存器中，而这些寄存器的形状和布局契合矩阵A和B进行矩阵乘法累积运算所需的要求。运用ldmatrix削减了加载指令的数量，提高了内存带宽运用率。因为ldmatrix指令将U8数据移动到寄存器中，加载后的布局契合U8U8的mma操作，不契合F16F16的mma操作。咱们完成了FragmentShuffler，运用shuffle（shfl.sync）操作在寄存器内重新排列数据，以完成布局一致性。
这项工作最重要的贡献之一便是经过寄存器shuffles完成了布局一致性，避免了在大局内存中进行离线预处理或更窄的位宽同享内存加载。此外，咱们供给了FastNumericArrayConvertor的完成，涵盖了从U8到F16、S8到F16、U8到BF16以及S8到BF16的数据类型转化。
4
功能表现
咱们在NVIDIA A100 SXM芯片上测量了该办法的八种混合输入变体的功能（如下图中的蓝色和赤色所示；根据矩阵A和B的数据类型不同而变化）以及两种混合精度数据类型（绿色显示）的功能。功能成果以FLOPS（数值越高表明功能越好））显示。
值得注意的是，相关于最终两个矩阵乘法，前八个需求额定的操作，因为混合精度变体直接针对硬件加快的张量中心操作，无需数据类型转化和布局一致性。即便如此，在混合输入矩阵乘法功能上，咱们的办法仅略低于或与混合精度相当。
在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表明由衷的感谢。
试用图片/视频生成加快引擎OneDiff:
github.com/siliconflow/onediff
