Reducing the GPU Memory Bottleneck with Lossless Compression for ML -- Extended
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Lossless Data Compression | GNN Dense | Average Space Savings10.4 | 9 | |
| Lossless Data Compression | DLRM Weights | Average Space Savings8.3 | 9 | |
| Lossless Data Compression | LLM Weights BF16 | Average Space Savings26.71 | 9 | |
| Lossless Data Compression | GNN Sparse | Average Space Savings92.9 | 9 | |
| Lossless Compression | LLM Weights BF16 | Compression Time (ms)94.9 | 9 | |
| Lossless Compression | GNN Dense | Compression Time (ms)53.2 | 9 | |
| Lossless Compression | DLRM Weights | Compression Time (ms)34.2 | 9 | |
| Lossless Compression | LLM KV BF16 | Compression Time (ms)40.9 | 9 | |
| Lossless Data Compression | LLM KV BF16 | Average Space Savings23.43 | 9 | |
| Lossless Compression | LLM KV FP16 | Compression Time (ms)42.8 | 9 |