大模型量化技术原理：FP6近年来，随着Transformer、MOE架构的提出，使得深度学习模型轻松突破上万亿规模参数，

本文为稀土掘金技术社区首发签约文章，30天内禁止转载，30天后未获授权禁止转载，侵权必究！

近年来，随着Transformer、MOE架构的提出，使得深度学习模型轻松突破上万亿规模参数，从而导致模型变得越来越大，因此，我们需要一些大模型压缩技术来降低模型部署的成本，并提升模型的推理性能。模型压缩主要分为如下几类：

剪枝（Pruning）
知识蒸馏（Knowledge Distillation）
量化（Quantization）

本系列将针对一些常见大模型量化方案（GPTQ、LLM.int8()、SmoothQuant、AWQ等）进行讲述。

之前对FP8进行了基本介绍，同时，讲述了FP8在不同推理框架中的应用，本文将围绕微软研究人员在FP6量化上相关的一些工作（ZeroQuant(4+2): Redefining LLMs Quantization with a New FP6-Centric Strategy for Diverse Generative Tasks、 FP6-LLM: Efficiently Serving Large Language Models Through FP6-Centric Algorithm-System Co-Design）讲述。

文章较长，建议先点赞收藏，后续再慢慢观看。另外，我撰写的大模型相关的博客及配套代码均整理放置在Github：llm-action，有需要的朋友自取。

ZeroQuant(4+2)

背景

近年来，LLMs 的广泛应用受到了其需要很大计算和内存需求的挑战，为了缓解这些挑战，业界诞生了很多量化技术来压缩模型的大小。但模型的压缩通常会带来模型质量下降。目前，一些先进的算法设计（例如：GPTQ 和 LoRC ）使质量下降大大减少。但是是针对较大的模型大小（大于 13B），对于较小的模型大小（例如：1B）会导致模型质量下降很多。同时，它们只关注zero-shot的评估。而在生产环境中，保持原始模型的不同任务的性能至关重要，模型质量的任何损失都是一个主要问题。现有方法虽然具有创新性，但并不能完全满足在实际应用中部署LLMs的实际要求。

为了应对这些挑战，论文主要的工作包括：

扩大模型量化评估范围。研究表明，现有的定量（如：GPTQ）方法往往会过度拟合校准数据集。本文扩大了 LLMs 中 4 位量化的评估范围，以包括zero-shot之外的任务，例如：代码生成和摘要。文中发现 INT4 量化通常表现不佳，尤其是在较小的模型中，甚至在参数高达 130 亿的情况下，例如：LLaMA-13b。
证明了 FP6 量化的卓越性能。本文说明了 FP6 采用基本的舍入到最近 (RTN) 算法和粗粒度量化方法，始终能够达到与全精度模型相当的精度，并在广泛的生成任务中证明了其高度有效。使用 FP6 量化的 StarCoder-13B 模型在代码生成任务中的性能，与 StarCoder-13B FP16 模型相匹配。而对于 406M 等较小模型在总结任务中，与基线结果接近。这超出了 INT4 量化的能力。
创新的4+2 FP6设计。本文为 FP6 引入了创新的 4+2 设计，克服了 AI 硬件上的集成和加速问题。该设计实现了与最先进的 INT4 细粒度量化类似的延迟，使 FP6 成为 LLMs 中现有 4 位量化方法的可行替代方案。

LLM量化需要进行全量的评估

本文通过Zero-Shot任务(困惑度、准确性)、生成任务（ROUGE或Pass@1）、摘要任务来评估模型的性能。全面比较了 RTN 和 GPTQ 方法的 INT4 精度的细粒度量化（FGQ）和粗粒度量化（CGQ），得出的结论是位精度和性能之间没有明确的联系。

同时发现虽然 GPTQ 在PTQ量化方面具有创新性，但它往往会过度拟合特定的数据集，尤其在其细粒度量化时。对于较大的模型（从 1B 到 13B 或 65B），过拟合现象不太严重。

LLMs的核心优势在于其生成序列的能力，通过评估大模型的代码生成能力发现，与标准基准测试相比，INT4 的性能存在显著差异。比如：

在 JavaScript 代码中， CodeLLaMA-34B 模型的性能从 45.05 (FP16) 下降到 43.45 (INT4, CGQ) 或 43.22 (INT4, FGQ)，分别下降了 1.6 和 1.83 个点。
对于较小的模型和 JavaScript 代码中，虽然 INT4 上的 FGQ 比 CGQ 有相当大的改进，但与 FP16 相比仍然存在差距。
但在 Python 代码中，FGQ 上的 INT4 CodeLLaMA-34B 得分为 46.88，超过了其基准，而 FGQ 上的 INT4 CodeGeeX2-6B 得分仅为 29.8，甚至落后于其 INT4-CGQ 性能。这凸显了 INT4 的不一致性。

因此，需要更深入探索 INT4 在复杂生成任务中的有效性。

补充：pass@k 评估指标

使用pass@k指标评估功能正确性，即每个问题生成k个代码样本，如果有任何样本通过单元测试，则认为问题已解决，并报告问题解决的总比例。然而，以这种方式计算pass@k会有很高的方差性。相反，为了评估pass@k，我们为每个任务生成n≥k的样本，计算通过单元测试的正确样本 $c≤n$ 的数量，并计算出无偏估计。

最佳解决方案：FP6

由于 INT4 量化在代码生成和摘要任务中的不稳定性以及之前的研究表明 FP8 在激活过程中的应用比 INT8 的使用具有显著的改进。受此启发，提高位精度（例如：提高到 5 位或 6 位）能否在生成任务中提供更稳定、更稳健的结果？本论文主要是探讨 FP6（FP5）的有效性程度及其对不同量化算法的适应性，为之前 INT4 量化挑战带来的困境提供潜在的解决方案。

之前在FP8量化一文中介绍过标准的浮点数表示格式，主要由三部分组成：符号位、指数位和尾数位。

$x = S \times 2^{E-b} \times M,$

其中，S 表示符号 (±1)，E 表示指数( $[0, 2^e-1]$ ，e 是指数位数），b 是指数的偏置（通常为 $2^{e-1}-1$ ），M 表示尾数 ( $[0, 2)$ )。

因此，对于 $FP6_{E3M2}$ 而言，可实现的最大/最小值为 $±28（±1×2^4×1.75）$ 。

FP16（或BF16）权重矩阵进行如下量化：

$\hat W_{fp16} \approx Quant(W_{fp16}) = S_{fp16} \times W_{fp6}$

其中， $W_{fp16}$ 是原始全精度权重，Quant(·)表示量化函数， $S_{fp16}$ 是缩放因子， $W_{fp6}$ 是 $FP6_{E3M2}$ 数。比例因子 $S_{fp16}$ 的计算公式为： $S_{fp16} = {max(abs(W_{fp16}))}/{28}$ ，从而确保在不影响精度的情况下，最佳地使用 $W_{fp16}$ 的范围。

为什么不用 INT6 而不是 FP6？

选择 FP6 而不是 INT6 是由两个关键因素驱动的：

FP 格式简化了转换过程，因为最终计算通常使用 FP16 或 BF16 执行。
ZeroQuant-FP的研究表明这些格式之间的准确性没有观察到差异，从而无需额外的实验验证。

FP6 性能评估

通过代码生成、摘要和Zero-Shot任务上进行了全面测试。总的来说，FP6 量化方法，特别是使用 CGQ 在不同任务、模型甚至量化算法（RTN 和 GPTQ）之间得到了高性能和鲁棒性的平衡，比 FP5 和 INT4 量化都有显著改进。同时，并没有发现位精度和性能之间存在明确的联系。

关于 FP6 系统支持探讨

4+2 格式用于 FP6

为了解决使用非标准 6 位数字格式的挑战，该格式偏离了传统的 2 次幂数字格式，本文提出了一种新颖的方法，此方法与两种通常考虑的策略不同：

第一种方法：直接将 6 位格式转换为 8 位浮点 (FP8) 格式。虽然这是一个简单的解决方案，但它抹杀了 6 位格式的主要优点：节省内存。
第二种方法：将多个 6 位数字组合在一个连续的内存块中，并使用 32 位整数 (INT32) 或 32 位浮点 (FP32) 格式表示它们。该方法保持了节省内存的优势，但增加了反量化过程的复杂性。

本文侧重于将 6 位数字划分为两个不同的子数字：第一个子数字代表最初的 4 位，第二个子数字代表剩余的 2 位。本文提出的“4+2”方法可以看作是第二种标准方法的高级变体。在此基础上，将 6 位数字分为两个部分：

第一部分由最初的 4 位组成，处理符号位和 3 位指数等元素。
第二部分包含剩余的 2 位，专用于 2 位尾数。

这种划分为 4+2 位的方式有利于这些子数的同时加载和反量化，生成最终的 16 位浮点 (FP16) 权重。

该方法创新地平衡了减少内存占用的需求与反量化的实用性，特别是在解决跨分段数字内存访问的方面。

偏置移位

在 GPU 上运行时将 FP6 反量化为 FP16 可能会占用大量资源，这主要是由于操作指数字段所涉及的复杂性。指数的偏置项通常由指数位确定，对于 FP6 为 3，对于 FP16 为 15。数学上，将FP6反量化为FP16（不包括符号）的过程表示为：

2^{E_{\text{FP16}} - 15} \times M_{\text{FP16}} = 2^{E_{\text{FP6}} - 3} \times M_{\text{FP6}}

其中，上标 FP16/FP6 表示各自的格式。

注意：对于细粒度（子行）量化方案，缩放因子反量化可以在矩阵乘法之后、累加之前进行；对于粗粒度（按行）量化方案，缩放因子反量化可以在累加之后进行。

虽然填充（padding）可以轻松调整尾数，但由于偏置的差异，对齐指数需要更多的努力。可以通过将 INT4 数字转换回对称 INT8 格式来进行类比：如果 INT4 采用对称格式（对于尾数），则零填充就足够了。然而，在非对称格式中，单独填充是不够的，还需要额外的步骤。

为了解决这个问题，论文定制了 FP6 格式，使用非标准指数偏置为 15。此修改不会影响精度或准确度，因为：

S_{\text{FP16}} \times 2^{E-3} \times M = (S_{\text{FP16}} \times 2^{12}) \times 2^{E-15} \times M

这意味着偏置位移可以无缝集成到缩放因子中。至关重要的是，由于 $S_{\text{FP16}}$ 小于 1，因此将其与 $2^{12}$ 相乘仍然可以通过简单的指数位移以 FP16 格式准确表示，从而避免数值错误。

下图 a 概述了原始对每个 FP6 权重进行反量化的两步过程。

第一步涉及将 $W_{fp6}$ 转换为 $W_{tmp}$ 。
第二步需要乘以量化比例 $S_{fp6}$ 。

此过程中要求最高的部分是重新计算 $W_{tmp}$ 的指数，其中涉及从 $W_{fp6}$ 中提取 $E_{fp6}$ ，加上 12，然后将其合并回 $W_{tmp}$ . 此外，对subnormal FP6 数进行反量化的过程进一步增加了运行时反量化的复杂性。

然而，利用本文的偏置移位策略，如下图 b 所示，指数调整变成了一个简单的bit-level填充过程。最初在步骤 1 中需要添加常量整数 12，现在可以推迟到步骤 2，从而消除任何运行时开销。这是可行的，因为量化比例与常数整数的乘法可以在模型量化之后和运行时之前静态地执行。此外，这种简化的方法还有效地适应了subnormal数的反量化。

系统评估

通过对比对仅权重量化 GPU kernels 的性能，FP6 kernel 通过 Bias-Shift 增强，速度比 cuBLAS 最高快 2.37 倍，平均快 1.92 倍。鉴于 LLM 推理通常受到 GPU DRAM 的限制，本文的方法通过最小化模型权重内存访问，有效地缓解了这一瓶颈。此外，FP6 kernel 在速度上优于 TensorRT-LLM 中最先进的细粒度 INT4 实现，FFN1平均快 1.06 倍，FFN2平均快 2.05 倍。另外，带有 Bias-Shift 的 FP6 kernel 比没有 Bias-Shift 的相同 FP6 kernel 平均快 1.36 倍，这凸显了 Bias-Shift 的关键作用。

小结

在ZeroQuant(4+2)中，探索了INT4量化技术(如GPTQ算法) 在LLMs中的表现能力。虽然这些技术可以减小模型大小和参数存储量，但由于过拟合问题, 它们在更一般的许多任务中往往表现不佳，包括代码生成和摘要等更多生成任务。因此，当前迫切需要新的方法来提高LLMs的效率和有效性。

通过对不同量化方法的探索，发现FP6精度格式在各种任务的性能和灵活性方面均表现出色。比如：使用FP6量化的模型，如StarCoder-15B，在代码生成方面达到了与FP16模型相当的结果，而较小的模型(如BART-406M)在摘要方面达到了标准FP16性能水平。但是FP6数据格式在当前AI硬件的高效支持中存在挑战。为了提高FP6在当前主流AI硬件上的执行效率，提出了一种新颖的 4+2 FP6 GPU kernel 方案，以实现与最先进的 INT4 细粒度量化相似的延迟。这一创新使FP6成为提高LLMs效率的有效途径。

FP6-LLM

背景

当今部署 LLMs 具有挑战性，一方面，它需要大量的 GPU 内存（FP16 中的 GPT-3 需要 326 GB）才能容纳模型权重，而 A100/H100 GPU 仅具有最多 80 GB 内存。另一方面，LLM推理在token生成过程中面临严重的"内存墙"问题，其中LLM推理的速度主要受到从 GPU DRAM 读取模型权重的时间限制。它使 LLM 推理内存受到限制，无法充分利用 GPU 的计算能力。最近的研究表明， LLM 部署时，FP6 是推理成本和模型质量之间的良好权衡（推理成本低于 8 位量化，模型质量比 4 位量化更好）。FP6 可以有效地减小 LLMs 的大小，并在不同的应用程序中保持一致的模型质量。

然而现有系统不提供对 FP6 量化的 Tensor Core 支持，并且在 LLM 推理期间难以实现性能改进。在 GPU 上支持 FP6 量化具有以下挑战：

(1) 位宽不规则的模型权重对内存访问不友好。
(2) 权重反量化的运行时开销较高。

为了解决这些问题，本文提出了TC-FPx，这是第一个全栈GPU KERNEL设计方案，具有统一的Tensor Core支持各种量化位宽的浮点权重。同时，将 TC-FPx kernel 集成到现有的推理系统中，为量化 LLM 推理提供新的端到端支持（FP6-LLM）。实验表明，FP6-LLM 仅使用单个 GPU 即可实现 LLaMA-70b 的推理，实现比 FP16 基线高 1.69× - 2.65× 的推理吞吐量。

补充： CUDA Cores 与 Tensor Cores

SIMT cores (CUDA Cores) 负责 GPU 中的通用处理任务，处理各种指令，包括整数运算、浮点运算、加载/存储操作等。SIMT cores 执行单个（或矢量）数据元素。

Tensor Cores 是为加速矩阵乘法而设计的专用硬件。在 A100 /H100 GPU 上，Tensor Cores 的 FLOPS 比 SIMT Cores 高 16.0 倍/ 14.8 倍。此外，Tensor Cores 以粗粒度工作，例如：使用单个 mma（矩阵乘法和累加）指令在形状为 16 × 16 和 16 × 8 的两个 FP16 矩阵之间执行矩阵乘法。

FP6 Kernel 设计的选择与挑战

设计选择

量化的主要目标是线性层的权重（即矩阵乘法），占总体 LLM 权重的99%以上。而现有对线性层的支持主要是针对位宽为2的指数的数据类型（例如：4位、8位和16位）而设计的。本文之前尚不清楚如何在现代 GPU 上有效支持 FP6。因此，本文提出了两个重要的设计选择。

启用 Tensor Cores：本文发现在执行量化 LLMs 推理时支持 Tensor Core 很有必要。通过评估 AWQ 纯 SIMT Cores 执行在各种批量大小上的性能，以测试其可扩展性。如下图所示，随着推理批量大小的增加，不带 Tensor Core 支持的线性层 (AWQ_W4A16_SIMT) 的运行时性能变得极低。主要有两点原因：

一方面，对于线性层执行，传统 SIMT Core 比 Tensor Core 慢一个数量级。
另一方面，SIMT Core 计算能力的很大一部分将用于在运行时对模型权重进行反量化，这进一步降低了 SIMT Core 用于计算矩阵乘法的可用计算能力。

因此，本文启用 Tensor Core 来进行矩阵乘法的密集计算，同时利用多功能 SIMT Core 进行权重反量化。

统一kernel而不是使用双kernel：WxA16量化的独特之处在于激活矩阵使用FP16，但权重矩阵以较窄的位宽存储。但是 Tensor Core 要求将权重矩阵和激活矩阵存储在同一数据类型中（例如：FP16/INT8/INT4）。简单的解决方案（即双kernel解决方案）添加一个额外的 GPU kernel，在调用普通 FP16 kernel之前将权重反量化为 FP16。然而，这样的推理速度会比没有量化的模型还要慢。如下图左边所示，将启动两个 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操作（按位操作）。因此，如何将反量化融合到线性层计算中而不影响整体性能也很重要。

设计方法论

本论文为了解决不友好的内存访问的挑战，提出了运行前 Bit-level 预包装。为了应对反量化高计算开销的挑战，实现了运行时 SIMT 高效 GPU 计算。同时，设计了高效流水线，使 SIMT Cores、Tensor Cores和 GPU 内存层次结构以最佳性能协同工作。

下图对比了本文设计的 TC-FPx（仅 x 位权重量化线性层 Kernel）与传统的通用矩阵乘法 (GEMM) ，其中两个输入矩阵均采用 FP16。

在 TC-FPx 中，模型权重的存储位数有所减少。同时，在寄存器层级引入了额外的反量化阶段 (Dequant W)，其中，使用 SIMT cores 在每个线程内将 FP6 权重本地反量化为 FP16。

注意：FP16 权重不会写回共享内存，而是存储在寄存器中以供将来使用，从而消除了对共享内存不必要的往返访问。

另一个区别是TC-FPx使用细粒度的lds指令将x位权重从共享内存加载到寄存器，而不是使用粗粒度的固有ldmatrix（加载矩阵），后者具有严格的布局要求并且灵活性更小。

运行前比特层级（Bit-level）预包装

为了解决对于现代 GPU 内存层次结构对于不规则位宽的权重的内存访问不友好的问题。本文提出可以每 32 个 x 位权重的内存读取合并，从而产生 x 请求（每个 GPU 线程 4 字节字）。在这种情况下，所有内存访问都将以 32 位字的粒度对齐，而不是不规则的位宽。

然而，由于 Tensor Core 严格的数据布局要求，将权重的内存读取合并起来并非易事，因为每个 GPU 线程所需的权重并不存储在连续的内存空间中。

为了解决这个问题，本文通过重新排序每个权重矩阵中的权重并提前预包装权重来优化运行时内存访问。由于模型权重是在模型训练和量化后静态确定的，因此可以提前对权重应用复杂的内存布局转换，从而不会引入运行时开销。

由于只需要预包装一次权重，因此预打包权重的开销可以通过每次推理服务有效摊销，并且可以忽略不计。

一般来说，权重预封装包括两个步骤。

第一步，收集每个 GPU 线程所需的所有权重，并在本地组合这些权重。鉴于每个 GPU 线程所需的权重最初并不位于每个权重矩阵内的连续位置，因此我们必须仔细选择每个 GPU 线程的权重。然后，为每个线程选择的权重在本地以相对时间顺序组合，因为它们在运行时被 Tensor Core 消费。

第二步，将整个GPU WARP（由32个GPU线程组成）所需的所有权重组合到一个统一的线性内存空间中，权重将在运行时按顺序存储在GPU DRAM中。为了完全消除共享内存库冲突，以锯齿状顺序组合每个线程的 32 位字。

注意：这里讨论的所有技术都独立于模型权重的实际位宽（始终使用 x 表示）。因此，我们的权重预封装可以自然地应用于任何位宽。

步骤一：逐线程权重收集

下图展示了T0(Thread#0)选取的权重以及组合它们的顺序。

假设 WARP 层级切片大小为 64 × 64 ，这意味着每个权重矩阵被分为 64 × 64 数据块，并以每个 WARP 的粒度加载到 GPU 的共享内存。然后，每个权重块进一步分为四个片，因为权重是从共享内存加载的，并逐片用于 Tensor Core 计算。每个切片被分为四个 16 × 16 块，因为 Tensor Core 在每条指令中处理 16 × 16 数据块。

在每个 16 × 16 块中，为 T0 选择四对 FPx 权重并组合。

如图所示，经过步骤1，我们得到了32组（即WARP大小）FPx权重。权重在每组中合并连续存储，并且每组权重将被某个GPU线程消耗。

综上所述，每个 64 × 64 权重块最终分配给 32 个线程（一个 WARP），每个线程将消耗 128 个 x 位权重。

步骤二：按 WARP 进行 Bit-level 封装

在步骤二中，将不同组的所有权重组装成统一的内存空间。

在这个 Bit-level 预包装过程中，将组合权重作为连续的数据进行复制，暂时忽略每个位的含义。具体来说，x 位的 128 个条目被视为 32 位的 4x 个条目。

本文建议按照图中的锯齿状顺序组装所有组的权重。

首先，将每个线程的第一个 32 位条目连接在一起。之后，每个线程的第二个 32 位条目被连接并附加到之前的结果中。通过重复这个过程，所有权重可以连续存储在线性内存空间中，并且对齐良好（128字节对齐）。这样，所有权重就可以以 128 字节块的粒度简单地从 DRAM 复制到共享内存，无需任何更改，轻松实现最佳 DRAM 访问。

此外，这些权重可以在运行时以最佳性能从共享内存加载。

具体来说，线程的 WARP 将为每个内存请求读取共享内存中的连续 32 位条目，完全避免了存储体冲突。

运行时高效的SIMT计算

并行反量化：

为了减少 FP-x 权重反量化的运行时开销，使用优化的按位 SIMT core 指令实现了 FP-x 反量化。此外，建议并行对多个 FPx 权重进行反量化，通过利用每个 32 位寄存器内的 bit-level 并行，进一步将 SIMT 开销减少 4 倍。

（1）优化的按位运算：当将 FPx 转换为等效的 FP16 时，FP16 的指数应为 $E^{fp16} = E^{fpx} + bias^{fp16}-bias^{fpx}$ 。为了简化这个过程，我们采用了Zeroquant(4+2)中的数学变换，用 $E^{fp16} = E^{fpx}$ 来计算FP16的指数。为了保持正确性，将 FP16 的结果与 FP16 常数 $2^{bias^{fp16}-bias^{fpx}}$ 相乘：

$cast(W_{fpx}) = new\_cast(W_{fpx}) \times 2^{bias^{fp16}-bias^{fpx}}$

下图a显示了优化后 FP6 到 FP16 转换。虽然只绘制从 FP6 到 FP16 的转换来进行演示，但它可以应用于任何位宽。 FP16 的符号字段与 FPx 的符号字段相同。此外，为了提高效率，指数字段的低位和尾数字段的高位可以一起从 FPx 复制到 FP16。此外，FP16 的其他位应该用零填充。

经过精心设计，仅用两个按位“and”、一个“shifting”和一个“or”就成功实现了从 FP6 到 FP16 的转换，如图b中❶所示。使用第一个“and”将符号字段从 FP6 复制到 FP16，同时将 FP16 的所有其他位初始化为零，从而无需稍后将零填充到指数和尾数字段。然后，FP6 的所有位都通过逐位“right shifting”向右移动。之后，首先通过FP6和位掩码“0x1f1f1f1f”之间的“and”选择FP6中指数的低位和尾数的高位，然后通过按位运算“or”复制到FP16。

(2) 比特层级（Bit-level）并行：考虑到可以利用每个 32 位字内的比特层级并行，建议并行地对多个 FPx 权重进行反量化，进一步减少反量化的运行时开销。在图b中，以 FP6 为例演示了详细设计。 32 位寄存器被视为四个处理槽，其中每个槽独立工作，使被相同的指令但输入不同的 FP6 数据。在开始反量化之前，应将四个 FP6 存储在 R1（Register#1）中，初始数据布局如图所示。通过代码片段❶，这四个FP6可以同时反量化为四个FP16，其中每个FP16仅前8位存储在R2中。之后，将第一个和第二个 FP16 提取到 R1，最后 8 位补零，即代码片段❷。最后，通过代码片段❸和❹，将第三个和第四个FP16提取到R2中。

权重分割和拼接：

下面将演示在 GPU 上通过精心设计的内存布局从 2+4 组合有效重建 6 位权重的方法，该方法也可以应用于其他位宽。

(1)运行前权重分割：为了将权重以良好对齐的方式存储在GPU的32位寄存器中，我们将每个权重分割成几个段，其中每个段的位宽为2，例如：每个 6 位权重可以分为 2+4 或 4+2。基于该方案，后续设计的指标计算得到显著简化。

(2)运行时权重拼接：在反量化之前，首先将权重从共享内存加载到寄存器。由于每个权重被分成几个段，因此需要在运行时在寄存器级别重建完整的权重。为了减少运行时开销，我们建议并行提取和拼接权重。如下图所示，使用两组寄存器来存储32个FP6权重，其中，Frag1_PTR 指向包含 32 个 2 位段的 2 个 32 位寄存器，而 Frag2_PTR 指向包含 32 个 4 位段的 4 个 32 位寄存器。通过我们的并行拼接，可以同时重建四个 FP6 权重，从而将 SIMT core指令的数量减少 4 倍。如下图所示，首先将四个 2 位段提取到Register#1 (❶)，然后将四个 4 位段提取到Register#2 (❷)。之后，Register#2 右移 (❸)，并将其有效位复制到Register#1 (❹)，从而得到完整的 6 位权重。

(3)位重新排序：为了并行提取和拼接权重，有必要强制执行上图中的初始数据布局。关键的观察是每四个连续段必须按图中所示的顺序放置，例如：前四个段必须按#2、#4、#1 和#3 的顺序存储。此外，每对 2/4 位段之间的步长应分别为 6/4。否则，不可能仅用四个SIMT核心指令同时拼接四个段。为了满足图中的初始数据布局要求，建议通过在运行之前重新排序权重段来确保这种布局，而无需运行时开销。

总体伪代码算法 ：

下面显示了包括并行反量化和权重拼接的伪代码（GPU 代码）。伪代码中的所有输入和输出变量都存储在寄存器中。如上图所示，代码总共对 32 个 FP6 权重进行反量化。对于每个外循环，会生成四个 FP16 权重，并将其存储在代码末尾的两个寄存器中。上图中的转换（❶、❷、❸ 和 ❹）分别是通过代码中第 6、7、9 和 10 行的 SIMT core 操作实现的。然后，Tensor Core 直接使用输出寄存器数组 (OutputReg) 作为输入。

全栈的高效流水线设计

为了减少 GPU 寄存器的使用，逐片（slice）对权重进行反量化。此外，将反量化过程无缝融合到线性层执行的传统软件流水线中，通过高效的指令并行完全隐藏了反量化的运行时开销。

逐片反量化

本文不是一次反量化所有权重，而是逐片反量化 FPx 权重。如下图a所示，我们假设 FPx 权重块和 FP16 激活块已从 DRAM 复制到共享内存。

然后，共享内存中的整个权重块将分几个步骤进行反量化。在每一步，仅将一部分 FPx 权重从共享内存加载到寄存器，使用运行时高效 SIMT GPU 计算将其反量化为 FP16 权重，然后存储在寄存器缓冲区 A1 或 A2 中作为 Tensor Core 的输入。然后使用 Tensor Cores 将 A 和 $B_{Slice}$ 相乘。

与一次对整个块进行反量化相比，逐片反量化将存储 FP16 权重所需的寄存器数量减少了 4 倍，从而显著降低了寄存器压力。此外，还为指令级并行创造了更多机会，因为一旦权重片被反量化，Tensor Cores 就可以立即用于计算，而不是等待整个块。

有效重叠

软件流程通过上图b中的时空图进行说明，其中，SIMT cores（负责反量化）、Tensor cores（负责矩阵乘法）和 GPU 内存层次结构协同工作，实现高指令级并行。

首先，全局内存读取是使用 cp.async 内部函数异步执行的，与其他操作完全重叠。内存屏障和线程块同步在处理第三个切片后（在 k=2 结束时）发出，确保下一个主循环的数据在共享内存中准备就绪，以便“De-quant”（反量化）并且当 k=3 时可以开始“ldmatrix”操作。

其次，共享内存读取也与Tensor cores操作重叠。当计算 islice 时，通过“De-quant”和“ldmatrix”同时从共享内存中读取第 (i + 1) 个切片的数据。

最后，用于权重反量化的 SIMT Core 操作也与 Tensor Core 操作有效重叠。在第 i 个 slice 的“反量化”过程中，首先使用硬件固有负载共享 (LDS) 将 FPx 权重从共享内存加载到寄存器，然后立即使用 SIMT Core 反量化为 FP16 权重。同时，Tensor Core 正在计算不依赖数据的 (i − 1) 切片。

TC-FPx Kernel 实现

FP6-LLM 中实现了矩阵乘法 C = A × B 的 TC-FPx 内核，其中 A 是形状 [M, K] 的权重矩阵，B 是形状 [K, N] 的激活矩阵。权重矩阵以自定义格式存储，输入和输出激活矩阵以列优先存储。

因此，TC-FPx Kernel 可以直接替代量化 LLMs 推理框架中的 cuBLAS 内核。

本文的 GPU Kernel 是在 Flash-LLM 代码之上使用超过 1.2K 行 CUDA 代码实现的。TC-FPx Kernel可以单独编译成 .so 动态可链接库，并且提供了一组 C++ API 来调用Kernel。因此，可以轻松使用和集成。此外，我们还提供了 C++ API 来预封装权重矩阵。

更重要的是，通过将 TC-FPx Kernel 集成到推理框架 DeepSpeed 中，为量化LLMs的端到端推理提供新的系统支持。

性能评估

不同形状下线性层加速对比：

对于每个模型，评估了每个 GPU Kernel 在三种典型推理批量大小（即: 8、16 和 32）下的延迟。TC-FPx(W6A16) 的性能优于 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×。

在线性层的执行过程中，对于cuBLAS基线，当推理批大小小于128时，DRAM带宽(如图10a中的黄线所示)几乎耗尽(>80%)，而GPU Tensor Cores(如图10a中的黄条所示)未被完全使用(<50%)。这是自回归推理一个普遍的问题。

使用本文提供的6位量化，DRAM访问显著减少(高达2.7倍)，缓解了DRAM带宽不足的瓶颈。因此，张量核心可以更有效地用于矩阵计算，如图10a中蓝色条与黄色条所示。总之，我们的kernel通过在Tensor Cores上支持6位量化，缓解了“内存墙”问题，实现了更高的计算效率(Tensor Cores的更高利用率)。

我们的内核可以超越 TensorRT-LLM 的 W8A16 kernel，因为我们可以更有效地减少模型权重的 DRAM 访问。当推理批量大小较大（大于 128）时，我们的 TC-FPx kernel、cuBLAS kernel和 TensorRT-LLM 的 W8A16 kernel的性能最终将收敛到相同的性能，因为它们的性能都将受到 Tensor Core 峰值计算能力的限制。

此外还观察到 BitsandBytes 始终比 cuBLAS 慢。经过进一步调查，我们发现BitsandBytes采用了双内核方法来支持FP4量化。在第一个内核的执行过程中，FP4 模型权重将首先从全局内存加载，反量化为 FP16，然后以 FP16 数据类型写回全局内存。cuBLAS 内核作为第二个内核启动计算矩阵乘法。因此，由于用于 FP4 反量化的额外 GPU 内核的开销，FP4 GPU 内核总是比原始 FP16 cuBLAS 内核慢。

FP6-LLM 与 4位量化方案对比：

在LLaMA-65b模型中运行四个不同线性层(即L1、L2、L3和L4)，TC-FPx和其他基线的延迟对比，TC-FPx 与 Fine-grained_W4A16 、Coarse-grained_W4A16都比cuBLAS_W16A16 快。TC-FPx 与 Fine-grained_W4A16 有相似的性能，仅比 Coarse-grained_W4A16 慢，但 TC-FPx 提供更高的模型质量。

对于端到端延迟对比，通过单GPU每秒的Token生成来标准化推理吞吐量，同时考虑执行时间和硬件成本（即使用的 GPU 数量）。

对于 LLaMA-70B，我们的 TC-FPx 内核（用于 FP6-LLM）平均比 cuBLAS 内核（用于 FP16 基线）快 1.20 倍。此外，使用 FP6-LLM 可以完全避免 NCCL 开销（跨 GPU 通信）；就单 GPU 每秒的Token而言，我们的 FP6-LLM 的吞吐量比 FP16 基线高出 2.65 倍。

对于OPT-30B，当批量大小设置为 1/2/4 时，与 FP16 基线相比，FP6-LLM 可以实现 1.91×/ 1.84×/ 1.72× 高的生成吞吐量。这些整体性能的改进主要来自执行线性层的时间的减少。TC-FPx Kernel 平均比 FP16 cuBLAS Kernel 快 2.39 倍。

对于LLaMA-13B，平均而言，与使用相同批量大小的 FP16 基线相比，FP6-LLM 可以实现 1.23 倍的生成吞吐量。由于非内核开销，与前两个模型相比，该模型的整体性能改进不太显著。TC-FPx 内核的线性层执行时间显著减少（平均快 2.11 倍）。然而，运行其他 GPU 内核的部分加上 GPU 空闲时间会增加，从而削弱了整体性能的提升。原因是随着模型尺寸变小，由于内核启动延迟和 GPU 同步，GPU 往往会拥有更大比例的空闲时间。

小结

FP6量化的一个挑战是缺乏针对这种不规则位宽的高效GPU KERNEL设计。在 FP6-LLM 中设计并实现了TC-FPx，第一个具有Tensor Core支持的用于FP6和各种量化位宽(6位、5位、3位等)的浮点权重的GPU系统设计方案，缓解了LLM推理期间的“内存墙”问题。TC-FPx打破了底层GPU硬件的限制，允许GPU支持涉及任意位宽模型权重的矩阵乘法计算。在TC-FPx中，Tensor Cores用于矩阵乘法的密集计算，而SIMT cores在运行时有效地用于权重反量化，将模型权重反量化为FP16类型，Tensor Core基于此进行计算。它具有以下关键创新:

运行前比特层级的数据排布转换。用以解决权重具有不规则位宽时不友好的内存访问挑战，实现GPU内存的最优访问；
运行时高效的SIMT计算。用以最小化权重反量化的运行时开销；
全栈的高效流水线设计。其SIMT计算、Tensor Core计算和GPU内存访问进行高效调度，最大程度提升性能。

平均而言，本方案的 FP6 kernel 在 NVIDIA A100 GPU 上进行（因decoder的矩阵形状狭长而导致参数矩阵的访存成为瓶颈的）矩阵乘法时，处理速度比FP16 cuBLAS 基准提高了2.1倍。值得注意的是，通过FP6量化实现的 FP6 kernel 使LLaMA-70b 模型能够在单个 A100 GPU 上运行。这一使得其在batch小于32的LLM推理任务中，性能比FP16基准高出1.69到2.65倍。目前，TC-FPx Kernel 仅支持 NVIDIA Ampere GPU，并且仅在 A100 GPU 上进行了测试和验证。

FP6 在 DeepSpeed 中的应用

目前，FP6 量化内核已经集成到了 DeepSpeed-FastGen 中，实现了运行时的即时量化。这一增强功能允许通过DeepSpeed-FastGen中的统一配置选项来高效量化和部署大语言模型。

通过里面的接口，用户可以输入HuggingFace模型名称或本地checkpoint目录。输入后，系统将启动指定模型的加载，对每个线性层实现FP6量化，并将量化的权重进行比特层级的数据排布转换。转换后的张量随后作为更新后的权重，而原始的FP16权重被丢弃以优化内存使用。

在推理阶段，FP6 kernel 将利用这些6位的权重进行计算。

FP6 量化为模型推理提供了以下好处：

它使大语言模型（LLMs）能够在更少的GPU上部署。例如，LLaMA-70b在单个A100-80G GPU上就能以FP6形式运行，而FP16模型至少需要两个GPU。
它显著加快了小batch之下内存访问为瓶颈的线性层计算。
FP6量化减少了模型权重的GPU内存需求，允许同时服务更多查询，从而提高了服务吞吐量。

随着生成序列长度的延伸，FP6与FP16之间的性能差异加大。这一趋势主要归因于解码长度扩展时，推理过程变得越来越受内存访问瓶颈限制，有利于我们的权重量化的GPU kernel，相对于FP16实现更大的kernel速度提升。

需要强调的是，较长解码场景中内存访问瓶颈增强的两个因素如下：

首先，KV缓存的内存使用随序列长度增加而增加，减少了可容纳的batch大小并导致线性层的矩阵计算瓶颈变为参数的访存。
其次，在DeepSpeed-FastGen的prefill-decoding-mixed-batch技术背景下，对于decoding较长的情况，用于和decoding进行mixed-batching的prefill切块会相对不足，这导致纯粹用于decoding的batch频率增加，进一步加剧了访存的瓶颈。

尽管FP6量化带来了显著的好处，但当前实现仍面临一些限制。值得注意的是，在GEMM 因batch较大或有充足的GPU内存而使得瓶颈变为Tensor Core计算时，本文的仅权重量化kernel可能无法保持其性能优势，尤其是与厂商的优化库如cuBlas相比。然而，本文系统的低内存占用仍是一个关键优势。目前仅支持非混合专家（Non-MoE）结构。此外，当前系统仅与FP16输入模型兼容，因为当前实现的FP6 kernel仅支持处理FP16的激活。

接下来应用FP6进行推理。

安装deepspeed、deepspeed-mii等

pip install deepspeed deepspeed-mii qtorch

2. 使用以下几行代码即可完成量化推理。

import mii
pipe = mii.pipeline("NousResearch/Llama-2-70b-hf", quantization_mode='wf6af16')
response = pipe(["DeepSpeed is", "Seattle is"], max_new_tokens=128)
print(response)

此外，DeepSpeed 官方提供了一套基准测试脚本，可以很方便进行不同推理及量化方案的延迟吞吐对比。

总结

本文通过微软的两篇论文介绍了FP6量化技术。其中，ZeroQuant(4+2) 中为FP6 引入了创新的 4+2 GPU kernel 设计方案，但没有进行全面的系统设计。而FP6-LLM则是第一个全栈GPU KERNEL设计，使用统一 Tensor Core 支持各种量化位宽的浮点权重，为 LLM 推理提供新的端到端支持。并在推理成本和模型质量之间实现了更好的权衡。同时，FP6-LLM通过一系列新颖的技术解决了硬件不友好的内存访问和反量化高计算开销的问题，以更少的GPU内存实现了更快的推理速度。最后，简要介绍了其在DeepSpeed中的应用。

参考文档：