返回 文章 apply CMS 文章

Unweight:Cloudflare 如何无损压缩 LLM 权重 22% 并加速推理

Unweight 通过压缩 LLM 权重中冗余的指数部分,在 H100 GPU 上实现无损推理加速,模型大小减少 22%,吞吐量开销约 30-40%。

LLM模型压缩无损压缩GPU推理
成长分 / 100 77 综合收获、行动、留存与影响

Unweight:Cloudflare 如何无损压缩 LLM 权重 22% 并加速推理
为什么值得读了解如何在不牺牲模型质量的前提下,通过压缩权重解决 GPU 内存带宽瓶颈。

学习一种结合霍夫曼编码和 GPU 内核融合的实用无损压缩方法。

关键洞察
  1. LLM 权重中指数部分高度冗余,前 16 个指数值覆盖超过 99% 的权重,可用霍夫曼编码压缩约 30%。
  2. Unweight 在 GPU 共享内存中解压权重并直接馈送给张量核心,避免通过主内存往返,减少内存总线流量。
  3. 系统提供四种执行流水线(完全解码、仅指数解码、调色板转码、跳过预处理),自动调优器为每个权重矩阵和批次大小选择最佳策略。
转成行动

深入阅读

正文与原文对照

原文保真覆盖:全文原文字符:20067

在覆盖全球 95% 互联网连接人口的范围内,在 50 毫秒内运行推理意味着必须对 GPU 内存进行极致高效的利用。去年,我们通过基于 Rust 的推理引擎 Infire 提高了内存利用率,并通过模型调度平台 Omni 消除了冷启动。现在,我们正在解决推理平台的下一个主要瓶颈:模型权重。

从 LLM 生成单个 token 需要从 GPU 内存中读取每个模型权重。在我们许多数据中心使用的 NVIDIA H100 GPU 上,张量核心处理数据的速度比内存传输数据的速度快近 600 倍,这导致瓶颈不在于计算,而在于内存带宽。如果权重更小,那么每个跨越内存总线的字节都是可以避免的。

为了解决这个问题,我们构建了 Unweight:一种无损压缩系统,可以将模型权重缩小 15–22%,同时保持比特精确的输出,且不依赖任何特殊硬件。这里的核心突破在于,在快速片上内存中解压权重并直接馈送到张量核心,避免了通过慢速主内存的额外往返。根据工作负载,Unweight 的运行时从多种执行策略中选择——有些优先考虑简单性,有些则最小化内存流量——自动调优器会为每个权重矩阵和批次大小选择最佳策略。

本文将深入探讨 Unweight 的工作原理,但本着提高透明度和鼓励这一快速发展领域创新的精神,我们还发布了 技术论文 并开源了 GPU 内核

我们在 Llama-3.1-8B 上的初步结果显示,仅多层感知器(MLP)权重就实现了约 30% 的压缩。由于 Unweight 选择性地对解码参数进行压缩,这导致模型大小减少 15-22%,并节省约 3 GB 的 VRAM。如下图所示,这使我们能够从 GPU 中榨取更多性能,从而在更多地方运行更多模型——使 Cloudflare 网络上的推理更便宜、更快。

感谢 Unweight,我们能够在单个 GPU 上容纳更多模型

为什么压缩比听起来更难

越来越多的研究探索如何以创造性的方式压缩模型权重,以使推理更快和/或在更小的 GPU 上运行。最常见的是量化,这是一种通过将大的 32 位或 16 位浮点数转换为较小的 8 位或 4 位整数来减小模型权重和激活值大小的技术。这是一种有损压缩形式:不同的 16 位浮点值可以转换为相同的 4 位整数。这种精度的降低会以不可预测的方式影响响应的质量。对于服务于多样化用例的生产推理,我们知道我们需要一种无损且能保留精确模型行为的方法。

最近的几个系统(Huff-LLMZipNNZipServ)已经表明 LLM 权重可以被显著压缩,但这些方法针对的问题与我们的不同。ZipNN 压缩权重用于分发和存储,解压在 CPU 上进行。HUff-LLM 提出了用于解码的自定义 FPGA 硬件。而 ZipServ 将解压与 GPU 推理融合,但针对的是消费级 GPU,这与我们的 H100 GPU 不兼容。这些方法都没有提供我们所需的东西:在 Hopper GPU 上进行无损推理时解压,并能与我们的基于 Rust 的 推理引擎 集成。

核心挑战并非简单的压缩——BF16权重的指数字节高度冗余,因此熵编码对其效果良好。真正的挑战在于解压速度足够快,以免拖慢推理。在H100上,张量核心大部分时间处于空闲等待内存的状态——但这种空闲能力不能简单地重新用于解压。由于共享内存的限制,每个GPU计算单元要么运行解压内核,要么运行矩阵乘法内核,无法同时进行。任何未能与矩阵乘法完美重叠的解码延迟都会直接累加到令牌延迟上。Unweight的解决方案是在快速的片上共享内存中解压权重,并将结果直接馈送给张量核心——但如何在不同批次大小和权重形状下高效实现这一点,才是真正的工程所在。

模型权重如何有效压缩

AI模型中的每个数字都以16位“脑浮点”(BF16)存储。每个BF16值包含三部分:

符号位(1位):正或负

指数(8位):数量级

尾数(7位):该数量级内的精确值

以下是其中一个权重的分解方式:

符号位和尾数在不同权重间变化不可预测——它们看起来像随机数据,无法有效压缩。但指数则不同。

指数出人意料地可预测

先前的研究已证实,在训练好的LLM中,256个可能的指数值中只有少数几个占主导地位。最常见的16个指数覆盖了典型层中超过99%的权重。信息论表明,表示这种分布只需要约2.6比特——远少于分配的8比特。如果查看典型LLM层中的指数值分布,可以看到前16个指数占据了所有模型权重的99%。

典型LLM层中的指数值分布

这就是Unweight利用的冗余。我们保留符号位和尾数不变,仅使用__霍夫曼编码__压缩指数字节——这是一种经典技术,为常见值分配短码,为稀有值分配长码。由于指数分布如此偏斜,这实现了指数流约30%的压缩。我们选择性地将其应用于MLP权重矩阵(门控、上投影和下投影),这些矩阵约占模型参数的三分之二,并在令牌生成期间主导内存流量。注意力权重、嵌入和层归一化则不被压缩。总体而言,这些优化使多层感知器(MLP)权重大小减少约20%,详见我们的技术报告。

少数具有稀有指数的权重被单独处理:如果64个权重的行中任何一个权重的指数超出前16调色板,则整行按原样存储。这种方法消除了热路径中的逐元素分支——我们不是检查每个权重的边界情况,而是每行预先做出一个决策。

GPU内存瓶颈

NVIDIA H100 GPU有两种相关的内存:

高带宽内存(HBM):容量大,但访问速度相对较慢。模型权重存储于此。

共享内存(SMEM):容量小,但速度极快。GPU在执行数学运算前在此暂存数据。

*推理过程中,生成每个 token 都需要从 HBM 读取完整的权重矩阵。HBM 与 SMEM 之间的内存总线是性能瓶颈——*而非数学计算本身。总线上传输的字节越少,token 生成速度越快。

推理过程中,生成每个 token 都需要通过内存总线从 HBM 读取完整的权重矩阵——这是瓶颈所在。H100 的张量核心处理数据的速度远快于 HBM 提供数据的速度。压缩之所以有帮助,是因为需要跨越总线的字节更少。但有一个问题:GPU 无法对压缩数据进行数学运算。权重必须先解压缩。

大多数先前的工作将整个权重矩阵解压缩回 HBM,然后执行标准矩阵乘法。这有助于存储容量,但对带宽没有帮助,因为每个 token 仍然需要从 HBM 读取完整的未压缩矩阵。

使用压缩权重的四种方式

在推理过程中使用压缩权重没有单一的最佳方式。正确的方法取决于工作负载——批次大小、权重矩阵的形状以及可用于解压缩的 GPU 时间。Unweight 提供了四种压缩执行流水线,每种在解压缩工作量和计算复杂度之间有不同的平衡:完全霍夫曼解码、仅指数解码、调色板转码或完全跳过预处理。

四种不同的执行流水线

这四种流水线形成了一个谱系。在一端,完全解码完全重建原始的 BF16 权重,并将其交给 NVIDIA 的 cuBLAS 库进行标准矩阵乘法。这是最简单的路径,cuBLAS 在普通数据上全速运行,但预处理步骤将最多的字节写回主内存。它在小批次大小下效果良好,此时矩阵乘法很小,自定义内核开销占主导地位。在另一端,直接调色板完全跳过预处理。权重在模型加载时预先转码为紧凑的 4 位格式,矩阵乘法内核从这些索引中即时重建 BF16 值。零预处理成本,但内核每个元素做更多工作。

中间有两条独立的路径:一条仅解码指数字节(将预处理流量减半),另一条在运行时转码为 4 位调色板索引(将其减少到四分之一)。两者都使用重建矩阵乘法——一个自定义内核,加载压缩数据,在快速共享内存中重建 BF16,并直接将其馈送到张量核心,无需通过主内存往返。

为什么没有单一流水线胜出

更少的预处理意味着更少的数据写入 HBM,从而更快地释放内存总线。但这将更多的重建工作转移到了矩阵乘法内核上。这种权衡是否值得取决于具体情况。

当批次大小较小时(即1-64个token),矩阵乘法规模很小,因此没有太多计算可以重叠,自定义内核的固定成本占主导地位。完整的解码+cuBLAS通常胜出,仅仅是因为cuBLAS的开销更低。当批次大小较大时(即256+个token),矩阵乘法运行时间足够长,可以吸收额外的重建工作。更轻量的预处理完成得更快,释放的总线带宽和计算重叠带来了收益。调色板或指数流水线领先。同一层内的不同权重矩阵可能偏好不同的流水线。“门”和“上”投影与“下”投影的维度不同,改变了矩阵乘法内部操作的顺序,这需要不同的性能权衡。

吞吐量与流水线策略

这就是为什么Unweight不硬编码单一策略。运行时根据自动调优过程(在目标硬件上测量实际端到端吞吐量,下文详述)为每个权重矩阵和每个批次大小选择最佳流水线。

重建矩阵乘法的工作原理

四个流水线中有三个使用自定义矩阵乘法内核,将解压缩与计算融合在一起。该内核从HBM加载压缩数据,在共享内存中重建原始的BF16值,并直接馈送到张量核心——所有这些都在一个操作中完成。重建的权重永远不会存在于主内存中。

传统解压缩 vs Unweight

使用Unweight,MLP权重矩阵通过内存总线的字节数减少约30%

在此内核内部,GPU的线程组分为两个角色:

生产者组使用专用内存复制硬件(TMA)将压缩输入从HBM加载到共享内存中。它暂存符号+尾数字节、指数数据(或调色板索引),以及对于具有稀有指数的行,逐字指数行。它领先于消费者运行,填充循环缓冲区,以便数据在需要之前就绪。

消费者组通过将指数与符号+尾数字节组合来重建BF16值,然后立即将结果馈送到Hopper的WGMMA张量核心指令中。重建的权重直接从汇编进入计算,无需离开共享内存。

重建矩阵乘法有多种变体,区别在于每个计算单元处理多少个输出瓦片以及循环缓冲区的深度。更宽的输出瓦片在大批次大小下改善数据重用;更深的缓冲区在小批次大小下隐藏内存延迟。自动调优器为每个工作负载选择最佳变体。

在解码和计算之间共享GPU

在两个融合流水线中,一个单独的预处理内核(霍夫曼解码器或调色板转码器)与重建矩阵乘法并发运行。但这些内核竞争GPU资源。

在Hopper上,每个计算单元(SM)有228 KB的共享内存。重建矩阵乘法需要约227 KB用于其流水线缓冲区和累加器瓦片。解码内核需要约16 KB用于其霍夫曼查找表。由于227 + 16 > 228,这两个内核无法共享同一个计算单元。分配给解码的每个SM都意味着矩阵乘法可用的SM减少一个。

这造成了一种平衡:更多的解码SM意味着更快的预处理但更慢的矩阵乘法,反之亦然。最佳分配是另一个可调参数——也是自动调优器测量实际吞吐量而非依赖启发式方法的另一个原因。

即使有SM分区约束,Unweight通过利用Transformer模型的结构,隐藏了大部分解压成本。

并非每一层都需要在运行时进行霍夫曼解码。Unweight将层分类为“硬”(需要霍夫曼预处理)或“易”(使用可直接供矩阵乘法使用的预转码调色板数据)。运行时在它们之间交替:

在引导、注意力和易MLP计算期间,解码在单独的CUDA流上运行。当硬层的MLP运行时,其预处理权重已在等待

当GPU计算一个不需要预处理的易层时,一组单独的CUDA流在后台解码下一个硬层的权重。当易层完成且轮到硬层时,其预处理数据已在等待。双缓冲预处理槽确保一个硬层的解码输出在仍被使用时不会被覆盖。

下投影从这种重叠中获益最多:它在MLP序列中最后被消费(在门控、激活和上投影之后),因此其解码有最长的完成时间。

通过四个流水线、多个矩阵乘法内核变体以及解码和计算之间可调的SM分割,配置空间很大。Unweight不是硬编码单一策略,而是使用一个自动调优器,在目标硬件上测量实际的端到端推理吞吐量。它在固定上投影和下投影的同时,扫描门控投影的候选配置,然后扫描上投影,再扫描下投影,重复直到没有进一步改进。结果是一个每模型配置文件,告诉运行时对于每个投影和每个批次大小,确切使用哪个流水线、矩阵乘法变体和SM分配——全部由测量性能而非启发式驱动。

编码格式、执行流水线和调度是独立的选择。同一个霍夫曼压缩模型包可以同时用于分发和推理:

对于分发,霍夫曼编码最大化压缩(总模型大小减少约22%),减少通过网络传输模型时的传输时间。

对于推理,霍夫曼编码的投影可以在模型加载时转码为调色板中间格式,从而实现最高效的运行时执行,而不限制分发格式。

单个模型包不需要在打包时承诺一种策略。运行时根据每个投影和每个批次大小动态选择最佳执行路径。

在Llama 3.1 8B(我们的主要测试平台)上,Unweight实现了:

推理包模型占用减少约13%(仅压缩门控/上MLP投影),或分发包减少约22%(压缩所有MLP投影,包括下投影)。所有压缩都是100%比特精确无损。外推到Llama 70B,根据配置,这可以转化为大约18–28 GB的节省。

在当前优化水平下,端到端在H100 SXM5上测量,吞吐量开销为30–40%。开销在批次大小为1时最大(约41%),在批次大小为1024时缩小(约30%)。三个已知来源——小批次固定成本、冗余权重块重建和排除的下投影——正在积极优化中。

这些是单个模型上的中间结果。压缩比应能推广到其他 SwiGLU 架构(指数统计在不同模型规模上保持一致),但吞吐量数字特定于当前的内核实现,并会随着优化继续而变化。我们尚未压缩注意力权重、嵌入或层归一化,这些会稀释整体压缩效果。

GPU 在多个维度上都很昂贵:显卡本身的成本、它们所需的高带宽内存以及巨大的功耗。

为了解决这个问题,几位研究人员展示了在完整模型上实现约 30% 压缩比的有前景的系统——但这些系统针对的是消费级 GPU 和研究框架,无法在生产规模下工作。Unweight 开发的关键洞察在于,多层感知器(MLP)构成了模型权重的大部分,并且在推理工作负载中占据了大量的计算成本。它仅压缩 MLP 权重(避免在压缩收益微小的层上产生开销),专门针对计算和内存紧密平衡的数据中心 H100 GPU 设计,并提供了四种适应批处理大小的执行流水线,而不是使用单一方法。

然而,我们想明确一点:Unweight 并非免费的午餐。片上重构增加了计算工作,而使用未压缩的权重则不会有这些工作。在 Llama 3.1 8B 上,推理配置在典型服务批处理大小下,以大约 30% 的吞吐量成本节省了约 13% 的总模型内存。这个差距在更大的批处理大小下会缩小(此时预处理重叠得到改善),并且随着我们的优化预计会进一步缩小——特别是,我们尚未压缩每个 MLP 层中的下投影(约占可压缩权重的三分之一),并且几个内核改进正在积极开发中。

对于 Cloudflare 的网络,Unweight 为我们提供了更好的容量:它使我们能够用更少的每实例 GPU 内存来服务最先进的模型,这转化为成本节约以及在更多地方部署更多模型的能力。对于模型分发,节省更大:霍夫曼压缩包大约小 22%,减少了将模型分发到全球边缘位置的传输时间。

展望未来,我们有三个具体的研究方向,我们认为将进一步提高我们的效率收益:

下投影压缩。 Unweight 目前压缩门和上 MLP 投影,但下投影约占可压缩权重的三分之一。由于其转置维度,这需要不同的内核变体,我们预计这将使总模型大小减少超过 22%。

内核优化。 当前 30–40% 的吞吐量开销有三个已知来源:重构矩阵乘法中小批量的固定成本、大批量下的冗余权重重构以及缺失的下投影。每个都有已知的缓解路径,我们在 技术论文 中进行了概述。

更多模型。 我们的结果是针对 Llama 3.1 8B 的,但底层指数统计在所有规模的 SwiGLU 架构上是一致的。我们正在努力将 Unweight 应用于我们通过 Workers AI 服务的更大模型。

长期来看,我们正在研究 Unweight 的架构对混合专家模型意味着什么,在这些模型中,冷专家必须按需获取,而减少存储将进一步降低成本。

这是一个快速发展的领域,因此我们很高兴在此开源我们的工作,并为压缩和GPU效率方面日益增长的研究成果做出贡献。Unweight是拼图的一部分,但我们希望其他研究人员能发现它是一个有用的范式,可以在此基础上进行构建!