MARLIN: Mixed-Precision Auto-Regressive Parallel Inference on Large Language Models
文章标题:MARLIN:大型语言模型上的混合精度自回归并行推理
作者/机构:Elias Frantar¹, Roberto L. Castro², Jiale Chen¹, Torsten Hoefler³, Dan Alistarh¹,⁴
A1 主要贡献
本文研究了在多客户端并行(即批处理)场景下,如何通过权重定点化技术加速大型语言模型(LLM)的推理。现有技术在单用户推理时能通过减少内存移动获得显著加速,但在批处理场景下,由于算术强度大幅增加,加速效果不明显。本文旨在解决这一问题,即设计一种GPU核心,使其在支持批处理工作负载所需的大量计算的同时,仍能保持为内存密集型(memory-bound)。
核心问题与研究目标:
* 核心问题: 现有的混合精度推理实现在批处理场景下无法提供显著的加速,因为高算术强度使得计算开销难以被减少的内存移动开销所掩盖。
* 研究目标: 设计一种新型的混合精度推理核心(kernel),使其在批处理场景下(如批量大小为16-32)也能实现接近理论极限(如4位定点化的4倍)的加速,并能在更大批量下(如64-128)仍提供可观的性能提升。
创新点与主要贡献:
1. MARLIN核心的设计与实现: 提出了一种名为MARLIN(Mixed-precision Auto-Regressive LINear kernels)的混合精度推理核心。MARLIN结合了多种先进技术,包括异步内存访问、复杂的任务调度与流水线技术,以及定制化的定点化支持。
2. 实现接近最优的批处理加速: 实验表明,在单个LLM层上,MARLIN能够在NVIDIA Ampere GPU上实现接近最优的性能。具体而言,对于4位权重推理,在NVIDIA A10 GPU上,当批量大小最高达到16-32时,MARLIN相对于FP16可获得约3.9倍的加速;当批量大小增加到128时,加速比逐渐降至1.5倍,此时问题转为计算密集型(compute-bound)。
3. 端到端性能提升: 将MARLIN集成到流行的LLM服务引擎vLLM中,在Llama和Falcon等主流开源LLM上进行了端到端测试。结果显示,在批量大小为16时,MARLIN相比vLLM的标准精度核心,可将多用户令牌生成速度提升高达2.8倍。
4. 对稀疏性的扩展: 提出了MARLIN的扩展版本Sparse-MARLIN,它支持NVIDIA的2:4稀疏Tensor Core格式。相对于原始的密集(dense)MARLIN版本,Sparse-MARLIN可带来高达65%的额外加速,端到端加速比最高可达3.2倍。
5. 开源贡献: MARLIN、Sparse-MARLIN及其vLLM集成的代码均已开源,为社区提供了在批处理场景下高效利用权重定点化LLM推理的实用解决方案。
A3 背景知识
2.1 图形处理器(GPU)
GPU架构概述。NVIDIA GPU由一组流式多处理器(Streaming Multiprocessor, SM)阵列构成,这些SM共享一个称为全局内存(Global Memory, GMEM)的DRAM内存和一个L2缓存。每个SM被划分为多个分区,每个分区包含各种处理块,如扭曲调度器(warp scheduler)、寄存器文件(Register File, RF)和L0指令缓存。SM内的处理块共享一个L1缓存,该缓存的一部分可配置为称为共享内存(Shared Memory, SMEM)的高速暂存存储器。每个处理块内有四种类型的计算单元:整数单元、特殊函数单元、浮点单元(FPU)/CUDA核心和张量核心单元(Tensor-Core Units, TCU)。
张量核心单元(TCU)的角色。TCU最早在Volta架构中引入,主要通过每个周期执行一次矩阵乘法累加(MMA)操作来针对机器学习工作负载,从而降低了获取和解码多条指令的成本。在Ampere架构中,TCU在FP16上的性能比在FPU上运行的融合乘加(FMA)操作高出多达16倍。
CUDA编程与执行模型。CUDA模型定义了三个粒度级别:线程块(thread blocks)、扭曲(warps)和线程(threads)。扭曲是CUDA中的基本调度单位,由32个并发执行的线程组成。线程块是一组扭曲的集合,被调度到同一个SM上执行。每个SM上同时运行的扭曲数量和线程块数量取决于硬件限制,例如扭曲调度器的数量、每个线程的寄存器数量或可用的SMEM大小。
2.1.1 现代张量核心单元(TCU)
Ampere架构TCU的扩展。Ampere GPU对TCU进行了扩展,以处理两个方面:1) 细粒度的结构化稀疏性,形成了稀疏张量核心单元(Sparse Tensor-Core Units, SPTC);2) 异步复制操作。首先,通过新的2:4格式支持结构化稀疏性,其性能承诺比原始TCU快2倍,比FPU快32倍。
2:4稀疏格式。2:4格式将左侧(LHS)矩阵划分为长度为4的向量,并对每个向量中的两个元素进行清零,从而形成一个50%稀疏但结构化的矩阵。图2展示了一个SPTC的简化表示。稀疏矩阵由两个数据结构表示:(1)一个包含非零值的“值”结构(蓝色表示);(2)一个包含每个非零值在每4个元素组中位置的“元数据”结构(紫色表示)。元数据结构被SPTC上的新硬件组件用来仅选择右侧(RHS)矩阵中计算所需的元素,从而跳过被清零的值。
异步复制指令。NVIDIA的Ampere微架构引入了数据获取的改进以增强Tensor Core性能,具体是通过一个新的异步复制指令,允许数据直接从GMEM加载到SMEM。如图2所示,在之前的架构中(图2左),数据需要先通过L1缓存加载到RF(全局加载指令),然后传输到SMEM(共享存储指令),最后再移回RF(共享加载指令)。Ampere的新异步复制指令通过避免中间的RF访问节省了SM内部带宽。该指令有两个变体:“access”变体将数据存入L1以供后续访问和重用(图2中),而“bypass”变体则跳过L1缓存(图2右)。
2.2 LLM上的混合精度推理
混合精度推理的动机。混合精度LLM推理通过静态压缩预训练模型权重,并在推理过程中根据需要动态解压,从而有潜力减少模型巨大的内存占用,并相应地加速内存密集型工作负载。
权重定点化。一种标准的LLM压缩方法是仅权重定点化(weight-only quantization),它降低了权重W
的存储精度,而保持层输入X
不变。这种方法非常流行,因为它即使在相对较高的压缩率下也表现出卓越的准确性鲁棒性。
均匀定点化。本文关注均匀定点化,即给定一个向量 $v \in R^n$,我们定义:
其中,$\lfloor·\rceil$表示四舍五入到最近的整数,$z = z(v) = \min(v)$ 映射到零,而 $s = s(v) = (\max(v) - \min(v)) / (2^b - 1)$ 是缩放因子(scale)。重构误差可以计算为 $\epsilon_r = \|v - Q(v, b)\|_2$。通过将v
划分为多个组并对每个组独立进行定点化(例如,每组128个连续值),可以存储每个组的s
和z
值,从而改善误差。
定点化算法。本文使用GPTQ【【10】,Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pretrained transformers. arXiv preprint arXiv:2210.17323, 2022】的一个变体来执行实际的权重定点化。该方法利用二阶信息来补偿定点化误差,从而仅导致微小的准确性下降。但需要强调的是,本文提出的核心(kernel)技术独立于任何特定的定点化算法。
A2 方法细节
3. The MARLIN Kernel
3.1 动机
理论基础。LLM权重定点化的动机在于现代GPU具有很高的浮点运算/字节比率(FLOPs/Bytes ratio),这意味着它们执行浮点运算的速度远快于从内存读取数据的速度。例如,A10 GPU的FLOPs/Bytes比率约为200【【17】,NVIDIA. Nvidia a10 datasheet. https: //www.nvidia.com/content/dam/en-zz/ Solutions/Data-Center/a10/pdf/ datasheet-new/nvidia-a10-datasheet. pdf, 2022a】。在单层矩阵乘法中,处理一个输入token每个权重需要2个FLOPs,而GPU在加载一个4位权重的时间内可以执行100个FLOPs。因此,只要输入批量大小小于理论最优值 $b_{opt} \approx 50$,内存加载将主导运行时。$b_{opt}$是延迟既不受内存限制也不受计算限制的批量大小,此时我们能以最大吞吐量实现最低延迟。在实践中,我们希望恰好在这个批量大小下运行:更小不会带来进一步的加速,更大则不会提高吞吐量。
实现挑战。然而,要实现这样一个能同时最大限度利用GPU所有资源(即计算和内存)的混合精度(FP16-INT4)矩阵乘法(matmul)核心是一个巨大的挑战。本文的目标是通过设计MARLIN(一个高度优化的混合精度自回归线性核心)来尽可能接近这个目标。
3.2 Ampere矩阵乘法
通用高性能矩阵乘法实现。本节描述在GPU上实现峰值性能(均匀精度)矩阵乘法核心的通用概念,特别是在Ampere级设备上。我们紧密遵循CUTLASS的分层并行化模型【【19】,NVIDIA. Efficient GEMM in CUDA. https: //github.com/NVIDIA/cutlass/blob/ main/media/docs/efficient_gemm.md, 2024a】。具体来说,我们考虑将一个 $M \times K$ 的矩阵A与一个 $K \times N$ 的矩阵B相乘,得到一个 $M \times N$ 的输出矩阵C。
* SM级别(SM Level):第一步,将矩阵A、B、C分别划分为 $M_{sm} \times K_{sm}$、 $K_{sm} \times N_{sm}$ 和 $M_{sm} \times N_{sm}$ 的块。由于矩阵乘法的性质,所有 $C_{sm}[i_{sm}, j_{sm}]$ 都可以通过累加所有 $k_{sm}$ 上的 $A_{sm}[i_{sm}, k_{sm}]B_{sm}[k_{sm}, j_{sm}]$ 的结果来独立计算。因此,可以通过将这些 $C_{sm}[i_{sm}, j_{sm}]$ 子问题分配到GPU的独立计算单元(即SM)上,轻松实现并行化。在此阶段,$A_{sm}$ 和 $B_{sm}$ 块必须从全局GPU内存加载。类似地,$C_{sm}$ 最终必须写回全局存储,但中间累加可以直接在寄存器中进行。
* Warp级别(Warp Level):在单个SM处理的子问题内部,会进行另一次等效的划分,这次的参数是 $M_{wa}$、 $K_{wa}$ 和 $N_{wa}$。这是为了将独立的 $C_{wa}[i_{sm}, j_{sm}][i_{wa}, j_{wa}]$ 输出累加任务分配给不同的warp。关键在于,SM块 $A_{sm}$ 和 $B_{sm}$ 可以临时存储在共享内存中,这样不同warp重复加载 $A_{wa}$ 和 $B_{wa}$ 的速度会快得多。同时,输出 $C_{wa}$ 保留在相应warp的寄存器中,消除了累加过程中的任何额外内存访问成本。
* Tensor Core级别(Tensor Core Level):最终,每个warp将重复乘以 $M_{wa} \times K_{wa}$ 和 $K_{wa} \times N_{wa}$ 的矩阵。虽然此级别的矩阵维度很小,但它们通常仍超过基本的Tensor Core ( $M_{tc}, K_{tc}, N_{tc}$ )形状。因此,需要最后一步划分。然而,与之前不同,$C_{tc}$ 是由单个warp顺序累加的。虽然此时所有数据都在寄存器中,没有内存访问成本,但将 $k_{tc}$ 的循环放在最外层仍然很重要。这是为了尽可能多地消除Tensor Core操作之间的顺序依赖,以最大化吞吞吐量。需要注意的是,实际使用Tensor Core还需要以非常特定的模式将矩阵元素分布到线程中,这是微架构强制的技术细节,而不是另一个灵活的并行化机会。
3.3 混合精度挑战
混合精度实现面临的挑战。将上述均匀精度矩阵乘法模型应用于混合精度场景,同时保持峰值性能(特别是在M中等大小,操作接近内存密集型时)是具有挑战性的,原因如下:
1. 并行化配置:必须非常仔细地配置各个并行化级别,以确保加载量化操作数B确实是核心的主要运行时瓶颈,而不是例如重复重新加载全精度 $A_{sm}$ 块。
2. 内存加载效率:由于运行时主要由内存加载决定,尽管B的表示被显著压缩,这方面必须达到峰值效率。
3. 计算与内存的重叠:对于中等大小的M,矩阵乘法计算的成本可能接近于总内存加载成本,因此需要极其小心的重叠才能接近理论性能。此外,还需要管理量化元数据,这使得这部分更加棘手。
4. 并行化限制:由挑战1带来的划分约束,加上M不是很大这一事实,显著限制了并行化选项。这使得在SM和warp级别上都难以实现峰值的内存加载和计算。现有模型的矩阵形状可能对特定GPU不利,进一步放大了这种效应。
MARLIN核心专门解决了以上所有挑战,最终使其在许多实际环境中能够实现接近峰值的性能。
3.4 核心设计
核心设计的前提。我们假设矩阵A为全FP16精度,而 $K \times N$ 的矩阵B已被(对称地)量化为INT4,可以为N列中的每一列使用一个FP16缩放因子,或者为每列中每G个连续权重使用一个缩放因子,总共有 $\lceil K/G \rceil N$ 个缩放因子。
受限于权重加载。理论上,执行我们的目标矩阵乘法需要接触 $16MK + 4KN + 16MN$ 位内存(读取两个操作数并写入结果),同时执行 $MKN$ 次乘法累加操作(每次计为2个FLOPs)。如果M相对较小,我们的问题算术强度较低。因此,它应该受限于从全局GPU内存中读取量化权重B的成本。为了在实践中保持这一理论,与之前研究的M=1的情况(其中A和C都非常小)不同,现在输入和输出具有不可忽略的大小,因为这些操作数的位宽比我们的权重要高4倍。因此,我们需要选择一个足够大的 $N_{SM}$ 来最小化昂贵的 $A_{sm}$ 块的重载。同时,这减少了 $C_{sm}[i_{sm}, j_{sm}]$ 子问题的数量,使得难以充分利用所有SM。解决这些问题的关键是利用GPU的L2缓存,它通常比全局内存快得多。此外,GPU可以同时从L2加载到L1和从全局加载到L2。因此,我们可以将这些加载流水线化,只要总内存流量不超过L2带宽,就可以完全隐藏 $A_{sm}$ 块加载的带宽成本。因此,我们通过将C划分为大小为 $M \times N_{sm}$ (其中 $N_{sm} \in \{64, 128, 256\}$)的块,即中等宽度的完整输入批次大小的瓦片,然后将每个相应的独立矩阵乘法子问题分配给一个SM。在 $N_{sm}=256$ 时,即使批次大小 $M=64$ 仍然受限于全局权重加载。更精确地说,只要从L2读取 $A_{sm}$ 和 $B_{sm}$ 块的速度更快,全局加载 $A_{sm}$ 块仍然是瓶颈,即:
其中 $B_{l2}$ 和 $B_{gl}$ 分别表示L2和全局带宽。
最大化加载带宽。为了最大化实际加载带宽,我们旨在利用尽可能宽的加载;在当前GPU上,每个线程16字节(128位)。这意味着一个warp可以通过一条指令加载 $32 \times 32 = 1024$ 个INT4权重。为达到峰值效率,我们需要一个warp中的8个线程各自从GMEM读取128字节的连续块(假设是128字节对齐的地址),即一个完整的缓存行。对于形状为 $M \times K_{sm}$ 的A块,这要求 $K_{sm}$ 至少为64。由于权重在推理期间是静态的,因此可以离线预处理,我们通过重新排列 $16 \times 64$ 的瓦片,使其在内存中连续布局,从而实现最优加载,并简化了相应的索引。尽管我们不断从L2缓存中重新加载 $A_{sm}$ 块,但B的每个元素只被访问一次。然而,每次读取都会被放入L2缓存,可能会驱逐一些SM仍然需要的A的部分。为了避免这种缓存污染,我们使用带有“evict first”缓存提示的 cp.async
指令,确保不必要存储的B数据在任何其他缓存行之前被丢弃。
共享内存布局。我们总是通过Ampere的 cp.async
指令从全局(或L2)异步加载到共享内存;这不需要临时寄存器,也使得将这些加载与计算重叠变得更加容易。由于我们对B进行了离线预处理,我们可以简单地以连续的方式复制到共享内存,避免了bank冲突。相比之下,处理A的片段需要更加小心:具体来说,我们需要确保每个 $16 \times 16$ FP16 A块的索引 $ij, (i+8)j, i(j+1)$ 和 $(i+8)(j+1)$ 对应的16字节向量存储在不同的内存bank中。只有这样 ldmatrix.sync
指令才能无冲突地执行。(这些指令加载A操作数数据并将其分布到warp线程中,为Tensor Core的使用做准备。)这可以通过将16字节元素ij存储在激活瓦片中对应共享内存瓦片的 $i(i \oplus j)$ 位置来实现,其中 $\oplus$ 表示XOR操作【【20】,NVIDIA. CUTLASS convolution. https://github. com/NVIDIA/cutlass/blob/main/media/ docs/implicit_gemm_convolution.md, 2024b】。此索引转换的另一个关键方面是,如果一个warp读取全局A瓦片的一个连续子瓦片(例如,前4行),那么它将被置换后仍然整体上连续地写入共享内存。虽然没有文档说明,但这似乎是避免写入时bank冲突所必需的,正如我们在分析NVIDIA性能分析器输出时观察到的。这些索引计算有些复杂,动态处理可能很慢;然而,由于它们只影响相对少量的共享内存位置,并且在主循环中保持静态,我们可以在寄存器中预计算它们,并配合适当的展开。
内存加载流水线。同时达到接近最大带宽和最大计算能力的关键是完全重叠内存加载和Tensor Core数学运算。对于全局到共享内存的加载,这可以通过 cp.async
操作实现,在每次迭代中预取将在 P-1 步后使用的 $A_{sm}$ 和 $B_{sm}$ 块,其中 P 是流水线深度(我们还需要一个缓冲区用于当前瓦片)。此外,我们可以在累加当前部分矩阵乘法之前,从共享内存中预取下一个子瓦片(大多数GPU操作在遇到依赖之前不会阻塞),该部分矩阵乘法的操作数已在上次迭代中被取到寄存器中——这种技术也称为双缓冲(double buffering)【【19】,NVIDIA. Efficient GEMM in CUDA. https: //github.com/NVIDIA/cutlass/blob/ main/media/docs/efficient_gemm.md, 2024a】。我们选择流水线深度 P=4 有两个原因:(a) 在我们所有的测试中,这足以完全隐藏延迟,并且即使在 M=64 时也能容纳在共享内存中;(b) 因为它可以被2整除。后者至关重要,因为它允许我们在整个流水线中平滑地展开,因为在 P 次迭代后,流水线和寄存器缓冲区的索引都将始终为0。这种展开使得所有共享内存寻址完全静态,通过使用我们可用的额外寄存器避免了缓慢的变换索引计算。最后,这也似乎是让CUDA编译器正确排序指令以实现真正双缓冲的最可靠方法。图3展示了MARLIN核心中使用的多层流水线。
Warp布局。在单个SM上计算 $C_{sm}$ 必须进一步在warp之间进行细分:如果直接进行,每个warp将计算输出的一个 $M \times (N_{sm}/\#warps)$ 的瓦片。为了达到峰值计算吞吐量,我们希望至少使用四个(因为Ampere GPU有四个warp调度器),理想情况下是八个warp【【26】,Sun, W., Li, A., Geng, T., Stuijk, S., and Corporaal, H. Dissecting tensor cores via microbenchmarks: Latency, throughput and numeric behaviors. IEEE Transactions on Parallel and Distributed Systems, 34(1):246–261, 2022】,以获得额外的延迟隐藏。然而,这会导致瓦片尺寸变小,尤其是在较小的 $N_{sm}$ 下。这不仅对我们之前讨论的内存重排有问题,也阻碍了Tensor Core的吞吐量,因为小瓦片宽度会给tensorcore操作带来更多的顺序依赖(因为这些连续操作必须使用相同的累加器),可能导致停顿。因此,我们将每个warp的子瓦片宽度固定为64,并进一步在 $K_{sm}$ 上分割计算;图4说明了这样的warp布局,算法1提供了相应的伪代码。因此,多个warp将在寄存器中累加相同 $C_{wa}[i_{sm}, j_{sm}][i_{wa}, j_{wa}]$ 的部分结果。这些结果最终必须在共享内存中进行归约,然后才能最终写出。但这可以通过对数并行归约(logarithmic parallel reduction)【【12】,Harris, M. et al. Optimizing parallel reduction in cuda. Nvidia developer technology, 2007】来完成,这通常只会带来最小的开销。
去量化与Tensor Cores。从INT4到FP16的简单类型转换速度很慢;因此,我们采用了Kim等人【【13】,Kim, Y. J., Henry, R., Fahim, R., and Awadalla, H. H. Who says elephants can’t run: Bringing large scale moe models into cloud scale production. arXiv preprint arXiv:2211.10017, 2022】二进制操作的修改版本。以最简单的情况为例:将一个INT16中位于12-15位的INT4转换为一个有符号的FP16值。首先,我们通过AND掩码仅提取对应于我们INT4的位,并将结果的1-7位变为0110010(通过OR操作);这可以在单个lop3
指令中完成。现在,我们得到了一个指数为50的FP16数,其尾数的最后4位对应我们的转换目标。因此,减去一个指数为50、尾数为0的FP16值,将得到我们目标4位的无符号浮点表示。为了使该值有符号,我们还需减去8,但这可以直接融合到我们减去的总值的最后3位中。现代GPU可以同时用一个32位寄存器中封装的两个独立的16位操作数进行计算。因此,我们可以使用上述过程同时高效地去量化一个INT32中的两个INT4。最后,我们希望直接去量化到后续Tensor Core调用所需的正确寄存器布局。为此,我们再次利用B可以离线预处理的优势,重新组织权重,使得每个线程读取的16字节向量精确包含其所需的4个独立的 $16 \times 16$ Tensor Core块的8个量化权重。此外,在一个INT32内,权重按模式64207531交错存储,以支持前面提到的并行解码。在最内层,我们累加一个 $M \times 16$ 乘以 $16 \times 64$ 矩阵乘法的结果。我们按列执行这个累加,发出 $16 \times 16$ 乘以 $16 \times 8$ 的Tensor Core mma.sync
指令。这比按行执行的优势在于,我们可以将下一个B操作数的去量化与当前列的Tensor Core数学运算流水线化。
分组与指令排序。对于逐输出量化,我们可以在全局写出前简单地对最终输出进行一次缩放。一个有趣的观察是,尽管这些加载与任何计算都不是异步的,但通过 cp.async
后紧跟一个即时的wait group
指令来执行它们仍然至关重要,以避免编译器对主循环进行不利的指令重排。对于分组量化(这对于保持良好准确性至关重要),我们必须在主循环期间加载并应用缩放。首先,我们以与量化权重类似的方式重组缩放因子的存储,使得同一类型线程所需的不同 $16 \times 16$ 块的缩放因子被打包在一起,并可以作为一个16字节向量从共享内存加载。原则上,对于组大小为128和 $B_{sm}$ 瓦片形状为 $64 \times 256$ 的情况,我们只需要每隔一个瓦片从全局和共享内存加载一次新的缩放因子(并且在这里只在第一个子瓦片期间加载一次)。然而,编译器似乎对代码最关键部分中的这种不规则性相当脆弱,导致在某些形状设置下指令排序不佳,整体 slowdown 10-20%。相反,我们发现为每个子瓦片从共享内存重新加载缩放因子可以保持峰值性能。这样做增加了一些技术上不必要的共享内存加载,但有足够的额外带宽来支持这一点而没有开销,同时它也保留了编译器为非分组量化精心设计的流水线指令排序。
条带化分区。通过上述所有技术,我们可以在矩阵较大且可以在N轴上完美地在所有SM之间划分时,达到接近最优的计算和带宽性能。在实践中,这种情况很少见。在这种情况下,标准的补救措施是在K维度上也进行分区,但对于许多流行的层形状和GPU组合,我们需要大量额外的分割才能达到均匀分布而没有显著的波浪量化。这反过来又增加了许多全局归约步骤,带来了额外的开销。相反,我们选择了一种条带化分区方案(striped partitioning scheme),其中一个SM处理的B的“条带”可以跨越多个 $C_{sm}$ 瓦片(参见图5)。具体来说,我们首先确定每个SM要处理的 $B_{sm}$ 瓦片数量 $T = \lceil \#tiles / \#SMs \rceil$,然后从左上角开始按列分配(最多)T个瓦片。关键是,如果我们到达一个瓦片列的底部,但当前SM尚未拥有T个瓦片,我们将继续从下一个瓦片列的顶部开始分配瓦片;换句话说,条带可以跨越多列。这确保了瓦片在所有SM之间的分布大致均匀,同时最小化了所需全局归约步骤的数量。该策略类似于stream-k分区【【21】,Osama, M., Merrill, D., Cecka, C., Garland, M., and Owens, J. D. Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the GPU. In ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, 2023】。
全局归约实现。我们以串行方式实现同一瓦片列中条带之间的全局归约,从下到上。后一种方法最有效,因为在存在任何列溢出的情况下,最底部的SM将最快得到其结果,而最顶部的SM最慢。我们直接在输出缓冲区中以FP16执行归约,以最大化L2缓存命中率,从而最小化任何全局读取开销。这也使操作基本上是就地的,只需要一个小的额外锁缓冲区用于同步。最后,我们注意到,对于批次大小远大于64的情况,我们可以为条带索引计算虚拟地复制B,然后通过模运算移回原始矩阵,并将A指针前进到相应的大小为64的输入批次段。这为大型输入批次大小(如LLM预填充期间发生的情况)显著减少了全局归约,并在此设置下提高了计算吞吐量。
3.5 GPTQ 修改
对GPTQ的改进。MARLIN使用的量化格式,专为峰值推理效率而设计,与原始GPTQ实现【【10】,Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pretrained transformers. arXiv preprint arXiv:2210.17323, 2022】略有不同,但仍能产生高精度模型。我们还向GPTQ集成了两个小改进:(a) 通过搜索最优的组级裁剪阈值来选择组缩放因子,类似于【【15】,Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023】;(b) 支持可变长度的校准序列。我们发现这些修改相比标准GPTQ带来了微小但一致的准确性提升,同时具有更高性能的优势。(我们也提供了一个在模型格式之间进行简单转换的脚本。)图6展示了我们的GPTQ变体与原始未压缩模型在困惑度(越低越好)与模型大小(以比特为单位)方面的对比。这表明,MARLIN量化的模型在与未压缩基线相同的困惑度下,大小减少了约3.33倍。虽然这不是无损的(在此位宽和组大小下,理想增益为3.87倍),但这是一个显著的改进,特别是考虑到MARLIN的高推理效率。
4. The Sparse-MARLIN Kernel
设计动机。为了进一步提高FLOPS/Byte比率,我们可以在4位量化权重表示的基础上集成2:4稀疏方案。NVIDIA Ampere架构中的稀疏Tensor Cores (SPTCs)为在专门硬件单元上执行50%稀疏矩阵提供了有效手段。然而,要利用SPTCs,需要对先前描述的MARLIN核心进行一些修改和扩展。
矩阵乘法重构。首先,为了适应mma.sp
指令的约束(该指令启用SPTCs,并要求稀疏矩阵作为张量操作中的左侧(LHS)操作数【【18】,NVIDIA. Nvidia instruction set. https://docs.nvidia.com/cuda/ parallel-thread-execution/index.html# warp-level-matrix-instructions-for-sparse-mma, 2022b】),我们设计了新的特定数据布局。具体来说,将A与B相乘的问题在底层被重构为求解 $(B^T A^T)^T$ 来产生C。然而,这种重构保留了先前描述的密集MARLIN核心设计中的所有技术和优化。请注意,B可以根据需要离线预处理,而A可以通过ldmatrix
指令及其.trans
可选限定符在SMEM中即时转置,而不会导致性能下降。
新的数据结构。接下来,我们描述在Sparse-MARLIN中编码2:4稀疏矩阵所需的两个新数据结构,以及为解决此特定问题而进行的调整:(1)非零值结构,和(2)元数据索引结构。
量化的非零值。图7左侧展示了一个大小为 $N \times K$ 的4位量化矩阵B,该矩阵已被修剪为2:4稀疏。该矩阵的压缩表示(图7中1)在内部维度上的大小将是原始矩阵的一半,即 $N \times K/2$。然而,由于每个值是4位元素,我们可以在此基础上应用密集的MARLIN压缩方法,将8个元素进一步压缩到一个32位值中(图7中1b),最终大小为 $N/8 \times K/2$。为最大化内存效率,由于权重在推理期间保持静态,每个 $64 \times 16$ 的瓦片都被重新排列,以便每个线程在连续的内存位置加载和存储元素,类似于密集的MARLIN。继续看图7,非零值结构中的彩色元素表示线程T0为一个 $64 \times 16$ 的B块提供的元素示例。成对的颜色表示在同一个mma.sp
指令内处理的元素,需要4次迭代才能计算所有元素。这种布局在从GMEM加载非零值结构时,确保了最宽的128位指令(每个线程4个连续的32位元素,如1b所示)。此外,由于乘法的重新定义,并且输出C将是一个形状为 $M \times N$ 的FP16矩阵,这种布局在将结果从RF转置存储到GMEM时,也确保了128位指令(例如,T0寄存器中存储的第0列中前八个连续输出元素)。因此,这种乘法的重构不仅在不降低性能的情况下转置存储结果,而且与基线密集设计相比,进一步提高了输出写入的效率。
元数据索引。为了选择稀疏计算中将需要的右侧(RHS)操作数A中的元素,需要一个包含原始矩阵中非零元素索引的元数据结构。图8左侧展示了B的元数据索引结构。由于这是2:4稀疏性,索引将在0~3范围内,用2位编码。基于图7中描述的数据布局,并考虑到mma.sp
指令的稀疏性选择器约束,我们为元数据结构提出了一种新的数据布局。稀疏性选择器指示四个连续线程组中的哪个线程(们)将为整个组提供元数据。在图8所示的示例中,稀疏性选择器可以是0(线程T0, T1)或1(线程T2, T3)。首先,我们必须根据先前描述的非零值结构重新排序行,如2a所示。这使我们能够使用128位加载指令从GMEM到SMEM。然后,为了无bank冲突地将数据从SMEM加载到RF,我们执行单个ldmatrix
指令,该指令将包含要执行的下四个mma.sp
操作的所有信息,并根据需要将信息分布到线程T0~T3。然而,需要进行先前的数据重排,如2b所示。这样,线程T0, T1将包含前两次迭代的信息,而线程T2, T3将包含剩余两次迭代的信息。需要注意的是,所有这些预处理都是一次性离线完成的,没有运行时开销。
A4 实验环境
- 模型架构: 实验中使用了多个开源大型语言模型,包括Llama-2-7B【【29】,Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023b】和Falcon【【27】,TII UAE. The Falcon family of large language models. https://huggingface.co/tiiuae, May 2023】系列模型。
- 硬件配置:
- GPU: 实验在多种NVIDIA Ampere架构的GPU上进行,包括专为推理优化的NVIDIA A10,旗舰级的NVIDIA A100,消费级的NVIDIA GeForce RTX 3090,以及工作站级的NVIDIA RTX A6000。
- 多GPU设置: 对于大型模型,实验还测试了vLLM支持的跨多GPU分片权重矩阵的场景。
- 软件配置:
- 核心实现: MARLIN核心基于CUDA编程模型实现。
- 推理引擎: 将MARLIN核心集成到流行的开源LLM服务引擎vLLM【【14】,Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles, pp. 611–626, 2023】中进行端到端评估。
- 对比基线: 与多个流行的开源4位推理核心进行了比较,包括PyTorch核心【【22】,Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019】、AWQ核心【【15】,Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023】、ExLlamaV2核心【【8】,ExLlamaV2. Exllamav2: A memory efficient fork of hf transformers optimized for llama models. https:// github.com/turboderp/exllamav2, 2024】和bits-and-bytes核心【【5】,Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. LLM.int8(): 8-bit matrix multiplication for transformers at scale. Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, 2022】。
A4 实验结果
5.1 核心基准测试
- 峰值性能对比 (图1): 在NVIDIA A10 GPU上,针对一个72k x 18k的大矩阵,MARLIN在批量大小为1时达到接近理论最优的3.87倍加速。与其他核心性能随批量大小增加而迅速下降不同,MARLIN在批量大小达到16-32时仍能维持接近最优的加速,之后随着问题转为计算密集型而逐渐减弱。
- 真实层形状性能 (图9): 在实际模型(Llama-2, Falcon)的线性层上,批量大小为16时,MARLIN在不同GPU(RTX 3090, A10, A100)上均展现了强大的性能。在消费级GPU(RTX 3090)上加速效果更好,而在旗舰级A100上由于其更高的带宽和算力,开销相对更大,加速比较低。
- 持续性能 (图10): 在锁定的基础GPU时钟频率下(模拟生产环境),其他核心的相对加速效果显著下降,而MARLIN的性能几乎不受影响,表现出强大的稳定性。
- 大批量性能: 在A100上针对“预填充”场景的大批量(高达1024)测试中,MARLIN的性能与未压缩的计算密集型矩阵乘法几乎相同,仅在更大输入尺寸下有约10%的性能下降。
- Roofline分析 (图11): 在A10 GPU上,Roofline分析证实了MARLIN的计算效率。批量大小小于64时为内存密集型,大于64时为计算密集型。MARLIN在各种矩阵尺寸和算术强度下都实现了很高的硬件利用率。
- Sparse-MARLIN性能 (图12, 图13): Sparse-MARLIN在峰值和持续性能测试中均展现了比密集版MARLIN更高的性能,验证了MARLIN设计对其他压缩格式(如2:4稀疏)的可扩展性。
5.2 端到端实验
- 模型准确性 (表1): 通过GPTQ和SparseGPT【【9】,Frantar, E. and Alistarh, D. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning (ICML), 2023】生成的INT4和INT4+2:4稀疏模型在Llama-2-7B上表现出良好的准确性恢复。
- vLLM集成性能 (图14): 在vLLM中,与FP16基线相比,MARLIN在Llama-2-7B上的端到端生成速度提升高达3倍,而Sparse-MARLIN在此基础上还能提供额外的1.2倍加速。
- 不同GPU和模型下的性能 (表2): 在多种GPU和模型(包括多GPU分片)的组合下,MARLIN均能提升性能。当推理是内存密集型(批量大小~16)且GPU较弱或数量较少时,加速效果最显著,表明MARLIN在资源受限的环境中尤其有益。
- 服务基准测试 (图15, 图16):
- TPOT (输出令牌时间): 在模拟的服务器-客户端设置中,随着QPS(每秒查询数)增加,MARLIN相比FP16基线实现了约2.8倍的延迟降低,Sparse-MARLIN则达到3.3倍。有趣的是,随着QPS增加,相对加速比也在增加,因为MARLIN的低延迟导致其在较低的平均批量大小下运行,从而放大了优势。
- TTFT (首个令牌时间): MARLIN在考虑提示处理的情况下,也能带来显著的TTFT改进。
A7 补充细节
6. 相关工作
LLM定点化推理支持。本文主要关注为量化LLM推理提供高效支持的相关工作。先前已有大量关于LLM权重精确量化的研究,流行的方法包括GPTQ【【10】,Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pretrained transformers. arXiv preprint arXiv:2210.17323, 2022】和AWQ【【15】,Lin, J., Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978, 2023】,以及对四舍五入(RTN)量化的探索【【4】,Dettmers, T. and Zettlemoyer, L. The case for 4-bit precision: k-bit inference scaling laws. arXiv preprint arXiv:2212.09720, 2022】,后者通常精度较低。MARLIN的并行化方法可以推广到这些量化方法。事实上,自本文核心发布GPTQ格式版本以来,vLLM中已独立引入了支持AWQ的MARLIN版本【【30】,vLLM Project Contributors. vllm project pull request #6612: Add support for awq marlin. https://github.com/ vllm-project/vllm/pull/6612, 2024】。
权重与激活值同时量化。更广泛地看,LLM量化方法也可以考虑同时压缩权重和激活值【【5】,Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. LLM.int8(): 8-bit matrix multiplication for transformers at scale. Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, 2022】,高级方法如SmoothQuant【【31】,Xiao, G., Lin, J., Seznec, M., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211.10438, 2022】或QuaRot【【1】,Ashkboos, S., Mohtashami, A., Croci, M. L., Li, B., Jaggi, M., Alistarh, D., Hoefler, T., and Hensman, J. Quarot: Outlier-free 4-bit inference in rotated llms. arXiv preprint arXiv:2404.00456, 2024】。然而,激活值的量化往往更复杂,因为会出现大的“离群值”【【5】,Dettmers, T., Lewis, M., Belkada, Y., and Zettlemoyer, L. LLM.int8(): 8-bit matrix multiplication for transformers at scale. Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, 2022】。因此,这些方法要么针对较低的8位精度,要么需要更复杂的额外处理步骤,例如通过旋转矩阵【【1】,Ashkboos, S., Mohtashami, A., Croci, M. L., Li, B., Jaggi, M., Alistarh, D., Hoefler, T., and Hensman, J. Quarot: Outlier-free 4-bit inference in rotated llms. arXiv preprint arXiv:2404.00456, 2024】。MARLIN的方法可以扩展到这种情况,例如,最近一个独立的后续工作将MARLIN的方法扩展到了激活值量化为8位,而权重量化为4位的情况【【33】,Zhang, Y., Zhang, P., Huang, M., Xiang, J., Wang, Y., Wang, C., Zhang, Y., Yu, L., Liu, C., and Lin, W. Qqq: Quality quattuor-bit quantization for large language models, 2024】。
A5 结论
本文提出了MARLIN,一种用于实现LLM生成式推理的混合精度核心的通用方法。该方法通过利用新的GPU硬件指令和并行化技术,实现了接近最优的效率。具体而言,MARLIN及其稀疏版本在单层上达到了接近最优的效率,并在实际部署场景中,以适度的精度影响为代价,带来了高达3倍的加速。
未来工作展望:
1. 支持更极端的压缩技术: 一个自然的方向是研究MARLIN对近期提出的更复杂的“极端”压缩技术的支持,例如向量量化【【3】,Chee, J., Cai, Y., Kuleshov, V., and Sa, C. D. Quip: 2-bit quantization of large language models with guarantees, 2023】和【【7】,Egiazarian, V., Panferov, A., Kuznedelev, D., Frantar, E., Babenko, A., and Alistarh, D. Extreme compression of large language models via additive quantization. arXiv preprint arXiv:2401.06118, 2024】,这些技术需要更复杂的解压过程。
2. 支持其他形式的混合精度: 另一个方向是研究MARLIN对其他形式混合精度的支持,例如由激活值压缩或稀疏性产生的混合精度。
方法细节中引用的参考文献汇总
- 【1】Ashkboos et al., Quarot: Outlier-free 4-bit inference in rotated llms, arXiv, 2024.
- 引用段落:A7. 相关工作
- 引用内容:作为一种更复杂的激活值量化方法,通过旋转矩阵处理离群值。
- 【3】Chee et al., Quip: 2-bit quantization of large language models with guarantees, 2023.
- 引用段落:A5. 结论
- 引用内容:作为未来工作方向,提到MARLIN可扩展支持向量量化等极端压缩技术。
- 【4】Dettmers & Zettlemoyer, The case for 4-bit precision: k-bit inference scaling laws, arXiv, 2022.
- 引用段落:A7. 相关工作
- 引用内容:作为一种准确性较低的量化方法(四舍五入,RTN)。
- 【5】Dettmers et al., LLM.int8(): 8-bit matrix multiplication for transformers at scale, NeurIPS, 2022.
- 引用段落:A4. 实验环境,A7. 相关工作
- 引用内容:作为对比的开源核心之一,并指出同时量化权重和激活值的方法,以及激活值中存在“离群值”的挑战。
- 【7】Egiazarian et al., Extreme compression of large language models via additive quantization, arXiv, 2024.
- 引用段落:A5. 结论
- 引用内容:作为未来工作方向,提到MARLIN可扩展支持加性量化等极端压缩技术。
- 【8】ExLlamaV2, Exllamav2: A memory efficient fork of hf transformers optimized for llama models, 2024.
- 引用段落:A4. 实验环境
- 引用内容:作为对比的开源4位推理核心之一。
- 【9】Frantar & Alistarh, Sparsegpt: Massive language models can be accurately pruned in one-shot, ICML, 2023.
- 引用段落:A4. 实验结果
- 引用内容:用于生成Sparse-MARLIN所使用的INT4 + 2:4稀疏模型。
- 【10】Frantar et al., Gptq: Accurate post-training quantization for generative pretrained transformers, arXiv, 2022.
- 引用段落:A3. 背景知识,A2. 方法细节,A7. 相关工作
- 引用内容:作为本文采用的量化算法的基础,并作为一种流行的权重精确量化方法。
- 【12】Harris et al., Optimizing parallel reduction in cuda, Nvidia developer technology, 2007.
- 引用段落:A2. 方法细节
- 引用内容:提到在warp布局中,多个warp累加的部分结果可以通过对数并行归约(logarithmic parallel reduction)高效地合并。
- 【13】Kim et al., Who says elephants can’t run: Bringing large scale moe models into cloud scale production, arXiv, 2022.
- 引用段落:A2. 方法细节
- 引用内容:MARLIN的去量化过程采用了该文二进制操作的修改版本,以实现从INT4到FP16的高效转换。
- 【14】Kwon et al., Efficient memory management for large language model serving with pagedattention, SOSP, 2023.
- 引用段落:A4. 实验环境
- 引用内容:作为MARLIN集成的LLM服务引擎。
- 【15】Lin et al., Awq: Activation-aware weight quantization for llm compression and acceleration, arXiv, 2023.
- 引用段落:A2. 方法细节,A4. 实验环境,A7. 相关工作
- 引用内容:作为对比的开源核心之一,一种流行的权重精确量化方法,其组缩放因子的选择方法被MARLIN借鉴和改进。
- 【17】NVIDIA, Nvidia a10 datasheet, 2022a.
- 引用段落:A2. 方法细节
- 引用内容:引用A10 GPU的FLOPs/Bytes比率(约200)作为MARLIN设计的动机。
- 【18】NVIDIA, Nvidia instruction set, 2022b.
- 引用段落:A2. 方法细节
- 引用内容:引用
mma.sp
指令要求稀疏矩阵作为LHS操作数的约束,这是Sparse-MARLIN重构矩阵乘法的原因。
- 【19】NVIDIA, Efficient GEMM in CUDA, 2024a.
- 引用段落:A2. 方法细节
- 引用内容:MARLIN的设计紧密遵循其分层并行化模型,并引用了其中双缓冲(double buffering)的技术。
- 【20】NVIDIA, CUTLASS convolution, 2024b.
- 引用段落:A2. 方法细节
- 引用内容:引用了为避免共享内存bank冲突而采用的XOR索引变换技术。
- 【21】Osama et al., Stream-k: Work-centric parallel decomposition for dense matrix-matrix multiplication on the GPU, PPoPP, 2023.
- 引用段落:A2. 方法细节
- 引用内容:提到MARLIN的条带化分区方案与stream-k分区策略相似。
- 【22】Paszke et al., Pytorch: An imperative style, high-performance deep learning library, NeurIPS, 2019.
- 引用段落:A4. 实验环境
- 引用内容:作为对比的开源4位推理核心之一。
- 【26】Sun et al., Dissecting tensor cores via microbenchmarks: Latency, throughput and numeric behaviors, TPDS, 2022.
- 引用段落:A2. 方法细节
- 引用内容:引用其观点,即理想情况下应使用八个warp以获得额外的延迟隐藏。
- 【27】TII UAE, The Falcon family of large language models, 2023.
- 引用段落:A4. 实验环境
- 引用内容:作为实验中使用的模型之一。
- 【29】Touvron et al., Llama 2: Open foundation and finetuned chat models, arXiv, 2023b.
- 引用段落:A4. 实验环境
- 引用内容:作为实验中使用的模型之一。
- 【30】vLLM Project Contributors, vllm project pull request #6612: Add support for awq marlin, 2024.
- 引用段落:A7. 相关工作
- 引用内容:证明MARLIN的方法已被社区独立扩展以支持AWQ格式。
- 【31】Xiao et al., Smoothquant: Accurate and efficient post-training quantization for large language models, arXiv, 2022.
- 引用段落:A7. 相关工作
- 引用内容:作为一种同时量化权重和激活值的高级方法。
- 【33】Zhang et al., Qqq: Quality quattuor-bit quantization for large language models, 2024.
- 引用段落:A7. 相关工作
- 引用内容:证明MARLIN的方法可以扩展到W4A8(权重4位,激活8位)的量化场景。
💬 评论讨论
欢迎在这里分享您的想法和见解!