跳到正文

用大约 30 行 Python 和 NVIDIA nvCOMP 降低检查点成本

本文介绍了大模型训练中检查点写入的成本来源,并说明如何用 NVIDIA nvCOMP 在 GPU 上进行无损压缩,以降低存储与 GPU 空转带来的开销。

功思 AI 编辑部

约 8 分钟读完

用大约 30 行 Python 和 NVIDIA nvCOMP 降低检查点成本

AI 生成的内容可能会不完整地概括信息。请核实重要信息。了解更多

训练 LLM 需要定期保存检查点。这些包含模型权重、优化器状态和梯度的完整快照会被写入存储,以便训练在中断后恢复。到了大规模训练时,这些检查点会变得非常庞大(70B 模型为 782 GB),而且保存频繁(每 15 到 30 分钟一次),成为训练预算中最大的一项支出之一。大多数 AI 团队关注 GPU 利用率、训练吞吐和模型质量,几乎没人去看检查点到底花了多少钱。

这是一个代价高昂的疏忽。仅 128 张 NVIDIA Blackwell GPU 训练 405B 模型时,同步检查点开销每月就可能达到 200,000 美元。只要加入一个用大约 30 行 Python 实现的无损压缩步骤,每月就能减少 56,000 美元的存储成本。混合专家(MoE)模型节省会更多。本文将拆解这一计算过程,并说明 NVIDIA nvCOMP 如何提升检查点效率。

单个检查点内部

在 1000+ GPU 规模下,硬件中断并不少见。Meta 报告称,在 16,384 张 NVIDIA H100 GPU 上进行 Llama 3 训练的 54 天中,出现了 419 次意外中断(约每 3 小时一次)。这也是为什么大多数团队每 15 到 30 分钟就保存一次检查点;它是承载训练可靠性的基础设施,而不是可有可无的额外负担。

很多人第一次看到这个拆分时都会感到意外。优化器状态,也就是 AdamW 的一阶与二阶矩估计,且两者都以 FP32 存储,体积是模型权重的 4 倍。它占据了每个检查点的大部分空间,这也是为什么检查点远大于部署后的模型。

按照容错的标准做法,每 30 分钟保存一次检查点,相当于每天 48 个检查点。持续训练一个月后:

782 GB × 48/天 × 30 天 = 每月向存储写入 1.13 PB

而大多数团队忽略的是这一点:每次进行同步检查点写入时,全部 8 张 GPU 都会完全闲置。检查点保存不会与其他工作重叠,训练循环会一直阻塞,直到最后一个字节写入存储。

按 4.40 美元/GPU/小时计算(可视作按需 Blackwell GPU 云价格的代表值),以及 5 GB/s 的共享存储吞吐(通过 InfiniBand 使用 Lustre 或 GPFS 的典型水平),可以算出这些等待期间 GPU 闲置的成本:

同步检查点带来的等待时间累计下来,每月会超过 2,200 美元,而且这还没算存储费用。

如果扩展到一个 64 GPU 集群,每月成本会升至 17,500 多美元。到 128 张 GPU、训练 405B 模型时,闲置成本会超过每月 200,000 美元。闲置 GPU 的成本比存储费用高出一个数量级。

异步检查点可以缓解一部分问题。不过,框架支持仍在逐步成熟,尚未广泛采用,而且内存水位管理仍然具有挑战。一个容易配合使用的补充方案是检查点压缩,它还有额外好处:在恢复状态时,由于冷启动本质上是串行的,因此也能减少冷启动时间。

NVIDIA nvCOMP 引入 GPU 加速压缩

核心思路很简单:在检查点离开 GPU 内存之前先压缩。无需绕到 CPU,也不需要额外的数据搬运。数据本来就在 GPU 上,所以可以直接在那里压缩。

NVIDIA nvCOMP 就是一套实现这一点的 GPU 加速无损压缩库。它在一个库中同时支持标准算法,如 Zstandard(ZSTD),以及经过高度优化、面向 GPU 的格式,如 gANS,从而在设备端原生处理数据瓶颈。开发者可以很容易地把高吞吐压缩直接集成到 Python 工作流中,例如 PyTorch 或 TensorFlow。

实测检查点压缩率

我们对两种模型架构(稠密 Transformer 和混合专家)进行了 50 步微调,保存完整训练检查点(权重 + AdamW 优化器状态 + 梯度),并在 NVIDIA H200 和 Blackwell GPU 上用 nvCOMP 对每个组成部分进行压缩。压缩率取决于数据,而不是硬件,因此在所有 GPU 上都相同。

ZSTD 是 Meta 开发并被广泛采用的通用压缩算法,在较强压缩率和合理速度之间取得平衡。Linux 命令行里的 ZSTD 背后就是这套算法,它也被广泛用于数据库、文件系统和数据管道。非对称数制系统(ANS)是一种现代熵编码技术,nvCOMP 将其实现为以吞吐为核心优化目标的 GPU 原生编解码器(gANS)。两者都是无损的,利用的是 1B/2B 字分布中的统计规律(熵编码),而不是仅仅匹配重复的字节序列。

关键区别在于权衡。ZSTD 可以获得略高一些的压缩率(在我们的基准中比 ANS 高 1% 到 2%),但在 Blackwell GPU 上的压缩速度约为 16 到 19 GB/s。ANS 则以近 10 倍的吞吐(181 到 190 GB/s)实现几乎相同的压缩率。正如下文所示,应该选哪一个,取决于你的存储速度。

这些区间反映了一个关键发现:压缩效果取决于模型架构,而不是硬件。并不是所有压缩算法都适用于浮点张量。像 LZ4 和 Bitcomp 这样的字节级编解码器会寻找重复的字节序列,就像在文档里找重复单词一样。但训练后的神经网络参数在字节层面看起来基本是随机的,因此这些编解码器几乎找不到可压缩内容(在我们的稠密检查点基准中约为 1.00×)。

ZSTD 和 ANS 使用熵编码,即便没有精确重复的序列,也能利用某些字节值出现频率上的统计规律。这也是为什么在同样的数据上,它们可以做到 1.25 到 1.40 倍压缩,而字节级编解码器几乎没有效果。两者中,ZSTD 的压缩率略强一些(在我们的基准中高 1% 到 2%),而 ANS 则以 10 倍压缩吞吐逼近这一效果。随着存储速度提高,这种权衡会变得重要,下文会具体说明。

我们的基准使用 BF16 权重和 FP32 优化器状态(AdamW),这也是当前大多数大规模训练的标准配置。使用 FP8 训练的团队(例如结合 NVIDIA Transformer Engine)会看到更低的压缩率,因为更低精度的数据具有更高熵、可供无损压缩利用的统计冗余更少。到了 FP4(NVFP4),量化已经去除了大部分冗余,无损压缩带来的额外收益可以忽略不计(约 5%)。不过,不论权重精度如何,优化器状态仍保持为 FP32,而检查点体积和压缩节省的大头也正来自这里。

计算过程:nvCOMP 如何省钱

以我们测得 ZSTD 压缩率为 1.27× 的 70B 稠密模型为例:

为什么 49 秒的压缩时间会“消失”?因为压缩和存储写入可以流水化:当一个分块写盘时,下一个分块会在 GPU 上压缩。只要编解码器压缩速度快于存储吸收输出的速度,压缩步骤就会完全隐藏在写入之后,GPU 等待时间只等于写入时间。在 5 GB/s 的共享存储下,ZSTD 的 16 GB/s 速度是写入速度的 3 倍,因此压缩可以完全重叠。等待时间从 156 秒降到 123 秒,文件缩小 21%,每个检查点快 33 秒。按一个月计算:48 × 30 = 少了 47,520 秒闲置时间,也就是回收了 13 小时以上。MoE 在 1.40× 压缩率下效果更好:文件缩小 29%,每个检查点快 44 秒,回收 17 小时以上。

当存储变快时,编解码器吞吐会更重要

你的存储速度取决于基础设施:共享网络文件系统(Lustre、GPFS、NFS)通常是 2 到 10 GB/s,或者使用本地 NVMe 的 GPUDirect Storage(GDS)可达 15 到 50+ GB/s。

当存储更快时,编解码器吞吐会决定压缩是带来收益还是带来负担:

为了说明这一点,我们比较了在 Blackwell GPU 上,70B 稠密模型(782 GB)在三种存储速度下的检查点写入时间。由于压缩和写入是流水化的,即一个分块压缩时,前一个分块正在写盘,因此 GPU 的总等待时间等于较慢的那个阶段:流水化等待时间 = max(压缩时间, 写入时间)。

在表 3 中,当 Blackwell GPU 上的存储速度约为 16 GB/s 时,ZSTD 的压缩步骤开始成为瓶颈,等待时间反而增加。而 181 到 190 GB/s 的 ANS 不会碰到这个问题。在 25 GB/s、64 张 Blackwell GPU 的情况下,ANS 每月可净省约 2,800 美元,而 ZSTD 只能净省约 240 美元。对于共享文件系统(5 到 10 GB/s),ZSTD 仍然是合适的默认选择。对于 15+ GB/s 的 GDS 或高性能存储,ANS 更明显占优。

图 4 中测得的节省,基于以下假设:GPU 等待时间减少 + 存储费用为 0.14 美元/GB/月,保留 96 个检查点,以及 5 GB/s 的共享存储:

节省幅度会随着模型规模增大而增加(检查点越大,可压缩内容越多),也会随着 GPU 数量增加而增加(等待期间闲置的 GPU 越多,每秒等待的成本越高)。第二个因素使得这个问题在大规模下尤其严重。256 张闲置的 Blackwell GPU 每小时成本为 1,126 美元。检查点写入每减少一秒,都能省钱。

行业正在转向 MoE 架构(如 DeepSeek-V3、Mixtral 和 Grok),这类模型会产生更大、也更容易压缩的检查点。压缩已经不再是可有可无的优化,增加几行 Python 代码就可能节省成本。

集成:大约 30 行 Python

集成工作量很小。下面给出一个可直接替换 `torch.save` / `torch.load` 的实现:

把 `torch.save` 换成 `save_compressed_checkpoint`,把 `torch.load` 换成 `load_compressed_checkpoint`。

就是这样。不需要改动训练循环、模型代码或优化器配置。

如果你使用的是带有自定义检查点钩子的训练框架(如 DeepSpeed、Megatron),同样适用这一模式:遍历 state dict,压缩 GPU 张量,然后序列化。

对于使用 NVIDIA GPUDirect Storage(GDS)的团队,还有更快的路径。nvCOMP 可以直接压缩到 GDS 缓冲区中,把压缩后的数据从 GPU 内存直接写入 NVMe 存储,整个过程没有 CPU 参与。

NVIDIA Blackwell 解压引擎

NVIDIA Blackwell GPU 包含专用的 Blackwell Decompression Engine(DE),可在不占用 SM 的情况下,以最高 280 GB/s 的速度解压 LZ4、Snappy 和 Deflate。不过,字节级编解码器在浮点张量上的压缩率约为 1.00×。对于检查点压缩,真正带来收益的是运行在 SM 上、基于熵编码的编解码器(ANS 和 ZSTD)。在恢复阶段,GPU 处于空闲并等待数据,随后才恢复训练,因此 SM 可用性并不是限制因素,而运行在 SM 上的 ANS 也能提供相近的吞吐(247 到 264 GB/s),同时生成明显更小的文件。

开始使用

检查点压缩是你能加到训练流水线中、投资回报率最高且最容易实施的优化之一:

检查点是训练流水线中最大的文件。把它们压缩。

参与讨论

正在确认登录状态…