极市导读
在 12B FLUX.1-dev 上,与 BF16 模型相比,它减少了 3.6 倍的 memory。通过消除 CPU offloading,在 16GB 笔记本电脑 4090 GPU 上时,它比 16 bit 模型提供了 8.7 倍的加速,比 NF4 W4A16 基线快 3 倍
太长不看版
量化文生图扩散模型的成功范式。
Diffusion Model 的部署一直都是个很大的问题,因为随着这种模型越做越大就需要更多的 memory 造成更大的时延,而这些对于部署而言都是不利的。
本文给出了一个很激进的量化方案:把扩散模型的 weight 和 activation 都量化为 4 bit。 在如此激进的量化层面上,weight 和 activation 都高度敏感,传统 LLM 的 post-training quantization 方法,如 smoothing 就不好使了。
因此,本文提出了一种 4-bit 量化算法:SVDQuant。SVDQuant 的核心操作不像 smoothing 那样在 weight 和 activation 之间重新分配异常值 (Outlier),而是使用一个 low-rank 分支来吸收这些异常值。SVDQuant 首先通过将异常值从 activation 转移到 weight 来巩固异常值,然后使用高精度的 low-rank 分支通过奇异值分解 (SVD) 获取权重异常值。这个过程使得 weight 和 activation 的量化都得到了简化。但由于激活的额外数据移动,简单地独立运行 low-rank 分支会产生显著的额外开销,使得量化很难获得实际加速。
为了解决这个问题,本文 co-design 了推理引擎 Nunchaku,把 low-rank 分支的 kernel 吸收到 low-bit 分支中,切断冗余的内存访问。它也可以无缝支持现成的低秩适配器 (LoRA),而无需重新量化。作者在 SDXL, PixArt-Σ, 和 FLUX.1 上都进行了实验。12B FLUX.1 模型的内存使用量减少了 3.5 倍,在 16GB 笔记本电脑 4090 GPU 上的 4-bit权重量化基线上实现了 3.0 倍的加速。
下面是对本文的详细介绍。
本文目录
1 SVDQuant:吸收异常值量化 4-bit 扩散模型
(来自 MIT 韩松团队,NVIDIA)
1 SVDQuant 论文解读
1.1 SVDQuant 研究背景
1.2 量化过程介绍
1.3 问题定义
1.4 SVDQuant 方法介绍
1.5 推理引擎 Nunchaku
1.6 SVDQuant 支持的模型数据集和评价指标
1.7 实验结果
1 SVDQuant:吸收异常值量化 4-bit 扩散模型
论文名称:SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models
论文地址:
http://arxiv.org/pdf/2411.05007
Quantization Library:
http://github.com/mit-han-lab/deepcompressor
Inference Engine:
http://github.com/mit-han-lab/nunchaku
1.1 SVDQuant 研究背景
扩散模型通过在海量数据上进行训练,可以从简单的文本提示生成令人惊叹的图像。为了追求更高的图像质量和更精确的文本到图像对齐,研究者开始扩大扩散模型。Stable Diffusion 1.4 只有 800M 参数,而 SDXL 扩展到 2.6B 参数。AuraFlow v0.1[1]进一步扩展为 6B 参数,最新模型 FLUX.1[2]扩展到了 12B 参数。如图2所示,与大语言模型 (LLM) 相比,扩散模型的计算量要大得多。它们的计算成本随模型大小的增加而迅速增加,给现实世界的模型部署带来了令人望而却步的内存和时延障碍,特别是对于需要低延迟的交互用例。
由于摩尔定律放缓,硬件供应商转向低精度推理以维持性能提升。例如,NVIDIA 的 Blackwell Tensor Cores 引入了一个新的 4-bit floating point (FP4) 精度,与 FP8 相比,性能翻倍。因此,使用 4-bit 推理来加速扩散模型很有吸引力。在 LLM 领域,研究人员利用量化来压缩模型大小并提高推理速度。
但是,扩散模型与 LLM 不同,LLM 的时延主要受权重加载的限制,对于小 Batch size 的扩散模型计算量仍然很大。仅仅量化权重不可以加速扩散模型。为了实现加速,weight 和 activation 必须量化为同等位宽。否则,较低精度的计算性能优势将被抹除。
因此,本文专注于将扩散模型的 weight 和 activation 量化为 4 bits。这种具有挑战性和激进的方案往往容易出现严重的质量下降。本文提出了一种新的通用量化范式 SVDQuant。
其可以用下图3概括。SVDQuant 的核心思想是引入一个低成本的分支来吸收两边的异常值。为了实现这一点, 首先通过平滑将它们从激活 迁移到权重 来聚合异常值。然后将奇异值分解 (SVD) 应用于更新的权重 , 将其分解为低秩分支 和残差 。low-rank 分支以 16-bit 运行, 允许仅将残差量化为 4 位, 这显着减少了异常值幅度。
但是直接运行 low-rank 分支会产生大量的内存访问开销,抵消了 4-bit 推理加速。为了克服这个问题,我们一起设计了一个专门的推理引擎 Nunchaku,将低秩分支计算融合到 4-bit 量化和计算 Kernel 中。这种设计能够在有额外的分支的情况下实现推理加速。
1.2 量化过程介绍
量化是加速网络中线性层的有效方法。给定一个张量 ,量化过程定义为:
式中, 是 的 low-bit 表示, 是比例因子, 是最大量化值。对于有符号的 位整数量化, 。对于 1 位尾数和 2 位指数的 4 位浮点量化, 。因此,去量化张量可以表述为 。对于输入 和权重 的线性层, 其计算可以近似为:
相同的近似适用于卷积层。为了加快计算速度, 现代算术逻辑单元需要具有相同位宽的 和 。
本文专注于用于加速的 W4A4 量化,其中权重和激活的异常值都有很大的障碍。抑制这些异常值的传统方法包括量化感知训练 (QAT)[3]和 Rotation[4][5][6]。QAT 需要大量的计算资源,特别是对于超过 10B 参数的模型 (例如,FLUX.1)。由于扩散模型中使用自适应归一化层,Rotation 不适用。归一化权重由于需要运行时候在线生成,所以没办法离线集成到投影层权重里。所以,在线生成的激活和权重都会产生显著的计算开销。
1.3 问题定义
SVDQuant 的核心思想是引入额外的 low-rank 分支,可以吸收 weight 和 activation 中的量化困难。最后,作者提供了一个具有 Kernel 融合的协同设计推理引擎 Nunchaku,以最小化 4 bit 模型中 low-rank 分支的开销。
设线性层为 , 权重为 , 则量化误差为:
式中, 是 Frobenius 范数。
Proposition 1 (误差分解): 量化误差可以分解为:
从命题中可以看出, 误差受权重和输入的大小 以及它们各自的量化误差 的限制。为了最小化整体量化误差, 目标是优化这4项。
1.4 SVDQuant 方法介绍
Smoothing 是减少异常值的有效方法。可以通过缩小输入 并使用每通道平滑因子 相应地调整权重矩阵 来平滑激活中的异常值。如图 4(a)(c) 所示, 平滑输入 表现出减小的幅度和更少的异常值,从而导致较低的输入量化误差。然而,在图 4(b)(d) 中,转换后的权重 在幅度和异常值的存在方面都显着增加,这反过来又提高了权重量化误差。因此,整体误差减少有限。
用 low-rank 分支吸收增加的权重异常值。 本文核心的做法是引入一个 16-bit low-rank 分支, 并进一步将权重量化难度迁移到该分支。具体来说,将转换后的权重分解为 ,其中 和 是秩 的两个低秩因子, 是残差。那么 可以近似为:
与直接进行 4-bit 量化, 即 相比, 本文的方法首先以 16-bit 精度计算 low-rank 分支 , 然后用 4-bit 量化逼近残差 。实践中, 一般设置 , 通常为 16 或者 32。因此, low-rank 分支的附加参数和计算可以忽略不计。现在, 仍然需要仔细的系统设计来消除冗余内存访问。
根据上式5,量化误差可以写成:
根据命题 4.1, 由于 已经没有异常值, 因此只需要专注于优化 及其量化误差
Proposition 2 (量化误差界): 对于式1 中描述的任何张量 和量化方法,如 。假设 的元素服从正态分布, 则有:
其中 表示 中的元素数。
因此,量化误差 受残差 大小的限制。
因此,本文目标是找到使 最小化的最优 ,这可以通过简单的奇异值分解 (SVD) 来解决。给定 的 SVD,最优解为
图5 绘制了原始权重 、转换后的权重 和残差 的奇异值分布。原始权重 的奇异值高度不平衡。平滑后,转换后的权重 的奇异值分布更加陡峭,前几个值明显更大了。
通过去除这些主导值, Eckart-Young-Mirsky 定理表明残差 的大小显着降低。因为有 , 原始的幅值为 , 其中 为 的第 个奇异值。
此外, 实践表明, 与 相比, 表现出更少的异常值, 具有显着压缩的值范围, 如图 4(d) (e)所示。在实践中,可以通过分解 并相应地调整 多次迭代来进一步减少量化误差,然后选择误差最小的结果。
因此,SVDQuant 的思路:
原始模型中,根据 Proposition 1 得到的误差分为 4 项,都要考虑。
使用了 Smoothing 操作之后, 此时 已经基本没有异常值了, 但是此时 中的异常值加剧。还是不好。
使用了本文提出的 low-rank 分支之后,将误差表达式变为了 5 式,继而将量化误差等效为 7 式 (由于此时 已经基本没有异常值了)。
根据 Proposition 2 得到 7 式的量化误差相当于以残差 为上界。 Eckart-Young-Mirsky 定理表明残差 大小相比原始的幅值 显著降低。
因此,使用 Smoothing 操作和 Low-Rank 分支以后,得到的量化误差显著降低了。
1.5 推理引擎 Nunchaku
尽管 low-rank 分支理论上引入的计算可以忽略, 但它作为一个单独的分支运行会产生显著的延时一大约 的 4-bit 分支延迟, 如图 6 所示。这是因为, 对于一个小的秩 , 即使计算成本大大降低,输入和输出激活的数据大小保持不变,把瓶颈从 Computation 转移到 Memory Access。这种情况, 当 GPU 的 L2 缓存装不下激活值的时候情况会更严重。比如, QKV 投影的 low-rank 分支中向上投影要慢很多, 因为其输出超过了可用的 L2 缓存, 并导致额外的加载和存储操作到 DRAM。
幸运的是,作者观察到低秩分支中向下投影 与低位分支中的量化内核共享相同的输入,而向上投影 与 4-bit 计算 Kernel 共享相同的输出, 如图 7 所示。通过将下投影与量化 Kernel 融合,将上投影与4位计算 Kernel 融合, low-rank 分支可以与低 bit 分支共享激活, 消除额外的内存访问, 并将内核调用的数量减半。因此, 低秩分支仅增加了 5 10% 的延迟, 使其几乎免成本。
1.6 SVDQuant 支持的模型数据集和评价指标
模型
1) FLUX.1
FLUX.1[7]是基于 SoTA 开源 DiT 的扩散模型。它由 19 个 joint attention block[8]和 38 个 parallel attention block [9]组成,总共 12B 参数。作者评估了两个模型,分别是 50 步引导蒸馏 (FLUX.1-dev) 和 4 步时间步蒸馏 (FLUX.1-schnell) 变体。
2) PixArt-Σ
PixArt-Σ[10]是另一个基于 DiT 的模型。它不是使用联合注意力,而是堆叠由 self-attention、cross-attention 和 FFN 组成的 28 个 attention block,相当于 600M 参数。作者在默认的 20 步设置中对其进行评估。
3) SDXL
Stable Diffusion XL (SDXL)[11]是一种广泛使用的基于 U-Net 的模型,具有 2.6B 参数。它预测具有 3 个分辨率尺度的噪声。最高分辨率阶段完全由 ResBlocks 处理,而其他 2 个阶段联合使用 ResBlocks 和 attention 层。与 PixArt-Σ 一样,SDXL 使用 cross attention 进行文本调节。作者评估其 30 步变体以及它的 4 步蒸馏变体 SDXL-Turbo[12]。
数据集
继之前的工作[13][14][15]之后,作者随机抽取 COCO Captions 2024[16]中的提示进行校准。为了评估本文方法的泛化能力,作者采用了两种不同的提示集,具有不同的样式进行基准测试:MJHQ-30K[17]和 Densely Captioned Images (DCI)[18]的 summarized 版本 (sDCI)。
作者将 SVDQuant 与以下 Post-Training Quantization (PTQ) 方法进行比较:
4 位 NormalFloat (NF4) 是仅权重量化的数据类型[19]。它假设权重服从正态分布,并且是信息论最优的 4-bit 表示。作者使用社区量化的 NF4 FLUX.1 模型[20]作为基线。 ViDiT-Q[15]使用 per-token 量化和 smoothing 缓解不同 batch 和 token 的异常值,并在 PixArt-Σ 上实现无损 8-bit 量化。 MixDQ[14]识别 text embedding 的句首标记中的异常值,并通过 16-bit 预计算保护它们。这种方法最多可实现 W4A8 量化,SDXL-Turbo 的性能退化可以忽略不计。 TensorRT[21]包含一个行业级别的 PTQ 工具包,用于将扩散模型量化为 8-bit。它使用 smoothing,并且仅使用百分位数方案在选定的时间步长范围内校准激活。
评价指标
对于图像质量评估,使用 FID[22]来衡量生成的图像和真实图像之间的分布距离。 使用 Image Reward 来近似生成图像的人工评分[23]。 使用 LPIPS (越低越好) 来衡量感知相似度[24]和峰值信噪比 PSNR 来衡量 16-bit 模型中图像的数值相似度。
实现细节。对于 8-bit 设置,作者对 activation 使用 per-token dynamic 量化,对 weight 使用 per-channel 的量化,low-rank 分支为 16。对于 4-bit 的设置,作者对 weight 和 activation 都采用 per-group 对称量化,low-rank 分支为 32。对于 FLUX.1 模型,自适应归一化中线性层的输入保持在 16-bit (即 W4A16)。对于其他模型,交叉注意中的键和值投影保持在 16-bit,因为它们的延迟仅占总运行时间的 5% 以上。
1.7 实验结果
作者在图 8 中报告了各种模型和精度级别的定量结果,并在图 9 中展示了相应的 4-bit 的定性比较。在所有的模型中,本文的 8-bit 结果可以完美地反映 16-bit 结果,实现了高于 21 的 PSNR,击败了所有其他 8-bit 基线。在 FLUX.1-dev 上,本文的 INT8 PSNR 在 MJHQ 上甚至达到了 27。
对于 4-bit 量化,在 FLUX.1 上,本文的 SVDQuant 在 Image Reward 上超过了 NF4 W4A16 基线。在 schnell 变体上,Image Reward 甚至超过了原始 BF16 模型,这表明人类偏好更强。
在 PixArt-Σ 上,虽然本文 INT4 Image Reward 显示出轻微的退化,但 FP4 模型获得了比 FP16 模型更高的分数。这可能是因为 PixArt-Σ 的小尺寸 (600M 参数),它已经高度紧凑,并且受益于较小的 group size。值得注意的是,本文的 INT4 和 FP4 结果在所有指标上都始终优于 ViDiT-Q 的 W4A8 结果。
对于基于 UNet 的模型,在 SDXL-Turbo 上,本文的 4-bit 模型显着优于 MixDQ 的 W4A8 结果,本文的 FID 分数与 FP16 模型相当,表明性能没有损失。在 SDXL 上,本文的 INT4 和 FP4 结果都实现了与代表了 8 位 SoTA 的 TensorRT 的 W8A8 性能相当。
显存节省和加速
在图 10 中,作者报告了 FLUX.1 的测量模型大小、内存节省和加速。
本文的 INT4 量化将原始 Transformer 大小从 22.2 GiB 降低到 6.1 GiB,包括 0.3 GiB 的 low-rank 分支,整体减少了 3.6 倍。
由于 weight 和 activation 都被量化,与仅 NF4 权重量化的变体相比,本文的推理引擎 Nunchaku 甚至节省了更多的内存占用,并且在桌面级和笔记本电脑级 NVIDIA RTX 4090 GPU 上提供了 3.0 倍的加速。
值得注意的是,虽然原始的 BF16 模型需要 16GB 笔记本电脑 4090 上的每层 CPU offloading,但本文的 INT4 模型完全适合 GPU 内存,通过避免 offloading,带来了 10.1 倍的加速。作者预计 NVIDIA 的下一代 Blackwell GPU 上 FP4 量化模型的加速比更大,因为它们固有地支持组量化的微缩放,而不需要专门的 GPU Kernel。
集成 LoRA
之前的量化方法需要融合 LoRA 分支并在集成 LoRA 时重新量化模型。相比之下,本文的 Nunchaku 消除了冗余内存访问,允许添加单独的 LoRA 分支。在实践中,可以通过稍微增加 rank 将 LoRA 分支融合到本文的 low-rank 分支中,进一步提高效率。图 11 展示了一些可视化示例,将 5 种不同风格的 LoRA (现实主义、Ghibsky 插图、动漫、儿童草图和 Yarn Art)应用于本文的 INT4 FLUX.1-dev 模型。INT4 模型成功地适应了每种风格,同时保留了 16-bit 版本的图像质量。
增加 rank 的影响
图 12 显示了 PixArt-Σ 上 SVDQuant 中不同 rank r 的结果。将 rank 从 16 增加到 64 可以显着提高图像质量,但会增加参数和延迟开销。在本文的实验中,作者选择了 32 的 rank,它提供了不错的质量,开销很小。
参考
^https://blog.fal.ai/auraflow/
^https://blackforestlabs.ai/
^EfficientDM: Efficient Quantization-Aware Fine-Tuning of Low-Bit Diffusion Models
^Quarot: Outlier-free 4-bit inference in rotated llms
^Spinquant–llm quantization with learned rotations
^Qserve: W4a8kv4 quantization and system co-design for efficient llm serving
^https://blackforestlabs.ai/
^Scaling Rectified Flow Transformers for High-Resolution Image Synthesis
^Scaling Vision Transformers to 22 Billion Parameters
^Pixart-sigma: Weak-to-strong training of diffusion transformer for 4k text-to-image generation
^SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis
^Adversarial Diffusion Distillation
^Q-Diffusion: Quantizing Diffusion Models
^abMixdq: Memory-efficient few-step text-to-image diffusion models with metric-decoupled mixed precision quantization
^abVidit-q: Efficient and accurate quantization of diffusion transformers for image and video generation
^Microsoft COCO Captions: Data Collection and Evaluation Server
^Playground v2.5: Three insights towards enhancing aesthetic quality in text-to-image generation
^A Picture is Worth More Than 77 Text Tokens: Evaluating CLIP-Style Models on Dense Captions
^QLORA: Efficient Finetuning of Quantized LLMs
^https://github.com/lllyasviel/stable-diffusion-webui-forge/discussions/981
^https://developer.nvidia.com/blog/tensorrt-accelerates-stable-diffusion-nearly-2x-faster-with-8-bit-post-training-quantization/
^Gans trained by a two time-scale update rule converge to a local nash equilibrium
^Imagereward: Learning and evaluating human preferences for text-to-image generation
^The unreasonable effectiveness of deep features as a perceptual metric