🤖 machine learning

Reducing the GPU Memory Bottleneck with Lossless Compression for ML -- Extended

本文介绍了不变位包装（Invariant Bit Packing, IBP），这是一种全新的无损压缩算法，能够无缝集成到机器学习流水线中，以消除 GPU 内存瓶颈，并在不产生有损压缩所带来的精度权衡的情况下，显著加速图神经网络（GNN）训练、深度学习推荐模型（DLRM）嵌入查找以及大语言模型（LLM）推理。

原作者： Aditya K Kamath, Arvind Krishnamurthy, Marco Canini, Simon Peter

发布于 2026-06-01

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Aditya K Kamath, Arvind Krishnamurthy, Marco Canini, Simon Peter

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

核心问题：“太大的行李箱”

想象你是一位顶级大厨（GPU），正在准备一场盛大的宴会（机器学习模型）。你的厨房非常高效，但你的冰箱（GPU 显存）却很小，一次只能装下少量的食材。

然而，你完成食谱所需的原料需要成千上万磅，这些原料都储存在城里一个巨大的仓库里（CPU 内存或硬盘）。

每当你需要一种新食材时，你都必须派一辆送货卡车（PCIền 总线）去仓库取货。问题在于，连接仓库和厨房的高速公路非常狭窄且缓慢。尽管你的厨房在切菜和烹饪方面速度极快，但你大部分时间都在等待卡车的抵达。这就是瓶颈。

旧的解决方案：“挤压”食材（有损压缩）

为了解决这个问题，人们尝试在把食材装上卡车之前先对其进行“挤压”。这被称为有损压缩。

类比： 想象你拿一个蓬松的枕头，把里面的空气全部挤出来，然后把它装进一个小盒子里。你在卡车上节省了很多空间。
代价： 当你到达厨房时，枕头已经变得又扁又硬。你无法再用它来完成食谱，因为它失去了形状。在 AI 世界中，这种“挤压”会改变数据，从而可能破坏模型的准确性。对于企业来说，哪怕是微小的精度下降也是不可接受的。

新的解决方案：“神奇的清单”（无损压缩）

本文的作者提出了一种不同的打包方式。他们称之为不变位打包（Invariant Bit Packing, IBP）。

他们不是在挤压食材，而是在寻找冗余。

类比： 想象你要打包 100 盒完全相同的谷物盒。你注意到每一盒的顶部都有同样的红色条纹。与其在 100 个盒子上都画上红条纹，不如在一份主清单（元数据/Metadata）上画好一个红条纹，然后告诉卡车司机：“嘿，这批货物里的每一个盒子顶部都有一个红条纹。”
结果： 你不再需要在盒子上画条纹了。你只需运送没有条纹的盒子和那份主清单。当盒子到达厨房时，大厨查看清单，记起“哦对了，红条纹应该在这里”，然后瞬间恢复盒子的原貌。没有任何信息丢失；只是打包得更加高效了。

IBP 是如何工作的（“神奇”步骤）

寻找模式： 系统观察一大堆数据（张量），并询问：“这些数字中哪些部分是始终保持不变的？”在 AI 数据中，某些位（信息的最小单位）在成千上经过千上万个不同数据点中往往保持不变，就像谷物盒上的红条纹一样。
剥离冗余： 系统从要发送的数据中移除这些“始终相同”的位。它在厨房的内存中保存了一个微小的笔记（掩码/Mask 和 位值/Bitval），上面写着：“对于这组数据，第 3 位始终为 1。”
快速交付： 因为数据变小了，卡车承载的重量更轻，在狭窄的高速公路上行驶得更快。
瞬间还原： 当数据到达 GPU 时，系统利用这个微小的笔记，能够瞬间重新插入缺失的位。因为 GPU 非常擅长并行处理任务，它可以几乎瞬间“重新充气”数据，其速度比卡车运送全量货物所需的时间还要快。

为什么它很特别

以往大多数尝试为 AI 压缩数据的尝试，要么需要复杂的数学运算从而拖慢 GPU 速度，要么面临破坏数据质量的风险。

无损（Lossless）： 它保证数据进去是什么样，出来就是什么样。没有任何精度损失。
对 GPU 友好（GPU-Friendly）： 作者设计了让“解包”过程在 GPU 内部通过其自身的超快工作单元（称为 warps）完成。这意味着 GPU 不需要等待缓慢的 CPU 来协助解包。
易于使用： 他们构建的工具可以适配现有的 AI 软件（如 PyTorch），因此开发者只需切换一个开关即可使用。

结果：更快的盛宴

团队在三种类型的 AI 任务上进行了测试：

GNN（图神经网络）： 用于社交网络或欺诈检测等领域。
- 结果： 训练速度提升了 74%。
DLRM（推荐模型）： 用于商店向用户推荐产品。
- 结果： 数据查询速度提升了 180%。
LLM（大语言模型）： 聊天机器人和写作助手。
- 结果： 推理（生成答案）速度提升了 24%。

总结

本文介绍了一种巧妙的 AI 数据打包方法：通过移除那些始终不变的“重复”信息并保留一份微小的笔记来节省空间。这使得数据在缓慢的高速公路（PCIe）上传输时体积更小，同时允许快速的厨房（GPU）在不损失任何质量的情况下瞬间恢复数据。这就像是发送了一辆更小的卡车，它到达得更快，从而让大厨能更快地烹饪。

技术摘要：利用无损压缩缓解机器学习中的 GPU 显存瓶限

问题陈述

现代机器学习（ML）工作负载，包括图神经网络（GNN）、深度学习推荐模型（DLRM）和大型语言模型（LLM），经常处理从数百 GB 到数 TB 不等的庞大数据集。这些数据量超过了可用 GPU 显存的容量，因此需要通过 PCIe 互连从 CPU 内存、磁盘或网络存储进行按需数据传输。

在这些场景中，主要的性能约束是 PCIe 带宽，其带宽比 GPU 显存带宽低一个数量级以上。因此，执行过程往往受限于数据传输延迟，而非 GPU 的计算能力。虽然诸如缓存和计算与传输重叠（overlapping）等技术可以提供部分缓解，但它们无法完全消除瓶颈，尤其是在发生缓存缺失（cache miss）时。

现有的解决方案通常依赖于有损压缩（例如量化）来减少数据量。然而，有损方法会引入不可预测的精度下降，且这种下降因模型和工作负载而异。在商业部署中，即使是微小的精度损失（例如 0.1%）也往往是不可接受的，因为这涉及收入风险，使得有损压缩难以作为一种通用的解决方案。

无损压缩提供了一种对部署更友好的替代方案，它能保持模型精度，但在 ML 流水线中面临着显著的实现挑战：

延迟竞争： CPU-GPU 传输延迟与 GPU 解压延迟处于同一量级；传统的压缩算法往往会增加开销，从而抵消了传输节省下来的时间。
内存约束： 解压过程通常需要大型中间缓冲区或复杂的元数据结构，这些结构会与模型输入竞争稀缺的 GPU 显存。
并行性失配： 许多标准的无损算法（如 Huffman 编码、LZ4）依赖于字典查找或不规则的内存访问模式，这阻碍了 GPU Warp 并行性和合并内存访问（coalesced memory accesses）。

方法论：不变位打包 (Invariant Bit Packing, IBP)

作者提出了 不变位打包 (IBP)，这是一种专门为 ML 张量传输设计的创新无损压缩算法。其核心洞察在于，即使在数值多样化的情况下，ML 张量在各组张量之间仍表现出低熵不变位（low-entropy invariant bits）。例如，在浮点数表示中，由于正则化和训练动态机制，符号位和指数位在许多张量中往往保持不变。

核心算法

IBP 分三个阶段运行：

预处理（元数据生成）：
- 系统分析一批输入张量，识别在较高比例（ $T$ ）的张量中保持不变（即具有相同值）的位位置。
- 生成固定大小的掩码（Mask）（指示哪些位位置是不变的）和 Bitval（存储不变的值）。
- 该元数据存储在 GPU 内存中（通常仅为千字节级别），并为整个数据集重复使用，避免了逐个张元的元数据开销。
压缩：
- 对输入张量进行分块处理。对于每个块，算法检查不变位位置是否与存储的 Bitval 相匹配。
- 如果匹配，则移除（消除）不变位，并将剩余位进行位移以填补空隙。
- 在每个块前添加一个参与位（participation bit），用以指示该块是被压缩了还是保持原样（以处理压缩未产生空间节省的情况）。
- 使用**零拷贝（zero-copy）**统一虚拟内存在 GPU 上执行压缩，允许 GPU 直接访问 CPU 内存，而无需针对小规模传输进行显式的 CPU 发起的 DMA 设置。
解压：
- 解压过程针对 GPU Warp 并行性进行了优化。一个包含 32 个线程的 Warp 协作式地从 CPU 内存中读取一段连续的压缩数据到共享内存工作区。
- 线程使用 Warp 原语（Warp 内的低延迟通信）执行前缀扫描（prefix scan），计算出每个线程负责的块的起始位偏移量。
- 然后，线程独立地重建其数据，通过将来自全局 Bitval 的不变位重新插入到正确的位置。
- 该过程利用异步内存访问来掩盖数据获取与解压逻辑之间的延迟，从而隐藏 PCIe 延迟。

系统集成

作者通过两种主要方式将 IBP 集成到 ML 流水线中：

缓存压缩： 压缩存储在 GPU 软件缓存（例如 GNN 框架中）中的张量，以增加有效缓存容量。
PCIe 传输压缩： 在传输前对 CPU 内存中的张量进行压缩。为了处理大小可变的压缩张量，系统使用原地压缩（in-place compression）（分配固定大小的缓冲区但仅读取压缩部分），并在 GPU 内存中维护一个**位掩码（bitmask）**来追踪哪些张量已被压缩。

核心贡献

ML 数据不变性的分析： 本文识别并量化了多种 ML 工作负载（GNN、DLRM、LLM）和数据类型（FP32、FP16、BF16）中不变位的存在情况，证明了结构化分布允许在不损失精度的情况下实现显著的无损压缩。
不变位打包 (IBP) 算法： 一种经过 GPU 优化的无损压缩技术，它移除不变位，将它们存储在极小的全局元数据中，并利用 Warp 并行解压来实现高吞流。
系统集成： 通过易用的 API（PyTorch 扩展和 CUDA 库）将 IBP 集成到主流 ML 框架（DGL、Legion、Colossal-AI、FlexGen、InfiniGen）中。
全面评估： 在真实世界的数据集上，针对最先进的 GPU 加速库（nvCOMP、ndzip-gpu）及原生系统进行了广泛的基准测试。

结果

评估是在配备 NVIDIA A100 GPU 和 PCIe Gen 4 的环境下进行的。

压缩性能： IBP 在稠密 GNN 数据集上实现了 10–12% 的空间节省，DLRM 权重为 8%，而在 LLM 权重和 KV-cache（特别是使用 BF16 数据类型时）上最高可达 27%。
吞吐量： IBP 的表现持续优于现有的 GPU 加速压缩库。对于稀疏 GNN 数据集，其传输吞吐量相比未压缩传输提升了 9.7 倍。
工作负载加速：
- GNN 训练： 通过减少 PCIe 传输时间和提高缓存命中率，平均加速 74%（在特定情况下高达 2.7 倍）。
- DLRM 推理： Embedding 查找的平均加速为 180%，这主要得益于高效的 GPU 发起传输和降低的 CPU 开销。
- LLM 推理： 卸载（offloaded）推理的平均加速为 24%。对于权重卸载（BF16），延迟降低了 22%；对于 KV-cache 卸载，延迟最高降低了 27%。
开销： 解压开销极小；对于空间节省低于 50% 的情况（典型的稠密 ML 数据），吞吐量仍保持在理想未压缩传输速率的 95% 以上。预处理时间适中，且不在关键路径上。

意义与主张

本文声称，IBP 提供了一种实用的、无损的解决方案，用于解决 GPU 显存瓶颈，且不存在量化或有损压缩带来的精度风险。通过利用 ML 数据固有的结构（不变位）并针对 GPU 架构进行优化（Warp 并行、合并访问、零拷贝），IBP 有效地将解压开销隐藏在了 PCIe 传输延迟之后。

作者强调，IBP 是一种通用型算法，适用于各种数据类型（float32、float16、bfloat16、稀疏向量）和工作负载（训练与推理）。能够通过极少的代码更改（通常少于 35 行 Python 代码）将 IBP 集成到现有框架中，展示了其在维持性能与精度的同时，扩展模型规模以超越 GPU 显存限制的商业应用潜力。