🤖 machine learning

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

This paper introduces Invariant Bit Packing (IBP), a novel lossless compression algorithm that integrates seamlessly into ML pipelines to eliminate GPU memory bottlenecks and significantly accelerate GNN training, DLRM embedding lookups, and LLM inference without the accuracy trade-offs associated with lossy compression.

Original authors: Aditya K Kamath, Arvind Krishnamurthy, Marco Canini, Simon Peter

Published 2026-06-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Aditya K Kamath, Arvind Krishnamurthy, Marco Canini, Simon Peter

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Too-Big Suitcase"

Imagine you are a master chef (the GPU) trying to cook a massive feast (a Machine Learning model). You have a very fast kitchen, but your refrigerator (the GPU memory) is tiny. It can only hold a few ingredients at a time.

However, the recipes you need to follow require thousands of pounds of ingredients stored in a giant warehouse across town (the CPU memory or hard drive).

Every time you need a new ingredient, you have to send a delivery truck (the PCIe bus) to the warehouse to fetch it. The problem is that the highway connecting the warehouse to your kitchen is narrow and slow. Even though your kitchen is incredibly fast at chopping and cooking, you spend most of your time just waiting for the truck to arrive. This is the bottleneck.

The Old Solution: "Squishing" the Ingredients (Lossy Compression)

To fix this, people tried to "squish" the ingredients before putting them on the truck. This is called lossy compression.

The Analogy: Imagine taking a fluffy pillow, squeezing out all the air, and packing it into a tiny box. You save a lot of space on the truck.
The Catch: When you get to the kitchen, the pillow is now flat and hard. You can't use it for the recipe anymore because it lost its shape. In the world of AI, this "squishing" changes the data slightly, which can ruin the accuracy of the model. For businesses, even a tiny drop in accuracy is unacceptable.

The New Solution: "The Magic Packing List" (Lossless Compression)

The authors of this paper propose a different way to pack the truck. They call their method Invariant Bit Packing (IBP).

Instead of squishing the ingredients, they look for redundancy.

The Analogy: Imagine you are packing 100 identical boxes of cereal. You notice that every single box has the same red stripe on the top. Instead of painting a red stripe on all 100 boxes, you paint a red stripe on one master list (the Metadata) and tell the truck driver, "Hey, every box in this shipment has a red stripe on top."
The Result: You don't paint the stripes on the boxes anymore. You just ship the boxes without the stripes, and the master list. When the boxes arrive at the kitchen, the chef looks at the list, remembers "Oh right, red stripe goes here," and instantly restores the boxes to their original state. Nothing is lost; it's just packed more efficiently.

How IBP Works (The "Magic" Steps)

Finding the Patterns: The system looks at a huge pile of data (tensors) and asks, "Which parts of these numbers are always the same?" In AI data, certain bits (the smallest units of information) often stay the same across thousands of different data points, just like the red stripe on the cereal boxes.
Stripping the Redundancy: The system removes those "always-the-same" bits from the data being sent. It saves a tiny note (the Mask and Bitval) in the kitchen's memory that says, "For this group of data, the 3rd bit is always a 1."
The Fast Delivery: Because the data is now smaller, the truck carries less weight and moves faster down the narrow highway.
Instant Restoration: When the data arrives at the GPU, the system uses the tiny note to instantly re-insert the missing bits. Because the GPU is so good at doing many things at once, it can "re-inflate" the data almost instantly, faster than the truck could have driven the full load.

Why This is Special

Most previous attempts to compress data for AI required complex math that slowed down the GPU, or they risked ruining the data quality.

Lossless: It guarantees the data comes out exactly as it went in. No accuracy is lost.
GPU-Friendly: The authors designed the "unpacking" process to happen inside the GPU using its own super-fast workers (called warps). This means the GPU doesn't have to wait for the slow CPU to help unpack the boxes.
Easy to Use: They built tools that fit into existing AI software (like PyTorch), so developers can just flip a switch to use it.

The Results: Faster Feasts

The team tested this on three types of AI tasks:

GNNs (Graph Neural Networks): Used for things like social networks or fraud detection.
- Result: Training became 74% faster.
DLRMs (Recommendation Models): Used by stores to suggest products.
- Result: Looking up data became 180% faster.
LLMs (Large Language Models): The chatbots and writing assistants.
- Result: Inference (generating answers) became 24% faster.

Summary

The paper introduces a clever way to pack AI data by removing "duplicate" information that is always the same, saving a tiny note instead. This makes the data smaller for the slow highway (PCIe) but allows the fast kitchen (GPU) to restore it instantly without losing any quality. It's like sending a smaller truck that arrives sooner, letting the chef cook much faster.

Technical Summary: Reducing the GPU Memory Bottleneck with Lossless Compression for ML

Problem Statement

Modern machine learning (ML) workloads, including Graph Neural Networks (GNNs), Deep Learning Recommendation Models (DLRMs), and Large Language Models (LLMs), frequently process datasets ranging from hundreds of gigabytes to terabytes. These data volumes exceed the capacity of available GPU memory, necessitating on-demand data transfers from CPU memory, disk, or networked storage via the PCIe interconnect.

The primary performance constraint in these scenarios is the PCIe bandwidth, which is over an order of magnitude lower than GPU memory bandwidth. Consequently, execution is often bottlenecked by data transfer latency rather than GPU compute capability. While techniques like caching and overlapping compute with transfers offer partial relief, they cannot fully mitigate the bottleneck, especially when cache misses occur.

Existing solutions often rely on lossy compression (e.g., quantization) to reduce data volume. However, lossy methods introduce unpredictable accuracy degradation that varies by model and workload. In commercial deployments, even minor accuracy losses (e.g., 0.1%) are often unacceptable due to revenue risks, making lossy compression difficult to adopt as a general-purpose solution.

Lossless compression offers a deployment-friendly alternative that preserves model accuracy but faces significant implementation challenges in ML pipelines:

Latency Competition: CPU-GPU transfer latencies are competitive with GPU decompression latencies; traditional compression algorithms often add overhead that negates transfer savings.
Memory Constraints: Decompression often requires large intermediate buffers or complex metadata structures that compete with model inputs for scarce GPU memory.
Parallelism Mismatch: Many standard lossless algorithms (e.g., Huffman coding, LZ4) rely on dictionary lookups or irregular memory access patterns that hinder GPU warp parallelism and coalesced memory accesses.

Methodology: Invariant Bit Packing (IBP)

The authors propose Invariant Bit Packing (IBP), a novel lossless compression algorithm designed specifically for ML tensor transfers. The core insight is that ML tensors, even when numerically diverse, often exhibit low-entropy invariant bits across groups of tensors. For example, in floating-point representations, sign and exponent bits often remain constant across many tensors due to regularization and training dynamics.

Core Algorithm

IBP operates in three main stages:

Preprocessing (Metadata Generation):
- The system analyzes a batch of input tensors to identify bit positions that are invariant (i.e., have the same value) across a high percentage ( $T$ ) of the tensors.
- It generates a fixed-size Mask (indicating which bit positions are invariant) and a Bitval (storing the invariant values).
- This metadata is stored in GPU memory (typically kilobytes in size) and reused for the entire dataset, avoiding per-tensor metadata overhead.
Compression:
- Input tensors are processed in chunks. For each chunk, the algorithm checks if the bits at the invariant positions match the stored Bitval.
- If they match, the invariant bits are removed (eliminated), and the remaining bits are shifted to fill the gap.
- A participation bit is prepended to each chunk to indicate whether it was compressed or left uncompressed (to handle cases where compression yields no space savings).
- Compression is performed on the GPU using zero-copy Unified Virtual Memory, allowing the GPU to access CPU memory directly without explicit CPU-initiated DMA setup for small transfers.
Decompression:
- Decompression is optimized for GPU warp parallelism. A warp of 32 threads cooperatively reads a contiguous segment of compressed data from CPU memory into a shared memory workspace.
- Threads use warp primitives (low-latency communication within a warp) to perform a prefix scan, calculating the starting bit offset for each thread's assigned chunk.
- Threads then independently reconstruct their data by re-inserting the invariant bits from the global Bitval into the correct positions.
- The process utilizes asynchronous memory accesses to overlap data fetching with decompression logic, hiding PCIe latency.

System Integration

The authors integrate IBP into ML pipelines in two primary ways:

Cache Compression: Compressing tensors stored in GPU software caches (e.g., in GNN frameworks) to increase effective cache capacity.
PCIe Transfer Compression: Compressing tensors in CPU memory before transfer. To handle variable-sized compressed tensors, the system uses in-place compression (allocating fixed-size buffers but only reading the compressed portion) and maintains a bitmask in GPU memory to track which tensors are compressed.

Key Contributions

Analysis of ML Data Invariance: The paper identifies and quantifies the presence of invariant bits across diverse ML workloads (GNN, DLRM, LLM) and data types (FP32, FP16, BF16), demonstrating that structured distributions allow for significant lossless compression without accuracy loss.
Invariant Bit Packing (IBP) Algorithm: A GPU-optimized lossless compression technique that removes invariant bits, stores them in minimal global metadata, and utilizes warp-parallel decompression to achieve high throughput.
System Integration: Implementation of IBP into popular ML frameworks (DGL, Legion, Colossal-AI, FlexGen, InfiniGen) via easy-to-use APIs (PyTorch extensions and CUDA libraries).
Comprehensive Evaluation: Extensive benchmarking against state-of-the-art GPU-accelerated libraries (nvCOMP, ndzip-gpu) and native systems across real-world datasets.

Results

The evaluation was conducted on an NVIDIA A100 GPU with PCIe Gen 4.

Compression Performance: IBP achieves space savings of 10–12% for dense GNN datasets, 8% for DLRM weights, and up to 27% for LLM weights and KV-caches (particularly with BF16 data types).
Throughput: IBP consistently outperforms existing GPU-accelerated compression libraries. For sparse GNN datasets, it achieves a 9.7x speedup in transfer throughput compared to uncompressed transfers.
Workload Speedups:
- GNN Training: Average speedup of 74% (up to 2.7x in specific cases) by reducing PCIe transfer time and increasing cache hit rates.
- DLRM Inference: Average speedup of 180% for embedding lookups, primarily driven by efficient GPU-initiated transfers and reduced CPU overhead.
- LLM Inference: Average speedup of 24% for offloaded inference. For weight offloading (BF16), latency was reduced by 22%; for KV-cache offloading, latency was reduced by up to 27%.
Overhead: Decompression overhead is minimal; for space savings below 50% (typical for dense ML data), throughput remains above 95% of the ideal uncompressed transfer rate. Preprocessing time is moderate and occurs off the critical path.

Significance and Claims

The paper claims that IBP provides a practical, lossless solution to the GPU memory bottleneck without the accuracy risks associated with quantization or lossy compression. By leveraging the inherent structure of ML data (invariant bits) and optimizing for GPU architecture (warp parallelism, coalesced access, zero-copy), IBP effectively hides decompression overhead behind PCIe transfer latency.

The authors emphasize that IBP is a general-purpose algorithm applicable to various data types (float32, float16, bfloat16, sparse vectors) and workloads (training and inference). The ability to integrate IBP with minimal code changes (often fewer than 35 lines of Python) into existing frameworks demonstrates its potential to be adopted in commercial ML deployments to scale models beyond GPU memory limits while maintaining performance and accuracy.