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

2026-05-28

AI Summary

This work addresses the severe bandwidth bottleneck in machine learning training and inference caused by frequent PCIe tensor transfers due to limited GPU memory capacity. The authors propose an efficient, lossless compression method tailored for ML workloads, centered on an Invariant Bit Packing (IBP) algorithm that identifies and eliminates invariant bits across tensor groups, achieving high-throughput, low-overhead compression with zero precision loss. By integrating warp-level parallel decompression, lightweight bit operations, and asynchronous PCIe data transfer, and providing an easy-to-integrate API, the approach seamlessly supports GNN, DLRM, and LLM frameworks. Experiments demonstrate significant performance improvements: 74% average speedup in GNN training, 180% acceleration in DLRM embedding lookups, and 24% enhancement in LLM inference throughput.

Abstract

Machine learning (ML) training and inference often process data sets far exceeding GPU memory capacity, forcing them to rely on PCIe for on-demand tensor transfers, causing critical transfer bottlenecks. Lossy compression has been proposed to relieve bottlenecks but introduces workload-dependent accuracy loss, making it complex or even prohibitive to use in existing ML deployments. We explore lossless compression as an alternative that avoids this deployment complexity. We identify where lossless compression can be integrated into ML pipelines while minimizing interference with GPU execution. Based on our findings, we introduce Invariant Bit Packing (IBP), a novel lossless compression algorithm designed to minimize data transfer time for ML. IBP identifies and eliminates invariant bits across groups of tensors, improving throughput through GPU-optimized decompression that leverages warp parallelism, low-overhead bit operations, and asynchronous PCIe transfers. We provide easy-to-use APIs, showcasing them by adding IBP support to GNN training, as well as DLRM and LLM inference frameworks. IBP achieves, on average, 74% faster GNN training, 180% faster DLRM embedding lookup, and 24% faster LLM inference.

Problem

• GPU memory bottleneck
• tensor transfer
• lossless compression
• machine learning
• PCIe

Innovation

• lossless compression
• GPU memory bottleneck
• Invariant Bit Packing
• tensor transfer optimization
• warp parallelism