使用约 30 行 Python 代码和 NVIDIA nvCOMP 降低检查点成本

训练 LLM 需要定期保存检查点。这些包含模型权重、优化器状态和梯度的完整快照会被保存到存储中，以便在中断后恢复训练。在大规模场景下，这些检查点会变得非常庞大（70B 模型为 782 GB）且频繁（每 15-30 分钟一次），成为训练预算中最大的支出项之一。大多数 AI 团队都在追求 GPU 利用率、训练吞吐量和模型质量。几乎没有团队关注检查点保存正在让他们付出多少成本。

这是一个代价高昂的疏忽。仅一个 405B 模型在 128 块 NVIDIA Blackwell GPU 上的同步检查点开销，每月就可能花费 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

而这正是大多数团队忽视的部分。在每一次同步 checkpoint 写入期间，所有 8 块 GPU 都完全空闲。没有任何操作会与 checkpoint 保存重叠——训练循环会阻塞，直到最后一个字节写入存储。

按 4.40 美元/GPU/小时（可代表按需 Blackwell GPU 云定价）和 5 GB/s 共享存储吞吐量（Lustre 或通过 InfiniBand 的 GPFS 的典型水平）计算，我们进行计算，以确定这些等待期间 GPU 空闲的成本：

• 每个 checkpoint 的写入时间：782 GB / 5 GB/s = 156.4 秒（约 2.6 分钟）
• 每月总等待时间：156.4 秒 × 48 次/天 × 30 天 = 225,216 秒 = 62.6 小时
• 这些等待期间闲置 GPU 的成本：62.6 小时 × 8 个 GPU × 4.40 美元/GPU/小时 = 2,203 美元/月

同步检查点带来的等待时间累计每月超过 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-190GB/s）实现了几乎相同的压缩率。正如我们将在下文展示的，选择哪一种取决于你的存储速度有多快。

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

ZSTD 和 ANS 使用熵编码，利用某些字节值出现频率中的统计模式——即使没有完全重复的序列也是如此。这就是为什么它们在相同数据上能达到 1.25-1.40x，而字节级编解码器则几乎没有压缩效果。在两者之间，ZSTD 的压缩比略高（在我们的基准测试中高出 1-2%），而 ANS 的压缩比与其接近，压缩吞吐量却是其 10 倍。随着存储速度变快，这种权衡会变得重要，正如我们将要展示的那样。

• 稠密 Transformer（Llama、GPT、Qwen）：所有参数都会参与每一次前向传播。约 0% 精确零值 → 约 1.25× ANS、约 1.27× ZSTD。
• MoE 模型（Mixtral、DeepSeek、OLMoE）：每个 token 只会激活一部分专家。12–14% 精确零值 → 约 1.39× ANS、约 1.40× ZSTD。

我们的基准测试使用 BF16 权重和 FP32 优化器状态（AdamW），这是当今大多数大规模训练的标准做法。使用 FP8 训练的团队（例如使用 NVIDIA Transformer Engine）会看到更低的压缩率，因为低精度数据具有更高的熵和更少的统计冗余，可供无损压缩利用的空间更小。在 FP4（NVFP4）下，量化已经去除了大部分冗余——无损压缩带来的额外收益可以忽略不计（约 5%）。不过，无论权重精度如何，优化器状态仍然保持为 FP32，而这正是检查点大小和压缩节省的主要来源。

数学原理：nvCOMP 如何节省成本

使用我们测得 ZSTD 压缩比为 1.27× 的 70B 稠密模型：

1. 不使用 nvCOMP：以 5 GB/s 的速度向磁盘写入 782 GB → GPU 等待 156 秒
2. 使用 nvCOMP ZSTD（1.27x）：以约 16 GB/s 的速度压缩 782 GB（约 49 秒），以 5 GB/s 写入 616 GB（123 秒）→ GPU 等待约 123 秒

为什么 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。

在更快的存储条件下，编解码器吞吐量决定了压缩是有益还是有害：

为便于说明，我们比较了一个 70B 稠密模型（782 GB）在 Blackwell GPU 上、三种存储速度下的 checkpoint 写入时间。由于压缩和写入是流水线式进行的——一个分块在压缩时，前一个分块写入磁盘——总 GPU 等待时间等于较慢的那个阶段：流水线等待时间 = max(压缩时间, 写入时间)。

在表 3 中，在 Blackwell GPU 上约 16 GB/s 时，ZSTD 的压缩步骤成为瓶颈，等待时间增加。ANS 在 181–190 GB/s 时从不会遇到这一限制。在 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 个 checkpoint，以及 5 GB/s 的共享存储：

稠密模型（在 Qwen2.5-3B 上测量）：

MoE 模型（在 OLMoE-1B-7B 上测量）：

节省会随着模型规模（更大的检查点 = 有更多内容可压缩）和 GPU 数量（等待期间空闲的 GPU 越多 = 每秒等待时间的成本越高）而扩大。第二个因素正是这在大规模场景下尤其严峻的原因——256 个空闲的 Blackwell GPU 每小时成本为 1,126 美元。你从检查点写入中节省的每一秒都在省钱。

行业正在转向 MoE 架构（如 DeepSeek-V3、Mixtral 和 Grok），这些架构会产生更大且更易压缩的检查点。压缩不再是锦上添花，添加几行 Python 代码就能节省成本。

集成：约 30 行 Python

集成工作量很小。下面是 torch.save / torch.load 的即插即用替代方案：

# CUDA 13.x (Blackwell):
pip install nvidia-nvcomp-cu13 cupy-cuda13x

# CUDA 12.x (Hopper):
pip install nvidia-nvcomp-cu12 cupy-cuda12x

import torch
import cupy as cp
from nvidia import nvcomp

codec = nvcomp.Codec(algorithm="Zstd")  # or "ANS" for near-ZSTD ratios at 10× throughput

def _compress_state(obj):
    """Recursively compress GPU tensors in a state dict tree."""
    if isinstance(obj, torch.Tensor) and obj.is_cuda:
        flat = obj.contiguous().reshape(-1).view(torch.uint8)
        # torch → cupy (zero-copy) → nvcomp (zero-copy) → GPU compress
        compressed = codec.encode(nvcomp.as_array(cp.from_dlpack(flat)))
        # .cpu() → bytes(): official nvcomp pattern (Cell 15 of nvcomp_basic)
        return {"__nvcomp__": True,
                "data": bytes(compressed.cpu()),
                "shape": list(obj.shape), "dtype": str(obj.dtype)}
    if isinstance(obj, dict):
        return {k: _compress_state(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return type(obj)(_compress_state(v) for v in obj)
    return obj

def _decompress_state(obj, device):
    """Recursively decompress nvCOMP tensors back to GPU."""
    if isinstance(obj, dict) and obj.get("__nvcomp__"):
        # codec.decode(bytes): official nvcomp pattern (Cell 27 of nvcomp_basic)
        decompressed = codec.decode(obj["data"])
        # nvcomp → cupy (zero-copy) → torch (zero-copy), stays on GPU
        dtype = getattr(torch, obj["dtype"].replace("torch.", ""))
        return torch.from_dlpack(cp.asarray(decompressed)).view(dtype).reshape(obj["shape"])
    if isinstance(obj, torch.Tensor):
        return obj.to(device)
    if isinstance(obj, dict):
        return {k: _decompress_state(v, device) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return type(obj)(_decompress_state(v, device) for v in obj)
    return obj

def save_compressed_checkpoint(model, optimizer, path):
    state = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
    torch.save(_compress_state(state), path)

def load_compressed_checkpoint(path, device="cuda"):
    return _decompress_state(
        torch.load(path, map_location="cpu", weights_only=True), device)

将 torch.save 替换为 save_compressed_checkpoint，将 torch.load 替换为 load_compressed_checkpoint。

就是这样。无需更改你的训练循环、模型代码或优化器配置。

如果你使用的是带有自定义 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.00x。对于检查点压缩，运行在 SM 上的基于熵的编解码器（ANS 和 ZSTD）才能带来真正的节省。在恢复期间，GPU 处于空闲状态，并在恢复训练前等待数据。SM 可用性并不是约束因素，而 SM 上的 ANS 可提供相近的吞吐量（247-264 GB/s），同时生成显著更小的文件。

开始使用

检查点压缩是你可以添加到训练流水线中的投资回报率最高的优化之一，也是最容易实现的优化之一：

• 安装 nvCOMP：pip install nvidia-nvcomp-cu13（Blackwell）或 pip install nvidia-nvcomp-cu12（Hopper）
• 立即试用：将本文中的 save_compressed_checkpoint / load_compressed_checkpoint 函数放入你的训练循环中。无需其他更改。
• 查看代码：GitHub 上提供了示例和基准测试。
• 阅读文档：nvCOMP 文档
• 深入了解：对于使用 GPUDirect Storage 的团队，nvCOMP 可以直接压缩到 GPU 缓冲区中，由 GDS 写入 NVMe，全程无需 CPU 参与。详情请参阅：https://docs.nvidia.com/cuda/nvcomp/ https://docs.nvidia.com/gpudirect-storage/api-reference-guide/

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