随着大模型工程技术的迅猛进步，提升大模型训练效率已成为推动其发展的关键要素。训练效率 = 训练吞吐 × 训练有效率 × 收敛效率，其中，训练有效率的保障离不开灵活且强大的模型恢复机制。Meta 的万卡集群在训练 Llama3.1 时，平均每 3 小时便遭遇一次故障，这凸显了完善的大模型 Checkpoint 保存与恢复功能对于提升训练效率的重要性。
为应对这一挑战，PaddleNLP 大语言模型套件针对大规模训练场景，研发并开源了大模型统一存储技术——Unified Checkpoint。该技术集训推一体、异步保存、快速恢复、无损压缩等多重优势于一身，显著优化了大模型训练流程。以下是其主要功能亮点：
1. 支持全分布式策略调整自适应转换，提升模型训练的灵活性与可扩展性
工业级的大模型训练根据机器数量的变化，会涉及多种分布式策略的灵活组合和调整。Unified Checkpoint 通过存储格式与分布式策略的解耦设计，支持 Checkpoint 在不同分布式策略间的自动切换，用户无需感知这一变化，可显著提升大模型恢复训练的灵活性与可扩展性。
2. 训练 - 压缩 - 推理统一存储协议，无需手动转换提升全流程体验
传统上，训练阶段 Checkpoint 的参数结构与推理阶段参数切分格式往往不相同，导致两者间额外的转换成本。PaddleNLP Unified Checkpoint 实现了训练与推理参数格式的无感切换，极大降低了两者衔接的复杂度，进一步提升了 PaddleNLP 套件的便捷性。
3. Checkpoint 无损压缩结合异步保存，实现秒级存储并降低模型存储成本
大模型 Checkpoint 的存储不仅耗时长，还占用庞大的磁盘空间。以 Llama-2-13B 模型为例，单个 Checkpoint 的存储耗时约 150 秒，磁盘空间消耗高达 182GB。PaddleNLP 通过异步存储和模型 Checkpoint 无损压缩算法最终可以实现秒级存储和降低 80% 左右的大模型存储空间。
一、Unified Checkpoint 统一训推存储，全面支持任意分布式策略组合
大模型 Checkpoint 存储技术面临诸多挑战，特别是对于大规模分布式训练而言：
传统 Checkpoint 存储方案无法适用复杂的大模型训练场景。在真实的大模型训练场景中，分布式策略组合非常多且复杂，同时训练资源的扩缩容导致 Checkpoint 的保存需要很好的灵活性和自适应性。例如用户因为训练集群缩容就需要 Checkpoint 恢复训练时能适应分布式策略带来的变化。
大模型训练和推理的 Checkpoint 格式不统一导致传统 Checkpoint 存储方案无法适用：与上面提到机器扩缩容类似，训练和推理的采用的分布式策略不同，导致传统的 Checkpoint 方法无法将训练保存模型参数直接拿来预测。
为了解决上述问题，我们提出了 Unified Checkpoint 统一存储的方案。
1.1 Checkpoint 存储格式
在 Checkpoint 存储格式上，我们在飞桨框架原先的存储协议上，针对大模型引入了 Safetensors 作为 Unified Checkpoint 的序列化保存格式。使用 Safetensors 格式具有非常多的优势：
Safetensors 的优势主要有：
安全性（Safe）： Safetensors 格式中不包含任何可以执行的代码，确保文件加载时不会触发任意代码执行。
零拷贝（Zero-copy）： Safetensors 格式具备接近零拷贝的效果，在将文件加载到内存时，通过内存映射技术将文件的数据直接映射到内存，从而可避免加载过程中多个进程重复加载拷贝同一文件的情况。
惰性加载（Lazy Loading）： 在分布式环境下，实现张量的部分加载是非常方便且效率更高的。这一点也非常有利于我们实现不同分布式策略切换的任意加载，而不需要为此创建新的权重副本，大大降低操作复杂度。
基于上述的这些特性，Unified Checkpoint 的存储格式采用了 Safetensors 格式。
1.2 统一 Checkpoint 存储，支持不同分布式策略切换
原始的 Checkpoint 存储方式，在存储时会按照实际使用的不同分布式策略进行保存，从而导致在进行分布式策略切换时难以灵活扩展。固定形式的 Checkpoint 存储格式使得我们难以灵活应对不同的分布式情况，往往需要人工定制相应转换脚本，导致可扩展性低。
Unified Checkpoint 设计了存储与分布式策略解耦的方案，做到 Checkpoint 灵活扩展的效果。为此，需要确保每个参数被保存时需要是一个完整的张量，而不是被多个节点切分保存，这样可以做到存储下来的 Checkpoint 文件与分布式策略解耦。如果同一个张量被多个节点保存，那在分布式策略发生变化时张量的切分将会极为复杂。我们通过 Gather 通信将被切分了的张量进行融合，同时为了加快模型参数的保存，每个 worker 进程都会并发保存部分的完整张量，提升保存效率。
在 Checkpoint 加载过程中，我们可以具体区分为两种情况，分别表示原地加载和动态加载。
- ，主要指当前训练进程所需要的参数文件均可以直接访问，此时主要依赖 Safetensors 格式的惰性加载特点，实现零冗余加载。每个 worker 根据参数路由表读取参数文件中所需的张量部分即可，无需进行跨卡或者跨机发送。原地加载的情况已经基本可以适配大多数的训练场景，在确保不同机器为共享存储的情况下，则可以灵活地进行不同分布式策略之间的切换。
