Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Unweight:我们如何在不牺牲质量的情况下将 LLM 压缩 22%

原文:Unweight: how we compressed an LLM 22% without sacrificing quality Source: https://blog.cloudflare.com/unweight-tensor-compression/

2026-04-17

要在距离全球 95% 联网人口 50ms 之内进行推理,意味着必须无情地节省 GPU 内存。去年我们用基于 Rust 的推理引擎 Infire 改进了内存利用率,用模型调度平台 Omni 消除了冷启动。现在我们正在解决推理平台中下一个大瓶颈:模型权重。

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

为了解决这个问题,我们构建了 Unweight:一个无损压缩系统,能让模型权重小 15–22%,同时保留 bit 精确的输出,且不依赖任何特殊硬件。这里的核心突破在于,在快速的片上内存中解压权重并直接馈给 tensor core,避免了通过慢速主内存的额外往返。根据工作负载,Unweight 的运行时从多种执行策略中选择 —— 一些优先简单性,另一些最小化内存流量 —— 一个 autotuner 为每个权重矩阵和批大小选择最佳方案。

这篇文章深入介绍 Unweight 是如何工作的,但本着更大透明度和鼓励这个快速发展空间创新的精神,我们也发表了一份技术论文并开源了 GPU kernel。

我们在 Llama-3.1-8B 上的初步结果显示,仅 Multi-Layer Perceptron(MLP)权重就实现了约 30% 的压缩。因为 Unweight 选择性地处理 decoding 的参数,这导致模型大小减少 15-22% 和约 3 GB 的 VRAM 节省。如下图所示,这使我们能从 GPU 中榨出更多,因此在更多地方运行更多模型——使 Cloudflare 网络上的推理更便宜、更快。

*Thanks\ to\ Unweight,\ we’re\ able\ to\ fit\ more\ models\ on\ a\ single\ GPU *

为什么压缩比听起来更难

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

最近几个系统(Huff-LLM、ZipNN 和 ZipServ)显示 LLM 权重可以被显著压缩,但这些方法针对的是与我们不同的问题。ZipNN 为分发和存储压缩权重,解压发生在 CPU 上。Huff-LLM 提出了用于 decoding 的自定义 FGPA 硬件。ZipServ 确实将解压与 GPU 推理融合,但目标是消费级 GPU,这与我们的 H100 GPU 不兼容。这些都没有给我们想要的:能与我们基于 Rust 的推理引擎集成的、Hopper GPU 上的无损推理时解压。

核心挑战不是普通的压缩——BF16 权重中的 exponent byte 高度冗余,所以 entropy coding 在它们上面工作得很好。挑战是解压得足够快以至于不会拖慢推理。在 H100 上,tensor core 大部分时间都闲着等内存——但那个空闲容量不能简单地被重新用于解压。每个 GPU 计算单元只能运行解压 kernel 或矩阵乘法 kernel,不能同时运行,这是由于共享内存的限制。任何不与矩阵乘法完美重叠的 decode 延迟都会直接累加到 token 延迟上。Unweight 的答案是在快速片上共享内存中解压权重,把结果直接馈给 tensor core——但要让这在不同批大小和权重形状下高效工作,才是真正的工程所在。

模型权重如何能被有效压缩

AI 模型中的每个数字都被存储为 16 位的 “brain float”(BF16)。每个 BF16 值有三个部分:

  • Sign(1 位):正或负

  • Exponent(8 位):大小

  • Mantissa(7 位):该大小内的精确值

下面是这些权重之一的分解:

Sign 和 mantissa 在权重之间不可预测地变化——它们看起来像随机数据,无法被有意义地压缩。但 exponent 讲述了一个不同的故事。

Exponent 出乎意料地可预测

之前的研究已经确定,在训练过的 LLM 中,256 个可能的 exponent 值中,只有少数几个占主导。前 16 个最常见的 exponent 覆盖了典型层中超过 99% 的权重。信息论说你只需要约 2.6 位来表示这个分布——远少于分配的 8 位。如果你看典型 LLM 层中 exponent 值分布,你可以看到前 16 个 exponent 占所有模型权重的 99%。

典型 LLM 层中的 exponent 值分布

这就是 Unweight 利用的冗余。我们让 sign 和 mantissa 不动,只用 Huffman coding 压缩 exponent byte——一种经典技术,将短代码分配给常见值,长代码分配给稀有值。因为 exponent 分布如此倾斜,这在 exponent 流上实现了大约 30% 的压缩。我们选择性地将其应用于 MLP 权重矩阵(gate、up 和 down 投影),它们大约占模型参数的三分之二,在 token 生成期间主导内存流量。Attention 权重、embedding 和 layer norm 不被压缩。所有的优化加起来,大约相当于整体 multilayer perceptron(MLP)权重大小减少 20%,我们的技术报告中有详尽解释。

具有稀有 exponent 的少量权重被分别处理:如果一行 64 个权重中的任何一个有不在前 16 个调色板中的 exponent,整行被原样存储。这种方法在热路径中消除了逐元素分支——而不是检查每个权重的边缘情况,我们提前每行做一个决定。

GPU 内存瓶颈

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

  • High Bandwidth Memory(HBM):大,但访问相对慢。这是模型权重所在。

  • Shared memory(SMEM):很小,但极快。这是 GPU 在做数学之前暂存数据的地方。

During\ inference,\ generating\ each\ token\ requires\ reading\ the\ full\ weight\ matrix\ from\ HBM.\ The\ memory\ bus\ between\ HBM\ and\ SMEM\ is\ the\ performance\ bottleneck\ – not\ the\ math\ itself.\ Fewer\ bytes\ across\ the\ bus\ =\ faster\ token\ generation.

在推理期间,生成每个 token 需要从 HBM 通过内存总线读取完整的权重矩阵——这就是瓶颈。H100 的 tensor core 处理数字的速度远快于 HBM 给它们喂数据的速度。压缩有帮助,因为更少的字节需要跨过总线。但有一个问题:GPU 不能在压缩数据上做数学。权重必须先被解压。

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

使用压缩权重的四种方法

在推理期间使用压缩权重没有单一的最佳方法。正确的方法取决于工作负载——批大小、权重矩阵的形状,以及解压可用的 GPU 时间。Unweight 提供四个压缩执行 pipeline,每个在解压努力和计算复杂度之间有不同的平衡:完整 Huffman decode、仅 exponent decode、palette transcode,或完全跳过预处理。

*Four\ different\ execution\ pipelines *

四个 pipeline 形成一个谱。在一端,full decode 完全重建原始的 BF16 权重并将它们交给 NVIDIA 的 cuBLAS 库进行标准的矩阵乘法。这是最简单的路径,cuBLAS 在普通数据上以全速运行,但预处理步骤将最多的字节写回主内存。它在批大小小、矩阵乘法很小且自定义 kernel 开销主导的情况下工作得很好。在另一端,direct palette 完全跳过预处理。权重在模型加载时被预转码为紧凑的 4 位格式,矩阵乘法 kernel 在运行时从这些索引重建 BF16 值。零预处理成本,但 kernel 每个元素做更多的工作。

中间是两条独立的路径:一条只 decode exponent byte(将预处理流量减半),另一条在运行时转码为 4 位 palette 索引(将其降为四分之一)。两者都使用一个重建式矩阵乘法——一个加载压缩数据、在快速共享内存中重建 BF16,并将其直接馈给 tensor core 而不通过主内存往返的自定义 kernel。

为什么没有单一 pipeline 胜出

更少的预处理意味着写入 HBM 的数据更少,这更早释放内存总线。但它将更多的重建工作转移到 matmul kernel 上。这种权衡是否划算取决于情况。

在小批大小下(即 1-64 个 token),matmul 很小,所以没有太多计算可以重叠,而自定义 kernel 的固定成本占主导。Full decode + cuBLAS 通常胜出,只是因为 cuBLAS 有更低的开销。在大批大小下(即 256+ 个 token),matmul 运行得足够长,可以吸收额外的重建工作。更轻量的预处理更快完成,释放出来的总线带宽和计算重叠收回成本。Palette 或 exponent pipeline 领先。同一层内不同的权重矩阵可能偏好不同的 pipeline。“gate” 和 “up” 投影的维度与 “down” 投影不同,改变了 matmul 内执行操作的顺序,这需要不同的性能权衡。

吞吐量 vs pipeline 策略

这就是为什么 Unweight 不硬编码单一策略。运行时为每个权重矩阵在每个批大小选择最佳 pipeline,由测量目标硬件上实际端到端吞吐量的 autotuning 过程提供信息(下面会讲更多)。

重建式 matmul 如何工作

四个 pipeline 中的三个使用一个自定义矩阵乘法 kernel,将解压与计算融合。这个 kernel 从 HBM 加载压缩数据,在共享内存中重建原始 BF16 值,并将它们直接馈给 tensor core——所有这些都在一个操作中。重建的权重从未存在于主内存中。

传统解压 vs Unweight

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

在这个 kernel 内部,GPU 的线程组被分成两个角色:

  • 一个 producer 组使用专用的内存复制硬件(TMA)将压缩输入从 HBM 加载到共享内存。它暂存 sign+mantissa byte、exponent 数据(或 palette 索引),以及——对于具有稀有 exponent 的行——逐字的 exponent 行。它跑在 consumer 之前,填充一个环形缓冲区,以便数据在需要前就准备好。

  • Consumer 组通过将 exponent 与 sign+mantissa byte 组合来重建 BF16 值,然后立即将结果馈入 Hopper 的 WGMMA tensor-core 指令。重建的权重直接从汇编到计算,而不离开共享内存。

重建式 matmul 有多个变体,在每个计算单元处理多少输出 tile 以及环形缓冲区有多深方面不同。更宽的输出 tile 在大批大小下改善数据重用;更深的缓冲区在小批大小下隐藏内存延迟。Autotuner 为每个工作负载选择最佳变体。

在 decoding 和计算之间共享 GPU

在两个融合 pipeline 中,一个独立的预处理 kernel(Huffman decoder 或 palette transcoder)与重建式 matmul 并发运行。但这些 kernel 竞争 GPU 资源。

在 Hopper 上,每个计算单元(SM)有 228 KB 的共享内存。重建式 matmul 需要约 227 KB 用于其 pipeline 缓冲区和累加器 tile。Decode kernel 需要约 16 KB 用于其 Huffman 查找表。由于 227 + 16 > 228,这两个 kernel 不能共享同一个计算单元。每个分配给 decoding 的 SM 都是 matmul 少一个的 SM。

这创造了一个平衡:更多的 decode SM 意味着更快的预处理但更慢的矩阵乘法,反之亦然。最佳分配是另一个可调参数——也是为什么 autotuner 测量真实吞吐量而不是依赖启发式的另一个原因。

跨层 pipeline

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

不是每一层在运行时都需要 Huffman decoding。Unweight 将层分类为“hard“(需要 Huffman 预处理)或“easy“(使用 matmul 可以直接消费的预转码 palette 数据)。运行时在它们之间交替:

Decode\ runs\ on\ separate\ CUDA\ streams\ during\ bootstrap,\ attention,\ and\ easy\ MLP\ compute.\ By\ the\ time\ a\ hard\ layer’s\ MLP\ runs,\ its\ preprocessed\ weights\ are\ already\ waiting

当 GPU 计算一个 easy 层时——它不需要预处理——一组独立的 CUDA stream 在后台 decoding 下一个 hard 层的权重。当 easy 层完成且 hard 层的轮次到来时,它的预处理数据已在等待。双缓冲的预处理槽位确保一个 hard 层的 decode 输出在被消费时不会被覆盖。

down 投影从这种重叠中受益最多:它在 MLP 序列中最后被消费(在 gate、activation 和 up 之后),所以它的 decode 有最长的时间窗口完成。

Autotuning

有了四个 pipeline、多个 matmul kernel 变体,以及在 decoding 和计算之间可调的 SM 分配,配置空间很大。Unweight 没有硬编码单一策略,而是使用一个 autotuner,测量目标硬件上的实际端到端推理吞吐量。它扫描 gate 投影的候选配置,同时保持 up 和 down 固定,然后扫描 up,然后是 down,重复直到没有进一步改进。结果是一个 per-model 配置文件,告诉运行时为每个投影在每个批大小确切使用哪个 pipeline、matmul 变体和 SM 分配——所有这些都由测量的性能驱动而不是启发式。

一种压缩格式,多种用途

编码格式、执行 pipeline 和调度是独立的选择。同一个 Huffman 压缩的模型 bundle 可以同时服务于分发和推理:

  • 对于分发,Huffman 编码最大化压缩(总模型大小减少约 22%),减少跨网络传输模型时的传输时间。

  • 对于推理,Huffman 编码的投影可以在模型加载时被转码为 palette 中间格式,启用最高效的运行时执行,而不限制分发格式。

单个模型 bundle 不需要在打包时就承诺一种策略。运行时在飞行中为每个投影和每个批大小选择最佳执行路径。

我们的结果

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

  • 推理 bundle 模型占用减少约 13%(只压缩 gate/up MLP 投影),或分发 bundle 减少约 22%(压缩所有 MLP 投影,包括 down)。所有压缩都是 100% bit 精确无损。外推到 Llama 70B,这可以转化为大约 18-28 GB 节省,取决于配置。

  • 当前优化级别下端到端在 H100 SXM5 上测量的吞吐量开销 30-40%。开销在批大小 1 时最大(约 41%),在批大小 1024 时收窄(约 30%)。三个已知来源——小批固定成本、冗余权重 tile 重建,以及被排除的 down 投影——正在积极优化中。

这些是单个模型上的中间结果。压缩比应该泛化到其他 SwiGLU 架构(exponent 统计在不同模型规模下是一致的),但吞吐量数字针对当前 kernel 实现,并将随着优化继续而变化。我们尚未压缩 attention 权重、embedding 或 layer norm,这稀释了整体减少。

为什么这很重要

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

为了对抗这一点,几位研究人员展示了完整模型上约 30% 压缩比的有希望结果——但这些针对消费级 GPU 和不在生产规模工作的研究框架。Unweight 开发的关键洞察是,multilayer perceptron(MLP)构成了模型权重的大部分,以及推理工作负载期间相当多的计算成本。它只压缩 MLP 权重(避免在压缩收益微薄的层上的开销),专为具有紧密平衡的计算和内存的数据中心 H100 GPU 设计,并配备四个根据批大小适应的执行 pipeline,而不是使用单一方法。

不过,我们想说清楚:Unweight 不是免费午餐。片上重建增加了未压缩权重不会有的计算工作。在 Llama 3.1 8B 上,推理配置节省了大约 13% 的总模型内存,在典型服务批大小下吞吐量成本约 30%。这个差距在更大的批下收窄(预处理重叠改善),并预期随着我们优化进一步收窄——特别是,我们尚未压缩每个 MLP 层中的 down 投影(约占可压缩权重的三分之一),还有几项 kernel 改进正在积极开发中。

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

接下来是什么

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

Down 投影压缩。 Unweight 今天压缩 gate 和 up MLP 投影,但 down 投影约占可压缩权重的三分之一。由于其转置维度,这需要一个不同的 kernel 变体,我们预计这将把总模型大小减少超过 22%。

Kernel 优化。 当前 30-40% 的吞吐量开销有三个识别的来源:重建式 matmul 中的小批固定成本、大批大小下的冗余权重重建,以及缺失的 down 投影。每个都有已知的缓解路径,我们在技术论文中概述。

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

更长远来看,我们正在调查 Unweight 的架构对 Mixture-of-Experts 模型意味着什么,在那里冷 expert 必须按需取回,减少的存储将进一步降低成本。

这是一个快速变化的领域,所以我们很高兴在这里开源我们的工作,并为压缩和 GPU 效率方面不断增长的研究语料库做贡献。Unweight 是拼图的一部分,但我们希望其他研究人员发现它是一个有用的范式,可以在此基础上构建!