NVIDIA检查点压缩实操：70秒通关存储优化

跨境从业伙伴们，新媒网跨境了解到，在训练大规模的语言模型（LLM）时，保存检查点（Checkpointing）是避免数据丢失的重要环节。但这一过程也带来了显著的成本压力。在数据中心规模化部署中，GPU由于检查点同步写入而出现闲置时间，可能会造成高额的时间和经济成本。今天，我们将分享如何利用NVIDIA的nvCOMP技术，通过数据压缩降低存储及成本，实现高性价比的解决方案。

跨境AI训练难题：LLM检查点中的隐形成本

在训练大规模语言模型时，我们需要定期保存模型的权重、优化器状态和梯度数据，这就是所谓的“检查点”。确保数据的安全和容错性是一个不可忽视的重要环节。然而，问题的关键在于这些数据量过大所带来的存储负担，以及存储写入导致的GPU闲置成本。

举个例子：

一个70B参数规模的模型，每次检查点写入会消耗约782GB存储空间。假设每天进行48次检查点，每月存储需求高达1.13PB。
对于一个8-GPU的小型集群，存储相关的GPU闲置成本可能达到每月约人民币1万5千元；而对于128-GPU的集群，比如训练405B参数模型，这个成本可能突破人民币130万元。

新媒网跨境认为，这些潜在成本对于企业主和跨境技术团队而言，既是预算上的挑战，也可能影响项目的长期规划。

检查点的组成：要解决问题，先了解问题

从一个典型的检查点来看，其数据结构包括以下几个部分：

模型权重：以BF16格式存储，占用130GB。
优化器状态：涉及动量和方差，以FP32格式存储，占用521GB，这是目前存储中占比最大的部分。
梯度信息：以BF16格式存储，与权重大小一致，占用130GB。

这种数据分布不难看出，优化器状态由于FP32精度占用了更多的存储空间，是压缩优化的核心方向。

NVIDIA nvCOMP助力：从技术到实践的压缩方案

NVIDIA nvCOMP是一个专门为GPU环境设计的无损压缩库，它能在GPU内直接对数据进行压缩，大幅减少检查点存储的大小，同时避免了数据转移到CPU所带来的额外开销。

新媒网跨境了解到，nvCOMP主要提供两种压缩算法：

ZSTD：适用于存储带宽在5-10GB/s的环境，压缩效果优秀，能实现约1.27倍的压缩比。
gANS：针对高性能存储环境（带宽超过15GB/s，如GPUDirect Storage），兼具压缩效果与高吞吐率，其在高性能存储场景中能真正释放存储潜力。

两种算法各有优劣：如果你的存储系统速度配合较慢（例如共享文件系统），建议优先选择ZSTD。而如果你拥有高性能的存储硬件，那么gANS可能是更优解决方案。

压缩效果与数据分析

根据测试，nvCOMP能有效减少检查点的总体文件体积。以下是一些具体数据：

对于70B规模的模型，使用ZSTD可降低存储体积约21%，而使用gANS则能进一步提升性能。
如果是MoE（专家混合）架构模型，由于路由算法的稀疏结构，其检查点能压缩至原始大小的71%以下，压缩比提升约28%-29%。

新媒网跨境预测，这样的存储优化不仅能缓解存储压力，更能帮助企业每月节省大量GPU闲置成本。

实战指南：30行Python代码轻松集成

今天我们为大家展示如何通过约30行Python代码实现nvCOMP的检查点压缩，从而快速降低文件存储需求。

import torch
import cupy as cp
from nvidia import nvcomp

codec = nvcomp.Codec(algorithm="Zstd")  # 可选择为“ANS”

def save_compressed_checkpoint(model, optimizer, path):
    # 保存模型和优化器状态数据，并进行压缩
    state = {"model": model.state_dict(), "optimizer": optimizer.state_dict()}
    compressed_state = {
        k: codec.encode(nvcomp.as_array(cp.from_dlpack(v.contiguous().flatten())))
        if isinstance(v, torch.Tensor) and v.is_cuda else v
        for k, v in state.items()
    }
    torch.save(compressed_state, path)

def load_compressed_checkpoint(path, device="cuda"):
    # 读取并解压压缩后的检查点
    state = torch.load(path, map_location="cpu")
    decompressed_state = {
        k: torch.from_dlpack(cp.asarray(codec.decode(v)))
        if isinstance(v, bytes) else v
        for k, v in state.items()
    }
    return decompressed_state

只需用以上代码替换原始的torch.save与torch.load方法，即可快速实现检查点压缩与解压，操作简单易行。