科学模拟和大预言模型（LLM）训练生成了大量数据，导致高性能计算（HPC）系统出现数据处理瓶颈。这一“大数据问题”促使领域科学家寻求更高效的数据减少技术。本文是已经发表在SC 2024上的《CUSZp2: A GPU Lossy Compressor with Extreme Throughput and Optimized Compression Ratio》旨在保证数据质量的前提下提高端到端的吞吐量。
1. 背景介绍
现代科学模拟和大型语言模型（LLM）的训练会产生海量数据，这给高性能计算（HPC）系统带来了严峻的挑战。传统的无损压缩技术虽然能够保证数据的完整性，但压缩率通常只有2:1，难以显著减小数据规模。为了应对这一问题，研究者们开始探索误差可控的有损压缩技术。这种方法通过引入用户可接受的误差，大幅提高了压缩率，已经在宇宙学模拟、量子电路模拟和地震成像等多个领域得到了成功应用。
随着HPC应用对超高速数据处理能力的需求日益增长，现有技术的局限性也愈发明显。例如，在斯坦福的线性相干光源（LCLS）实验中，每秒约250 GB的原始数据采集速率远超传统CPU压缩器的处理能力。同样，在分布式LLM训练中（如使用2048块NVIDIA A100 GPU训练LLaMA模型），大量的GPU间数据传输会成为系统性能的瓶颈。如果依赖CPU进行计算或数据传输，势必会显著拖慢训练进度，带来巨大的时间和经济成本。
然而，目前的GPU有损压缩方法存在一些不可忽视的不足。混合设计需要CPU和GPU的协作，频繁的数据移动和同步操作导致吞吐量严重受限。而纯GPU压缩器虽然摆脱了对CPU的依赖，但由于内存访问效率低和同步延迟问题，其性能仍未能充分发挥（见表1）。例如，cuSZ、cuSZx 和 MGARD-GPU虽然核心压缩算法在GPU上执行，但它们仍需依赖昂贵的CPU计算来完成全局同步、生成Huffman树或处理GPU内核间的通信，这显著增加了开销。另一方面，诸如cuZFP、FZ-GPU和cuSZp等纯GPU设计的压缩器虽然不依赖CPU，但它们通常未能充分利用内存带宽或受限于延迟问题，从而严重影响了GPU内核的吞吐量。
表1 现有GPU有损压缩器在吞吐量方面的关键设计
理想的GPU有损压缩器应满足以下几点：
纯GPU设计：实现完全基于GPU的操作，避免任何CPU计算和数据传输开销。
极高的吞吐量：充分利用GPU的内存带宽，并采用高效的延迟控制技术。
高压缩率和数据质量：在保证用户期望数据质量的同时，实现显著的压缩性能提升。
针对上述问题，本文提出了CUSZP2，一种完全基于GPU的误差可控有损压缩器，专为高吞吐量、高压缩率和高数据质量的需求而设计。CUSZP2摒弃了CPU的参与，彻底消除了数据传输带来的开销。通过优化内存访问模式和同步延迟控制技术，CUSZP2最大限度地发挥了GPU的内存带宽潜力，并在端到端吞吐量上大幅超越现有的解决方案。此外，CUSZP2通过创新的无损编码技术显著提高了压缩率，为大规模HPC应用提供了一种性能优越、效果卓越的解决方案。
2. CUSZP2: 概览
CUSZP2是一种基于GPU的误差控制有损压缩器，采用块级压缩策略，并将所有操作整合到单个GPU内核中，从而实现高吞吐量和优化的压缩率。其主要工作流程包括1有损转换（Lossy Conversion）；2无损编码（Lossless Encoding）；3全局前缀和（Global Prefix-Sum）；4块级拼接（Block Concatenation）四个步骤（见图1）。
图1 CuSZp2压缩kernel工作流
3. CUSZP2主要设计
3.1 异常固定长度编码
不同于传统的固定长度编码（Plain-FLE）为每一个数据分配一个相同的位宽，会导致异常值占用更多编码空间如图2所示。CUSZP2提出了一种针对数据块中异常值优化的编码策略，称为异常固定长度编码（Outlier-FLE），旨在提高整体压缩率，同时兼顾高效性。
图2 GPU 分块设计和压缩数据块时普通固定长度编码的缺点。
图3 异常固定长度编码
如果数值大小超过一个范围则视为异常值，异常值的存储和正常值使用不同的存储位宽，如图3所示，使用异常固定长度编码只需要3字节而使用固定长度编码需要5字节。CuSZp2的每一个块会使用一个字节来存储三个部分如图4所示，1bit存储是否使用异常值固定长度编码，2bits标识异常值存储字节大小1、2、3、4字节，5bits表示正常值的存储位宽。
图4 块头设计
3.2 向量化内存访问
CUSZP2 提出了两种优化内存访问的关键方法，用于提升压缩吞吐量：在每个数据块内，使用向量变量减少内存指令数量。在数据块级别，通过合并全局内存访问的方式提升内存带宽利用率。
减少内存指令数量：向量化后的代码（右图）：每次迭代将四个元素组合为一个float4变量，执行批量内存操作。内存加载（LD.E.128）和存储（ST.E.128）指令只需要执行 ele_num/4 次，从而显著减少了内存指令数量。
合并全局内存访问：每个线程迭代性地压缩多个数据块（如一次处理32个线程所覆盖的块）。合并一个warp中所有线程对全局内存的访问，使得总事务数量显著减少。
图5 向量化的方法减少指令数量
3.3 解耦回溯计算全局前缀和
全局前缀和计算是压缩流程中的重要步骤，用于确定每个数据块在最终压缩字节数组中的位置。传统方法通常采用普通链式扫描（Plain Chained-Scan）的方式，其中每个线程块必须等待其前一个线程块完成后才能继续。这种设计导致了高延迟，尤其是在处理大规模HPC数据集时，显著降低了吞吐量（如图6(a)所示）。为了克服这一问题，CUSZP2引入了解耦回溯策略（Decoupled Lookback Strategy）。与普通链式扫描不同，在串行扫描执行的同时，CUSZP2中的每个线程块都会执行回溯操作，查找其前驱线程块的状态。如果找到一个已完成的线程块，它会直接聚合该线程块的局部归约长度，而无需等待全局同步完成（如图6(b)所示）。CUSZP2的线程块在任何时刻可能处于三种状态之一：已完成（Finished）、回溯（Looking Back）或等待（Waiting）。已完成的线程块表示其本地归约和存储操作已经完成；回溯状态的线程块则会查找其前驱线程块以聚合局部结果；而等待状态的线程块则表示其本地扫描或压缩尚未完成。通过解耦回溯机制，当一个回溯线程块聚合了一个等待线程块的结果后，设备级同步能够跳过原始的链式结构，直接处理下一个未完成的线程块，从而显著减少了同步延迟。
图6 普通链式扫描、CUSZP2中的解耦回溯计算前缀和对比
4. 实验
设备：NVIDIA A100 (108 SMs, 40 GB) GPU；AMD EPYC 7742 CPUs；1 TB DDR4 内存
数据集：见表2
表2 数据集
Baseline：纯GPU：FZ-GPU、cuSZp、 cuZFP；CPU-GPU：cuSZ、cuSZx、MGARD-GPU；CUSZP2-P (采用固定长度编码)、CUSZP2-O (采用异常固定长度编码)
实验结果:图7表示CUSZP2在压缩和解压缩吞吐量方面显著优于现有GPU有损压缩器（如cuZFP、FZ-GPU和cuSZp）。实验分别在不同的误差范围（REL 1E-2、1E-3和1E-4）下，对多个典型HPC应用数据集（如CESM-ATM、HACC、RTM、SCALE等）进行测试，同时比较了CUSZP2的两种编码模式（CUSZP2-P和CUSZP2-O）。结果显示，无论在压缩还是解压缩阶段，CUSZP2的性能均表现出明显优势。
图7 吞吐量对比
表3表明CUSZP2-O在压缩率方面表现显著优于FZ-GPU和CUSZp，并展示了良好的适应性和鲁棒性。在REL 1E-2条件下，CUSZP2-O在CESM-ATM数据集上实现了 18.84~82.41 的压缩率（平均46.98），在RTM数据集上的压缩率达到了 30.12~104.18（平均61.48），而在NYX数据集上更是达到了 13.46~127.80（平均69.14）。即使在更严格的误差范围（REL 1E-4）下，CUSZP2-O的压缩率依然保持领先，例如在SCALE数据集上的压缩率为 6.31~43.19（平均19.75），在QMCPack数据集上为 3.79~7.25（平均4.25）。这种表现得益于Outlier-FLE的异常值处理能力，使其在处理含有复杂异常值的数据集时表现优异。
表3 压缩率对比
5. 总结
本文提出了CUSZP2，一种基于GPU的高吞吐量、有损压缩器，旨在解决高性能计算（HPC）和机器学习中海量数据压缩的瓶颈问题。CUSZP2通过创新设计实现了显著的性能提升，包括向量化内存访问、解耦回溯策略（Decoupled Lookback）以及动态编码模式选择（Plain-FLE和Outlier-FLE）。这些优化使其在压缩和解压缩阶段都能充分利用GPU的内存带宽和计算能力。实验结果表明，CUSZP2在多种HPC数据集上相较于现有方法（如FZ-GPU、cuZFP和CUSZp）展现了显著的吞吐量和压缩率优势。在REL 1E-2至1E-4的误差范围下，CUSZP2实现了最高的压缩率和接近GPU内存带宽上限的吞吐量（如1330.24 GB/s），同时保证了数据恢复质量。通过灵活适应多样化数据分布和处理复杂异常值，CUSZP2为HPC和机器学习领域的高效数据压缩提供了强有力的解决方案。
 |
重庆大学时空实验室（Spatio-Temporal Art Lab，简称Start Lab），旨在发挥企业和高校的优势，深入探索时空数据收集、存储、管理、挖掘、可视化相关技术，并积极推进学术成果在产业界的落地！年度有3~5名研究生名额，欢迎计算机、GIS等相关专业的学生报考！
图文｜李 政
校稿｜朱明辉
编辑｜李佳俊
审核｜李瑞远
审核｜杨广超
