37 KiB
通过以 FP6 为中心的算法系统协同设计高效服务大语言模型
六位量化(FP6)可以有效地减小大型语言模型(LLMs)的大小,并在不同的应用程序中保持一致的模型质量。
然而,现有系统不提供对 FP6 量化的 Tensor Core 支持,并且在 LLM 推理期间难以实现实际性能改进。在 GPU 上支持 FP6 量化具有挑战性,因为
(1) 位宽不规则的模型权重的内存访问不友好,
以及 (2) 权重反量化的运行时开销较高。
为了解决这些问题,我们提出了TC-FPx,这是第一个全栈GPU内核设计方案,具有统一的Tensor Core支持各种量化位宽的浮点权重。
我们将 TC-FPx 内核集成到现有的推理系统中,为量化 LLM 推理提供新的端到端支持(称为 FP6-LLM),其中推理之间有更好的权衡实现了成本和模型质量。
实验表明,FP6-LLM 仅使用单个 GPU 即可实现 LLaMA-70b 的推理,实现比 FP16 基线高 1.69×- 2.65× 的归一化推理吞吐量。
部署 LLMs 具有挑战性
一方面,它需要大量的 GPU 内存(FP16 中的 GPT-3 需要 326 GB)才能容纳模型权重,而 A100/H100 GPU [20, 21] 仅具有最多 80 GB 内存。
另一方面,LLM推理在token生成过程中面临严重的“内存墙”问题[12, 36],其中LLM推理的速度主要受到读取模型权重的时间限制来自 GPU DRAM。它使 LLM 推理内存受到限制,无法充分利用 GPU 的计算能力。
最近的研究表明,6 位量化是 LLM 部署的推理成本和模型质量之间的良好权衡 [30, 35]。 然而,现代 GPU 上仍然没有有效的系统支持 6 位线性层执行(即矩阵乘法)。迫切需要开发支持6位量化的系统,充分利用GPU的计算能力。
一方面,与较大位量化(例如 8 位)相比,6 位量化可以实现更高效的 LLM 推理。
首先,可以节省更多的 GPU 内存,例如如果部署具有 6 位而不是 8 位量化的 GPT-3 模型,可以节省大约 40 GB 内存。 其次,由于可以有效减少从GPU DRAM读取模型权重的时间,因此可以进一步加速LLM推理。如图 1 所示,使用我们新提出的 6 位量化系统设计 (TC-FPx_W6A16) 实现的线性层始终更快(高达 1.45 倍)。
使用我们新提出的 6 位量化系统设计 (TC-FPx_W6A16) 实现的线性层始终比最先进的 8 位量化支持 (TensorRT_LLM_W8A16) 快 1。
另一方面,尽管最近的 4 位技术 [6,39] 在以最小的压缩 LLMs 方面表现出了潜力,但与较小位量化(例如 4 位)相比,6 位量化可以更有效地保持模型质量。
最近的研究 [35] 表明,在超出零样本测量的任务中,例如代码生成和总结,4 位方法表现不佳且缺乏鲁棒性,而 6 位量化方法则主要针对零样本评估。在这些不同的应用程序中显示出强大且一致的性能。
受上述观察的启发,我们提出了 TC-FPx,这是第一个全栈 GPU 系统设计方案,具有统一的 Tensor Core [20,21] 支持各种量化位宽(6 位、5 位)的浮点权重、3 位等),减轻 LLM 推理过程中的“内存墙”问题。
TC-FPx打破了底层GPU硬件的限制,让GPU支持涉及任意位宽的模型权重的线性层计算。
在 TC-FPx 中,Tensor Core 用于矩阵乘法的密集计算,而 SIMT 核则有效地用于权重反量化,在运行时将 x 位模型权重转换为 FP16 类型,然后再将其馈送到 Tensor Core。
我们提出提前位级预打包(第 5.2 节)来解决不规则位宽权重的不友好内存访问的挑战(第 4.2.1 节),从而实现最佳 GPU 内存访问。
此外,我们提出 SIMT-Efficient GPU Runtime(第 5.3 节)来最小化权重反量化的运行时开销(第 4.2.2 节)。
最后但并非最不重要的一点是,我们展示了 TC-FPx 内核的软件管道,其中 SIMT 内核、Tensor 内核和 GPU 内存层次结构以高性能高效协作。
我们将 TC-FPx 内核集成到最先进的推理系统 [19] 中,为量化 LLM 推理提供新的端到端支持(称为 FP6LLM),其中更好地权衡实现了推理成本和模型质量。
目前,FP6-LLM主要支持流行的LLMs的6位量化(FP6),例如各种尺寸的LLaMA [32]、OPT [41]。
评估表明,FP6-LLM 仅使用单个 GPU 即可实现 LLaMA-70b 的推理,实现比 FP16 基线高 1.69×2.65 倍的归一化推理吞吐量。此外,FP6-LLM将OPT-30b的推理吞吐量提高了1.72×-4.05×。
我们确定了在现代 GPU 上支持 FP6 量化的重要性和主要挑战。
我们提出了TC-FPx,这是第一个全栈GPU内核设计方案,具有统一的Tensor Core支持各种位宽的浮点权重,例如FP6。
我们通过TC-FPx的集成为量化LLMs提供新的端到端推理支持,在推理成本和模型质量之间实现更好的权衡。
我们在各种 LLM 模型上评估 FP6-LLM 并证明它大大优于基线。
量化的主要目标是线性层的权重(即矩阵乘法),占总体LLM权重的99%以上。激活也可以在推理过程中进行量化。
SIMT cores (CUDA cores) 负责 GPU 中的通用处理任务,处理各种指令,包括整数运算、浮点运算、加载/存储操作等。
SIMT cores 执行在单个(或矢量)数据元素。
Tensor Cores [20, 21] 是为加速矩阵乘法而设计的专用硬件。在 A100 [20]/H100 [21] GPU 上,Tensor Cores 的 FLOPS 比 SIMT 核心高 16.0 倍/ 14.8 倍。
此外,Tensor Cores 以粗粒度工作,例如:使用单个 mma(矩阵乘法和累加)指令在形状为 16 × 16 和 16 × 8 的两个 FP16 矩阵之间执行矩阵乘法。
动机:
与 8 位和 4 位量化相比,FP6 量化可以在推理成本和模型质量之间实现更好的权衡。
(I) 推理成本低于 8 位量化。与 8 位量化相比,通过更积极的 6 位量化可以进一步降低部署 LLMs 的成本,而不会出现明显的精度下降。 一方面,LLM权重的大小可以显着减小,比 FP16 基线小近 2.7 倍。存储模型权重需要更少的 GPU 内存,从而需要更少的 GPU 并降低部署 LLMs 的服务成本。
另一方面,6位量化也可以更有效地加速LLMs的推理。鉴于 LLM 推理在令牌生成期间通常受内存限制,因此可以通过减少模型权重的 GPU DRAM 访问来实现更快的 LLM 推理。
如图 1 所示,
与最先进的 8 位量化支持(TensorRT-LLM_W8A16 )相比,我们新提出的 6 位量化系统设计 (TC-FPx_W6A16) 中 llama-65b 模型 [32] 中线性层的执行速度始终更快(高达 1.42 倍)。 [26])。鉴于线性层是大型语言模型中最耗时的部分,这种加速将直接转化为端到端推理场景的性能改进(参见第 7.3 节)。
(II) 比 4 位量化更好的模型质量。
虽然 4 位量化更积极地减少内存占用和 DRAM 访问,但它不可避免地会导致模型质量下降。相比之下,通过 6 位量化可以实现近乎无损的模型压缩。如表 1 和表 2 所示,FP6 在各种任务中显示出强大且一致的性能,包括代码生成和零样本困惑性能。它还在各种模型尺寸(例如 1B、13B 和 65B LLaMA [32] 模型)中表现出高稳健性。
我们还发现 INT4 量化严重依赖细粒度量化 (FGQ) 方法来保持较高的模型质量,而我们的 FP6 量化已经在粗粒度量化上表现良好。请注意,表 1 和表 2 中的数据点选自[35]。欲了解更多详细信息,请参阅本文。总之,在算法层面,FP6 量化是一种实用的替代方案,可以进一步民主化 LLMs 的部署,而不会显着牺牲复杂任务和各种模型大小的模型质量。
设计的选择与挑战
尽管对训练后 FP6 量化的高性能支持的需求不断增加,但目前还没有这样有效的以 FP6 为中心的系统设计,可以实现上述 4 位和 8 位量化的权衡。具体来说,现有对线性层的支持主要是针对位宽为2的指数的数据类型(例如4位、8位和16位)而设计的。鉴于目前尚不清楚如何在现代 GPU 上有效支持 FP6,我们在本节中说明两个重要的设计选择。
启用 Tensor Core 的必要性。
我们发现有必要在执行量化 LLMs 推理时支持 Tensor Core。例如,我们评估了 AWQ [14, 15] 纯 SIMT 核心执行在各种批量大小上的性能,以测试其可扩展性。如图 1 所示,随着推理批量大小的增加,不带 Tensor Core 支持的线性层 (AWQ_W4A16_SIMT) 的运行时性能变得极低。这背后的原因是双重的。
一方面,对于线性层执行,传统 SIMT 核心比 Tensor Core 慢一个数量级,如第 2.3 节所述。
另一方面,SIMT 核心计算能力的很大一部分将用于在运行时对模型权重进行反量化,这进一步降低了 SIMT 核心用于计算矩阵乘法的可用计算能力。
这促使我们启用Tensor Core来进行矩阵乘法的密集计算,同时利用多功能 SIMT core 进行权重反量化。
统一kernel而不是双kernel的解决方案
这WxA16量化的独特之处在于激活矩阵使用FP16,但权重矩阵以较窄的位宽存储。但是,Tensor Core 要求将权重矩阵和激活矩阵存储在同一数据类型中,例如FP16/INT8/INT4。
简单的解决方案(即双kernel解决方案)添加一个额外的 GPU kernel,在调用普通 FP16 kernel之前将权重反量化为 FP16。然而,这样的推理速度会比没有量化的模型还要慢。如图 2(左)所示,将启动两个 GPU kernel用于线性层执行,反量化的 FP16 权重将再由第二个 GPU Kernel 读取之前写入 GPU DRAM,从而导致 2 倍 DRAM 访问。
将反量化和矩阵乘法过程融合到单个 GPU kernel中更加高效,从而消除了反量化权重的读/写(FP16 中的 W')。
设计挑战
在现代 GPU 上设计支持 FP6×FP16 矩阵乘法的统一 GPU Kernel 具有挑战性。一方面,现代 GPU 内存系统天然不支持不规则位宽(不是 2 的指数),因为 GPU 全局/共享内存的最小访问大小是每个线程 8/32 位,并且要访问的内存地址必须对齐。并且 Tensor Core 复杂的数据布局要求使得不规则位宽更具挑战性。另一方面,反量化计算成本昂贵,因为它需要大量复杂的bit-level操作。因此,如何将反量化融合到线性层计算中而不影响整体性能也很重要。
4.2.1 硬件不友好的内存访问
在现代 GPU 上执行线性层期间,应先将模型权重从 DRAM 加载到寄存器,然后才能执行相应的乘法计算。
通常,模型权重分两步加载,以隐藏 DRAM 的高访问延迟以获得高性能。具体来说,模型权重首先从 GPU DRAM 加载并缓冲到片上存储(例如:共享存储)中以供数据重用。之后,缓存的权重从共享内存读取到寄存器以进行实际计算。
鉴于每个GPU线程不能直接访问其他GPU线程的寄存器,每个线程必须将自己需要的模型权重放入自己的私有寄存器中。
考虑到 Tensor Core 严格的数据布局要求,当权重以不规则位宽(不是 2 位,例如 6 位)存储时,这个过程可能会变得极具挑战性。
如图 3a 所示,FP16 Tensor Core 的最小输入是现代 GPU 架构中的 8 × 8 子矩阵,每个 GPU 线程应在其寄存器中保存一对权重。
正常情况下,每个权重以16位存储,每对权重可以以32位字的粒度自然地从共享内存中读取。
然而,在我们的工作中,每个权重都以 x 位存储,这使得内存访问对于现代 GPU 内存层次结构极其不友好。
具有未使用位的片上存储访问:
我们以 6 位量化为例,展示访问不规则位宽权重的低效率。
如图3b所示,权重已经缓冲在共享内存中,每个GPU线程需要从共享内存中读取一对权重(12位,2 * 6位)。然而,共享内存有 32 个内存组,每个内存组在 GPU 上的每个内存请求上输出一个 32 位字。
因此,从共享存储读取的大部分位将未被使用,导致共享存储带宽显著浪费。例如,图 3b 中的 T0(线程 #0)仅需要 12 位。然而,将读取 32 位字 (W1),导致 32 位中的 20 位 (62.5%) 未使用。
由于现代 GPU 内存层次结构中,有对齐内存访问的要求,未使用位的浪费可能会变得更加严重。如图 3b 所示,T2(线程#2)所需的位分布在两个 W1 和 W2 中。
因此,T2 需要读取 W1 和 W2,从共享内存中读取 2 * 32 位。然而,最终只会使用 6 * 2 位,导致 64 位中有 52 位(81.25%)未使用和浪费。
还值得注意的是,由于位宽不规则,GPU DRAM 和寄存器上的内存管理和访问也遇到类似的问题。
4.2.2 反量化计算开销高
FPx-FP16 反量化的运行时开销可能非常高,这很容易减慢整体执行速度。
一方面,大量的模型权重需要在运行时进行反量化,例如对于 LLaMA-70b [33] 推理,每个 LLM 解码步骤应反量化 700 亿个 FPx 权重。
另一方面,反量化每个 FPx 权重的运行时开销很高,需要复杂的按位操作。根据公式 2,新的符号、指数和尾数都需要在运行时计算,以获得与给定 FPx 等效的 FP16。
FPx-FP16 反量化的运行时开销可能非常高,这很容易减慢整体执行速度。一方面,大量的模型权重需要在运行时进行反量化,例如: 对于 LLaMA-70b 推理,每个 LLM 解码步骤应反量化 700 亿个 FPx 权重。
另一方面,反量化每个 FPx 权重的运行时开销很高,需要复杂的按位操作。根据公式 2,新的符号、指数和尾数都需要在运行时计算,以获得与给定 FPx 等效的 FP16。
在公式2中,bias = 15,bias = 2− 1。
FP16的符号字段与FPx的符号字段相同,并且FP16的尾数也可以通过向FPx的尾数补零来计算。更重要的是,FP16 的指数应该是 E=E+bias−bias ,这在计算上更加昂贵。
综上所述,如何有效地对 FPx 值进行反量化也成为一个重大挑战。
- 设计方法论
在本节中,我们首先在 5.1 节中概述我们的设计。为了解决不友好的内存访问(第 4.2.1 节)的挑战, 我们在第 5.2 节中提出了运行前 Bit-level 预包装。 为了应对反量化高计算开销的挑战(第 4.2.2 节),我们在第 5.3 节中介绍了实现 SIMT 高效 GPU 运行时的设计。
最后,我们在第 5.4 节中介绍了我们的软件管道设计,其中 SIMT 核心、Tensor 核心和 GPU 内存层次结构以最佳性能协同工作。
图 4 比较了 TC-FPx(我们设计中的仅 x 位权重量化线性层内核)与通用矩阵乘法 (GEMM) 的传统设计,其中两个输入矩阵均采用 FP16。 TC-FPx 的模型权重以减少的位数存储。因此,在寄存器级别引入了额外的反量化阶段 (Dequant W),其中使用 SIMT cores 在每个线程内将 FP6 权重本地反量化为 FP16。
值得注意的是
FP16 权重不会写回共享内存,而是存储在寄存器中以供将来使用,从而消除了对共享内存不必要的往返访问。
另一个区别是TC-FPx使用细粒度的lds(加载共享)指令将x位权重从共享内存加载到寄存器,而不是使用粗粒度的固有ldmatrix(加载矩阵),后者具有严格的布局要求并且更少灵活性。
5.2 运行前 Bit-level 预包装
如 4.2.1 节所述,对具有不规则位宽的权重的内存访问对于现代 GPU 内存层次结构是不友好的。
为了解决这个问题,我们提出可以将每 32 个 x 位权重的内存读取结合起来,从而产生每个 GPU 线程 4 字节字的 x 请求。
可以合并每 32 个 x 位权重的内存读取,从而产生每个 GPU 线程 4 字节字的 x 请求。
在这种情况下,所有内存访问都将以 32 位字的粒度对齐,而不是不规则的位宽。
然而,由于 Tensor Core 严格的数据布局要求,将权重的内存读取结合起来并非易事,因为每个 GPU 线程所需的权重并不存储在连续的内存空间中。
为了解决这个问题,我们建议通过重新排序每个权重矩阵中的权重并提前预打包权重来优化运行时内存访问模式。
由于模型权重是在模型训练和量化后静态确定的,因此可以提前对权重应用复杂的内存布局转换,从而不会引入运行时开销。
此外,我们只需要预打包一次权重,因此预打包权重的开销可以通过每次推理服务有效摊销,并且可以忽略不计。
一般来说,权重预包装包括两个步骤。
第一步,我们收集每个 GPU 线程所需的所有权重,并在本地组合这些权重。鉴于每个 GPU 线程所需的权重最初并不位于每个权重矩阵内的连续位置(参见图 3a),因此我们必须仔细选择每个 GPU 线程的权重。
然后,为每个线程选择的权重在本地以相对时间顺序组合,因为它们在运行时被Tensor Core 消耗。
第二步,我们将整个GPU WARP(由32个GPU线程组成)所需的所有权重组合到一个统一的线性内存空间中,权重将在运行时按顺序存储在GPU DRAM中。为了完全消除共享内存库冲突,我们建议以锯齿状顺序组合每个线程的 32 位字。
值得注意的是,本小节中讨论的所有技术都独立于模型权重的实际位宽(始终使用 x 表示)。 因此,我们的权重预打包可以自然地应用于任何位宽度。
步骤 1:每线程权重收集
图 5 演示 规定 T0(线程 #0)选取的权重以及组合它们的顺序。
我们假设 WARP 级切片大小为 64 × 64 ,这意味着每个权重矩阵被分为 64 × 64 数据块,并以每个 WARP 的粒度加载到 GPU 的共享内存。
然后,每个权重图块进一步分为四个片,因为权重是从共享内存加载的,并逐片用于 Tensor Core 计算。
更重要的是,每个切片被分为四个 16 × 16 块,因为 Tensor Core 在每条指令中处理 16 × 16 数据项。
在每个 16 × 16 块中,为 T0 选择四对 FPx 权重并组合。
如图5所示,经过步骤1,我们得到了32组(即WARP大小)FPx权重。权重在每组中组合并连续存储,并且每组权重将被某个GPU线程消耗。
综上所述,每个 64 × 64 权重图块最终分配给 32 个线程(一个 WARP),每个线程将消耗 128 个 x 位权重。
步骤 2:按 WARP 进行 Bit-level 组装
在步骤 2 中,我们将不同组的所有权重组装成统一的内存空间。
在这个Bit-level 预打包过程中,我们将组合的权重视为要复制的连续数据,暂时忽略每个位的含义。具体来说,x 位的 128 个项目被视为 32 位的 4x 个项目。
我们建议按照图 5 所示的锯齿状顺序组装所有组的权重。
首先,将每个线程的第一个 32 位项连接在一起。之后,每个线程的第二个 32 位项被连接并附加到之前的结果中。
通过重复这个过程,所有权重可以连续存储在线性内存空间中,并且对齐良好(128字节对齐)。
这样,所有权重就可以以 128 字节块的粒度简单地从 DRAM 复制到共享内存,无需任何更改,轻松实现最佳 DRAM 访问。
此外,这些权重可以在运行时以最佳性能从共享内存加载。
具体来说,线程的 WARP 将为每个内存请求读取共享内存中的连续 32 位项目,完全避免了存储体冲突。
运行时高效的SIMT计算
SIMT-高效 GPU 运行时
并行反量化为了减少 FP-x 权重反量化的运行时开销,我们使用优化的按位 SIMT 核心指令实现了 FP-x 反量化。
此外,我们建议并行对多个 FPx 权重进行反量化,通过利用每个 32 位寄存器内的位级并行性,进一步将 SIMT 开销减少 4 倍。
(1) 优化的按位运算:如第 4.2.2 节所述,当将 FPx 转换为等效的 FP16 时,FP16 的指数应为 E = E +bias −bias 。为了简化这个过程,我们采用了[35]中的数学变换,用E = E来计算FP16的指数。为了保持正确性,将结果 FP16 与 FP16 常数 2 相乘:
图 6a 显示了优化的 FP16 到 FP6 转换。虽然我们只绘制从 FP6 到 FP16 的转换来进行演示,但它可以应用于任何位宽。 FP16 的符号字段与 FPx 的符号字段相同。此外,为了提高效率,指数字段的低位和尾数字段的高位可以一起从 FPx 复制到 FP16。此外,FP16 的其他位应该用零填充。 经过精心设计,我们仅用两个按位“与”、一个“移位”和一个“或”就成功实现了从 FP6 到 FP16 的转换,如图 6b 中❶所示。使用第一个“和”将符号字段从 FP6 复制到 FP16,同时将 FP16 的所有其他位初始化为零,从而无需稍后将零填充到指数和尾数字段。然后,FP6 的所有位都通过逐位“右移”右移。之后,首先通过以下方式选择 FP6 中指数的低位和尾数的高位:
之后在 FP6 和位掩码“0x1f1f1f1f”之间进行“与”,然后通过按位运算“或”复制到 FP16。
(2) 位级并行性:考虑到我们可以利用每个 32 位字内的位级并行性,我们建议并行地对多个 FPx 权重进行反量化,进一步减少反量化的运行时开销。
图 6b 以 FP6 为例演示了详细设计。 32 位寄存器被视为四个处理槽,其中每个槽独立工作,使用相同的指令但输入不同的 FP6 数据。在开始去量化之前,应将四个 FP6 存储在 R1(寄存器#1)中,初始数据布局如图所示。通过代码片段❶,这四个FP6可以同时反量化为四个FP16,其中每个FP16仅前8位存储在R2中。之后,将第一个和第二个 FP16 提取到 R1,最后 8 位补零,即代码片段❷。最后,通过代码片段❸和❹,将第三个和第四个FP16提取到R2中。
权重分割和缝合
然后我们将演示在 GPU 上通过精心设计的内存布局从 2+4 方案 [35] 有效重建 6 位权重的方法,该方法也可以应用于其他位宽。
(1)提前权重分割:为了将权重以良好对齐的方式存储在GPU的32位寄存器中,我们将每个权重分割成几个段,其中每个段的位宽为2,例如每个 6 位权重可以分为 2+4 或 4+2。基于该方案,后续设计的指标计算得到显着简化。
请注意,第 5.2 节中描述的技术可以应用于任何位宽度,因此可以根据第 5.2 节分别有效地预打包 2 位和 4 位段。
(2)运行时权重拼接:
在反量化之前,首先将权重从共享内存加载到寄存器。由于每个权重被分成几个段,因此需要在运行时在寄存器级别重建完整的权重。为了减少运行时开销,我们建议并行提取和拼接权重。如图7所示,使用两组寄存器来存储32个FP6权重,
其中,Frag1_PT R 指向包含 32 个 2 位段的两个 32 位寄存器,而 Frag2_PT R 指向包含 32 个 4 位段的 4 个 32 位寄存器。通过我们的并行拼接,可以同时重建四个 FP6 权重,从而将 SIMT 核心指令的数量减少 4 倍。如图 7 所示,首先将四个 2 位段提取到寄存器 #1 (❶),然后将四个 4 位段提取到寄存器 #2 (❷)。之后,寄存器#2 右移 (❸),并将其有效位复制到寄存器 #1 (❹),从而得到完整的 6 位权重。 (3)位重新排序:为了并行提取和缝合权重,有必要强制执行图7中的初始数据布局。关键的观察是每四个连续段必须按图中所示的顺序放置,例如前四个段必须按#2、#4、#1 和#3 的顺序存储。此外,每对 2/4 位段之间的步长应分别为 6/4。否则,不可能仅用四个SIMT核心指令同时拼接四个段。为了满足图 7 中的初始数据布局要求,我们建议通过在运行时之前重新排序权重段来确保这种布局,而无需运行时开销。此外,该技术应该作为附加通道叠加在第 5.2 节中描述的技术上。
总体伪代码算法 1 显示了包括并行去量化和权重拼接的伪代码(GPU 代码)。伪代码中的所有输入和输出变量都存储在寄存器中。如图 7 所示,算法 1 总共对 32 个 FP6 权重进行反量化。对于每个外循环,会生成四个 FP16 权重,并将其存储在代码末尾的两个寄存器中。图 7 中的转换(❶、❷、❸ 和 ❹)分别是通过算法 1 中第 6、7、9 和 10 行的 SIMT 核心操作实现的。然后,Tensor Core 直接使用输出寄存器数组 (OutputReg) 作为输入。
5.4 软件管道设计
为了减少 GPU 寄存器的使用,我们逐片(slice)对权重进行反量化。此外,我们将反量化过程无缝融合到线性层执行的传统软件管道中,通过有效的指令并行性完全隐藏了反量化的运行时开销。
逐片反量化
我们不是一次反量化所有权重,而是逐片反量化 FPx 权重。如图 8a 所示,我们假设 FPx 权重图块和 FP16 激活图块已从 DRAM 复制到共享内存。然后,共享内存中的整个权重块将分几个步骤进行反量化。
在每一步,仅将一部分 FPx 权重从共享内存加载到寄存器,使用 SIMT 高效 GPU 运行时(第 5.3 节)将其反量化为 FP16 权重,然后存储在寄存器缓冲区 A1 或 A2 中作为 Tensor Core 的输入。然后使用 Tensor Cores 将 A 和 B_{Slice} 相乘。
与一次对整个图块进行反量化相比,我们的逐片反量化将存储 FP16 权重所需的寄存器数量减少了 4 倍,从而显着降低了寄存器压力。此外,还为指令级并行性创造了更多机会,因为一旦权重片被反量化,审查核心就可以立即用于计算,而不是等待整个图块。
有效重叠
软件流程通过图 8b 中的时空图进行说明,其中 SIMT 核心(负责反量化)、Tensor 核心(负责矩阵乘法)和 GPU 内存层次结构协同工作,实现高指令级并行性。
首先,全局内存读取是使用 cp.async [20] 内部函数异步执行的,与其他操作完全重叠。
内存屏障和线程块同步在处理第三个切片后(在 k=2 结束时)发出,确保下一个主循环的数据在共享内存中准备就绪,以便“De-quant”(de-quant)量化)并且当k=3时可以开始“ldmatrix”操作。
其次,共享内存读取也与张量核心操作重叠。当计算 islice 时,通过“De-quant”和“ldmatrix”同时从共享内存中读取第 (i + 1) 个切片的数据。
最后但并非最不重要的一点是,用于权重反量化的 SIMT 核心操作也与 Tensor Core 操作有效重叠。
在 第i个slice 的“反量化”过程中,首先使用硬件固有负载共享 (LDS) 将 FPx 权重从共享内存加载到寄存器,然后立即使用 SIMT 内核反量化为 FP16 权重。同时,Tensor Core 正在计算不依赖数据的 (i − 1) 切片。
我们实现了支持矩阵乘法 C = A × B 的 TC-FPx 内核,其中 A 是形状 [M, K] 的权重矩阵,B 是形状 [K, N] 的激活矩阵。
权重矩阵以我们在 5.2 节中描述的自定义格式存储,输入和输出激活矩阵以列优先存储。
因此,我们的 TC-FPx 内核可以直接替代量化 LLMs 推理框架中的 cuBLAS 内核。
我们的 GPU 内核是在 Flash-LLM 代码之上使用超过 1.2K 行 CUDA 代码实现的 [37]。
我们的 TC-FPx 内核可以单独编译成 .so 动态可链接库,并且我们提供了一组 C++ API 来调用内核。
因此,我们的内核可以轻松使用和集成。此外,我们还提供了 C++ API 来预打包权重矩阵(参见第 5.2 节)。
更重要的是,我们通过将我们的内核集成到最先进的推理框架 DeepSpeed [19] 中,为量化 LLMs 的端到端推理提供新的系统支持。
评估令牌生成阶段LLMs内线性层的性能。使用 NVIDIA Nsight Compute [23] 测量运行时期间每个 GPU 硬件单元的利用率(第 7.1 节)。
对于端到端评估,我们在具有 CUDA 11.8 的 NVIDIA A100-SXM4-80GB DGX 平台上进行典型 LLMs 的推理。
推理延迟和延迟分解(第 7.3 节)是使用 NVIDIA Nsight 系统 [24] 测量的。
对于每个模型,我们评估了每个 GPU 内核在三种典型推理批量大小(即 8、16 和 32)下的延迟。
TC-FPx 的性能优于 BitsandBytes (W4A16)、cuBLAS (W16A16) 和 TensorRT_LLM(W8A16)高达 8.9 倍、2.6 倍和 1.9 倍。平均而言,当批量大小为 8/16/32 时,TC-FPx 的性能分别优于 BitsandBytes、cuBLAS 和 TensorRT_LLM 7.6×/7.5×/6.6×、2.2×/ 2.2×/2.0× 和 1.3×/1.3×/1.2×。
性能分析 通过广泛的内核分析,我们展示了每个 GPU 硬件单元的利用率,并提供了对性能改进来源的更深入的见解。
在线性层的执行过程中,对于 cuBLAS 基线,当推理批大小小于128时,DRAM 带宽(如图 10a 中的黄线所示)几乎耗尽(>80%),而 GPU Tensor Core(如图 10a 中的黄bar所示),未充分利用(<50%)。
这是大型语言模型推理过程中的常见问题,由大型语言模型的自回归推理方案引起。
在我们对6位量化的支持下,DRAM访问量显着减少(高达2.7倍),缓解了DRAM带宽不足的瓶颈。因此,Tensor Core 可以更有效地用于矩阵计算,如图 10a 中的蓝色条与黄色条所示。
总而言之,我们的内核通过支持 Tensor Core 上的 6 位量化,缓解了“内存墙”问题并实现了更高的计算效率(Tensor Core 的利用率更高)。
此外,它还解释了我们的内核可以超越 TensorRT-LLM 的 W8A16 内核,因为我们可以更有效地减少模型权重的 DRAM 访问。请注意,当推理批量大小较大(大于 128)时,我们的 TC-FPx 内核、cuBLAS 内核和 TensorRT-LLM 的 W8A16 内核的性能最终将收敛到相同的性能,因为它们的性能都将受到 Tensor Core 峰值计算能力的限制。
我们还观察到 BitsandBytes 始终比 cuBLAS 慢,平均速度比 cuBLAS 慢 29.6%。 经过进一步调查,我们发现BitsandBytes采用了双内核方法(在4.1节中讨论)来支持FP4量化。在第一个内核的执行过程中,FP4 模型权重将首先从全局内存加载,反量化为 FP16,然后以 FP16 数据类型写回全局内存。之后,正常
cuBLAS 内核作为第二个内核启动计算矩阵乘法。
因此,由于用于 FP4 反量化的额外 GPU 内核的开销,FP4 GPU 内核总是比原始 FP16 cuBLAS 内核慢。
动态反量化分析图 10b 显示
FP6 到 FP16 反量化的开销有两个方面。 一方面,即使采用我们的 SIMT 高效设计,FP6 到 FP16 去量化也会引入大量按位运算。结果,算术/逻辑单元(ALU)的利用率平均从 6.36% 增加到 38.8%。这也是强有力的证据,表明用于反量化的 SIMT 高效设计(第 5.3 节)至关重要。
另一方面,FP6 到 FP16 去量化还引入了更多浮点乘法,计算权重和量化尺度之间的乘法。一般,
FMA装置利用率从0.33%提高到16.64%。
鉴于ALU和FMA单元都是SIMT核心的一部分,反量化操作不会消耗Tensor Core的计算能力。
更重要的是,通过将这些 SIMT 指令与其他操作重叠,可以有效地隐藏 SIMT 内核的运行时开销,我们的新颖设计如第 5.4 节中所述。
工作负载 如第 3 节所述,在保持模型质量方面,6 位量化比 4 位量化更具吸引力。然而,我们仍然将 W6A16 内核的性能与最先进的 W4A16 内核进行比较,充分证明我们的 6 位量化可以实现与现有 4 位量化方法相当的推理速度。我们评估了 LLaMA-65b 模型 [32] 中线性层在不同批量大小下的性能。
基线 这里的主要基线包括来自 TensorRT-LLM [26](提交:6837c81)的 W4A16 对行量化的支持(Coarse-grained_W4A16)和 W4A16 对分组量化的支持(Finegrained_W4A16)最先进的性能。我们在这里还包括 cuBLAS [22] 作为性能基线,清楚地显示了每种量化方法的优点。 结果图 11 显示了 TC-FPx 和在 LLaMA-65b 模型中运行四个不同线性层(例如 L1、L2、L3 和 L4)的其他基线的延迟加速。我们使用 cuBLAS 的性能来规范其他 GPU 内核的性能。
如图 11 所示,TC-FPx_W6A16、Finegrained_W4A16 和 Coarse-grained_W4A16 表现优于 cuBLAS_W16A16 高达 2.4×、3.0× 和 3.3×。
更重要的是,TC-FPx 与 Finegrained_W4A16 实现了相似的性能,当以批量大小 8/16/32 运行所有这些线性层时,TC-FPx 分别比 Fine-grained_W4A16 快 1.06×/ 1.04×/ 0.94×。
此外,在批量大小为 8/16/32 时,TC-FPx 仅比 Coarse-grained_W4A16 慢 16% / 17% / 24%。
由于 6 位量化可以提供显着更高的模型质量,因此这是值得的权衡。
工作负载 我们评估了 FP6-LLM 在各种模型大小的大型语言模型上的端到端推理性能,例如LLaMA-13b [33]、OPT-30b [41] 和 LLaMA-70b [33]。对于每个模型,我们评估了不同批量大小下的令牌生成吞吐量,从 1 开始,直到 GPU 内存耗尽。
公制。我们使用每 GPU 秒的度量标记来表示标准化推理吞吐量,同时考虑执行时间和硬件成本(即使用的 GPU 数量)。
其计算公式如下:
N表示生成的令牌数量,而N和T表示GPU数量和第i个GPU执行所花费的时间。
我们使用这个指标来评估本节中的端到端推理性能。
设置和基线 我们将每个请求的预填充/提示长度设置为 0.5K,并为每个请求生成 1.5K 令牌,忽略“EOS”(序列结束)令牌。我们将 TC-FPx 内核集成到 DeepSpeed [19] 中进行端到端评估,并将这个新系统支持称为 FP6-LLM。比较的基准是原始 DeepSpeed 系统的 FP16 执行。使用我们的 FP6-LLM,仅使用单个 80GB A100 GPU 来推理所有工作负载,包括 LLaMA-70b 模型 [33]。相比之下,LLaMA-70b 使用两个 80GB A100 GPU 进行推理 FP16 基线的模型,因为模型权重(约 130 GB)无法适合单个 GPU。
LLaMA-70b 图 12a 显示了使用我们的 FP6-LLM (FP6LLM-1GPU) 和 FP16 基线 (FP16-2GPU) 的 LLaMA-70b 模型上的令牌生成吞吐量。根据我们的实验,我们的 FP6-LLM 和 FP16 基线在 GPU 内存耗尽之前最多可以将推理批量大小设置为 32,而 FP6-LLM 仅需要单个GPU 和基线使用两个 GPU。
结果表明,FP6-LLM 可以实现比 FP16 基线高 1.69×- 2.65× 的归一化推理吞吐量。 我们对这个端到端推理过程进行了仔细的延迟分解。如图 12b 所示,即使 GPU 数量减半,我们的 TC-FPx 内核(用于 FP6-LLM)平均比 cuBLAS 内核(用于 FP16 基线)快 1.20 倍。
此外,使用 FP6-LLM 可以完全避免 NCCL [25] 开销(跨 GPU 通信),因为只需要单个 GPU。我们还注意到 FP16 基线加速了具有 2 路张量并行性的 MHA(多头注意力)的计算 [31]。总体而言,就每 GPU 秒的令牌而言,我们的 FP6-LLM 的吞吐量比 FP16 基线高出 2.65 倍。
Inference\_Performance = \frac{N_{token}}{\sum_{i=1}^{N_{GPU}} T_i}
OPT-30b 图 13a 显示了使用 FP6-LLM (FP6-LLM-1GPU) 和 FP16 基线 (FP16-1GPU) 的 OPT-30b 模型上的令牌生成吞吐量。
根据我们的实验,FP6-LLM 在 GPU 内存耗尽之前最多可以将推理批量大小设置为 16,而 FP16 基线最多可以在一个批次中服务 4 个请求。因此,FP6-LLM 在批量大小为 16 时最多可以实现每 GPU 秒 319.1 个令牌(高出 4.05 倍),而 FP16 基线在批量大小为 4 的情况下最多可以实现每 GPU 秒 78.8 个令牌,给定相同的 GPU 预算。
此外,当批量大小设置为 1/2/4 时,与 FP16 基线相比,FP6-LLM 可以实现 1.91×/ 1.84×/ 1.72× 高的生成吞吐量。这些整体性能的改进主要来自执行线性层的时间的减少。
如图 13b 所示,TC-FPx 内核平均比 FP16 cuBLAS 内核快 2.39 倍。
LLaMA-13b 图 14a 显示了使用 FP6-LLM (FP6LLM-1GPU) 和 FP16 基线 (FP16-1GPU) 的 LLaMA-13b 模型上的令牌生成吞吐量。
根据实验,FP6-LLM 和 FP16 基线在内存耗尽之前最多可以将推理批量大小设置为 32。平均而言,与使用相同批量大小的 FP16 基线相比,FP6-LLM 可以实现 1.23 倍的生成吞吐量。 由于非内核开销,与前两个模型相比,该模型的整体性能改进不太显着。
根据图 14b,TC-FPx 内核的线性层执行时间显着减少(平均快 2.11 倍)。然而,运行其他 GPU 内核的部分加上 GPU 空闲时间会增加,从而削弱了整体性能的提升。
原因是,随着模型尺寸变小,由于内核启动延迟和 GPU 同步,GPU 往往会拥有更大比例的空闲时间。
相关设计技术[40]和[35]先前提出了4+2权重分割。然而,他们只提出了直观的想法,而没有进行全面的系统设计。
Zeroquant(4+2): Redefining llms quantization with a new fp6-centric strategy for diverse generative tasks. arXiv preprint arXiv: 2312.08583, 2023.
[40] Zhewei Yao, Xiaoxia Wu, Cheng Li, Stephen Youn, and Yuxiong He. Zeroquant-v2: Exploring post-training quantization in llms from comprehensive study to low rank compensation, 2023.
在本文中,我们介绍了TC-FPx,这是第一个全栈GPU内核设计方案,统一张量核心支持各种量化位宽的浮点权重,缓解LLM 推理提供新的端到端支持(称为 FP6-LLM),其中在推理成本和模型质量之间实现了更好的权衡。 FP6-LLM通过一系列新颖的技术解决了硬件不友好的内存访问和反量化高计算开销的问题,以显着更少的GPU内存实现了更快的推理速度。评估表明,FP6-LLM 仅使用单个 GPU 即可实现 LLaMA-70b 的推理,实现比 FP16 基线高 1.69×-2.65 倍的归一化推理吞吐量。此外,FP6-LLM将OPT-30b的推理吞吐量提高了1.72×-4.05×。