随着大模型工程技术的迅猛进步，提升大模型训练效率已成为推动其发展的关键要素。训练效率 = 训练吞吐 × 训练有效率 × 收敛效率，其中，训练有效率的保障离不开灵活且强大的模型恢复机制。据悉，Meta的万卡集群在训练Llama3.1时，平均每3⼩时便遭遇⼀次故障，这凸显了完善的大模型Checkpoint保存与恢复功能对于提升训练效率的重要性。

为应对这一挑战，PaddleNLP大语言模型套件针对大规模训练场景，研发并开源了大模型统一存储技术—— Unified Checkpoint。该技术集训推一体、异步保存、快速恢复、无损压缩等多重优势于一身，显著优化了大模型训练流程。以下是其主要功能亮点：

    传统上，训练阶段Checkpoint的参数结构与推理阶段参数切分格式往往不相同，导致两者间额外的转换成本。PaddleNLP Unified Checkpoint实现了训练与推理参数格式的无感切换，极大降低了两者衔接的复杂度，进一步提升了PaddleNLP套件的便捷性。

     大模型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格式具有非常多的优势，具体可⻅下图：

https://github.com/huggingface/safetensors

Safetensors的优势主要有：

• 安全性（Safe）：Safetensors格式中不包含任何可以执行的代码，确保文件加载时不会触发任意代码执⾏。

• 零拷⻉（Zero-copy）: Safetensors格式具备接近零拷⻉的效果，在将文件加载到内存时，通过内存映射技术将文件的数据直接映射到内存，从而可避免加载过程中多个进程重复加载拷⻉同一文件的情况。

• 惰性加载（Lazy Loading）: 在分布式环境下，实现张量的部分加载是⾮常方便且效率更高的。这一点也非常有利于我们实现不同分布式策略切换的任意加载，而不需要为此创建新的权重副本，大大降低操作复杂度。

基于上述的这些特性，Unified Checkpoint的存储格式采用了Safetensors格式。

1.2 统一Checkpoint存储，支持不同分布式策略切换

原始的Checkpoint存储⽅式，在存储时会按照实际使用的不同分布式策略进行保存，从而导致在进行分布式策略切换时难以灵活扩展。具体⻅下图，固定形式的Checkpoint存储格式使得我们难以灵活应对不同的分布式情况，往往需要⼈工定制相应转换脚本，导致可扩展性低。下面是原始Checkpoint格式无法扩展的示意图。

Unified Checkpoint 设计了存储与分布式策略解耦的方案，做到Checkpoint灵活扩展的效果。为此，需要确保每个参数被保存时需要是一个完整的张量，而不是被多个节点切分保存，这样可以做到存储下来的Checkpoint文件与分布式策略解耦。如果同一个张量被多个节点保存，那在分布式策略发生变化时张量的切分将会极为复杂。在下图中我们通过Gather通信将被切分了的张量进行融合，同时为了加快模型参数的保存，每个worker进程都会并发保存部分的完整张量，提升保存效率，下面是张量融合的示意图。

在Checkpoint加载过程中，我们可以具体区分为两种情况，分别表示原地加载和动态加载。

• 原地加载的情况，主要指当前训练进程所需要的参数文件均可以直接访问，此时主要依赖Safetensors格式的惰性加载特点，实现零冗余加载。每个worker 根据参数路由表读取参数文件中所需的张量部分即可，无需进行跨卡或者跨机发送。原地加载的情况已经基本可以适配大多数的训练场景，在确保不同机器为共享存储的情况下，则可以灵活地进行不同分布式策略之间的切换。

• 动态加载的情况，主要指训练进程所需要的参数文件不在当前机器，需要进行跨卡或者跨机之间的张量发送。这种情况往往发生在机器为飞共享存储，且需要切换分布式策略的情况。为了解决这一问题，我们实现了自适应动态加载机制，通过构造了张量动态路由表，用于确认哪些worker需要负责Tensor的发送，哪些进程需要负责接收Tensor。在构造完成动态路由表之后，即可进行不同卡之间的Send、Recv异步通信。在每个worker获取到所需要的参数之后，即可完成Checkpoint的加载，实现自适应动态加载，进而实现不同分布式策略、不同机器之间的切换。

实现上述加载操作的核心，本质上是依赖张量路由表的实现。在原地加载场景下，每个进程可以直接访问如下图的参数路由表，根据实际需要惰性加载数据（图中以模型并行为2的情况进行举例）。可以看到，参数路由表包含了参数名称和参数文件。

而在动态加载情况下，参数路由表进一步升级，需要包含张量发送表和张量接收表，以指明哪些Worker负责发送相应的张量，以及哪些Worker负责接收相应的张量以及接收张量中被切分的哪一部分。

1.3 训练和推理模型存储格式无缝切换

我们将模型的 Checkpoint 通过 Unified Checkpoint 技术转换为 Safetensors 格式进行保存，这种统一的存储⽅案具备高度的灵活性和兼容性，使得模型参数可以无缝衔接到模型推理的不同环节，⽆论是在动态图直接推理，还是在将动态图转换为静态图的过程中，均能实现平滑切换。这一特性消除了用户手动处理参数合并或切分的繁琐操作，不仅简化了使用流程，还显著提升了开发和部署效率，为⼤模型下游应⽤提供了更加便捷和高效的支持。

⼆、Checkpoint存储优化，存储耗时减少95%，空间最高可节省78.5%

除了统⼀训练推理存储格式，支持灵活切换不同的分布式策略以外，我们还针对Checkpoint存储速度和存储空间进行了优化。⼀般的Checkpoint存储⽅式往往⾯临两⼤挑战：

• 存储耗时⻓，影响训练效率：对于拥有百亿乃至千亿参数的大型模型，每次Checkpoint的保存可能需要几分钟甚至十几分钟。若设置较短的保存间隔，会频繁中断训练进程，不仅浪费训练资源，还会显著降低整体训练效率。

• 存储空间庞⼤，增加存储成本：以⻜桨大模型训练框架为例，一个原始的Checkpoint所占用的存储空间可能是模型权重本身所占空间的7倍。对于70B的模型，这意味着需要接近1TB的磁盘空间来存储一个 Checkpoint。

为了解决上述两个问题，我们分别进行了针对性解决和优化，最终使得Unified Checkpoint的存储节省大模型存储时间高达95%，存储空间最高可节省78.5%。

2.1 异步存储大幅减少存储耗时95%

通常情况下，模型训练和模型存储是同步进行的。然而，在保存模型Checkpoint时，传统的做法会中断GPU训练进程以进行模型保存，这往往导致训练效率显著降低。具体来说，保存Checkpoint通常包含两个关键阶段：

1. 显存拷贝操作（D2H Copy）：将GPU显存中的参数拷贝到主机内存。
2. 模型参数持久化（Save to Disk）：将内存中的模型参数通过序列化方式保存到硬盘中。

由于这两个阶段均涉及IO操作，因此存在将计算操作（训练）与IO操作（模型保存）并行的可能性，通过重叠（Overlap）这两种操作来减少存储开销。

为了实现这一目标，我们采用了异步存储的流水线方法。通过开辟异步存储⼦进程，我们显著降低了主进程被阻塞的时间，从而提升了系统整体的吞吐量。值得注意的是，显存拷贝操作（D2H Copy）需要在训练进程中执行，而异步存储子进程则需要访问主进程中的模型参数内存。因此，在第一次进行Checkpoint保存 时，我们会分配一段连续且固定的共享内存作为主机端的存储位置。主进程将GPU端的数据拷贝到这段共享内存后，异步进程会从共享内存中读取数据，并将其持久化到磁盘。最终，PaddleNLP Unified Checkpoint的异步存储架构可简单描述如下：

在实际的存储测试中，我们对不同规模的模型进行了存储对⽐测试。以下表格对⽐中，保存介质均为NVMe SSD。我们目前测试情况下，发现存储耗时最大可减少95%。

(模型的保存时间和模型尺⼨、分布式策略有关，不⼀定和尺⼨完全成正⽐例关系)

2.2 Checkpoint压缩最高可节省 78.5%存储空间

在大模型训练过程中，为了能够在训练中断后从最新的Checkpoint快速恢复，我们需要不断地保存模型训练的Checkpoint。理论上，保存的Checkpoint频次越⾼，恢复训练时所需重训的step就越少。但是大模型的 Checkpoint会占用大量的磁盘空间，从⽽限制了能够保存的Checkpoint数量。那么，我们是否可以对模型的 Checkpoint进行压缩以节省空间呢？当我们使⽤AdamW优化器和混合精度策略来训练一个参数量为Φ的⼤模型时，模型参数和优化器状态的总占⽤预计会达到14Φ的大小。72B参数的模型，会占用接近1TB的磁盘空间。从下图展示的Checkpoint各权重参数所占⽤的空间量可以看出，FP32 格式的模型参数、动量1和动量2占据了绝大部分的磁盘空间。因此，这部分是我们需要重点优化的对象。

2.2.1 AdamW优化器更新方式

大模型训练常⽤的优化器为 Adam 优化器，其中会维护一二阶优化器动量权重：momentum 和 variance，每次对模型权重的更新需要先根据一二阶动量计算出 Adam 更新率：

Adam_ratio = momentum/(sqrt(variance) + ϵ)

最终以下公式更新模型权重：

θt = θt−1 − γ ∗ Adam_ratio

2.2.2 Checkpoint压缩方案

通过大量的推导和实验我们在PaddleNLP套件提供了两种Checkpoint量化的⽅法，即 Checkpoint O1、 O2 方法，分别对优化器权重参数进行 INT8、INT4 的优化器压缩。

Checkpoint O1 压缩和恢复方法

• 压缩步骤

     使⽤channel-wise对称均匀量化方法，以权重的channel维度为粒度计算每个 channel 的最大绝对值，得到这个 channel 的 scales 张量（FP32）。通过 scales 张量，以公式：

quant_weight = clip(round(weight/scales ∗ bcnt), − bnt − 1, bnt)

将优化器的一阶动量（momentum）从 FP32 量化为 INT8 数据格式。最终 Checkpoint 中只需要存储 INT8 类型的一阶动量与其对应的量化 scales。

实验中发现，独立量化一阶动量和二阶动量，恢复后计算出来的 Adam 更新率会因量化损失大幅振荡。因此本方法不直接量化二阶动量，⽽采取直接量化 Adam更新率中的部分：

1/(sqrt(v) + ϵ)

这⾥采取飞对称均匀量化法将其从 FP32 量化为INT8数据格式。量化完成后，原本 FP32 数据类型的一二阶动量量化为 INT8 数据类型的格式保存，存储空间变为原来的 25%。

• 恢复步骤：

    在Checkpoint恢复阶段，⾸先将 INT8 数据类型的一阶动量和 Adam 更新率根据各自的scales反量化为对应的 FP32 格式。接着需要通过 FP32 的 Adam 更新率重新计算二阶动量，使用公式：

variance = (1/(ratio + ϵ) − ϵ)2

Checkpoint O2 压缩和恢复方法

• 压缩步骤

    O2 压缩⽅式的压缩对象和 O1 压缩一致，即一阶动量与 Adam 更新率，但压缩方法有所不同。O2 压缩⽅式中，对一阶动量进⾏ block-wise 的 INT4 对称均匀量化，对 Adam 更新率进行 block-wise 的 INT4⾮对称均匀量化，此外为了更好的保存数据，本方法通过二进制移位的⽅式将两个 INT4 数据合并成⼀个 INT8 数据来保存，进⼀步提升Checkpoint压缩率。

• 恢复步骤

    与 O1 压缩恢复类似，但在Checkpoint恢复阶段，首先会将 INT8 数据类型的一阶动量和 Adam 更新率合并张量根据位运算恢复成 2 份 INT4 张量，再根据各自的scales反量化为对应的 FP32 格式。最后通过 FP32 的 Adam 更新率重新计算二阶动量。

总方法示意图如下：

对于压缩可能带来的精度损失，我们使⽤ Llama3-8B 在 6 个数据集上进行了共计 1000 个 steps 的SFT训 练，其中会在第 500 个 steps 时进行 Checkpoint 保存并强行让训练从 Checkpoint 中恢复，由下表所示，使⽤Checkpoint O1 和 Checkpoint O2 压缩算法后，模型训练精度是基本无损的，最终评估精度误差在千分位以下。

从上面的表格来看 Checkpoint O1 和 Checkpoint O2算法在恢复训练后整体对模型的精度无影响。

三、只需三行代码，无缝升级

想让原本1000GB 的 Checkpoint 瘦身成不到 300GB 么，只需三行！通过以下命令可快速体验 Unified Checkpoint 的使用。

• unified_checkpoint：总开关，控制是否使⽤ Unified Checkpoint 存储工具。

• unified_checkpoint_config：加载和存储的优化配置:

1. skip_save_model_weight: 当master_weights存在时，跳过保存模型权重，与remove_master_weight 不兼容开启。
2. master_weight_compatible: 
    1. 仅当optimizer需要master_weights时，才进⾏加载;
    2. 如果Checkpoint中不存在master_weights，则将model weight作为master_weights进⾏加载。
3.  remove_master_weight: 是否保存 master weight, 如果Checkpoint中不存在master_weights，则将model weight作为master_weights进⾏加载。
4.  async_save: 在保存Checkpoint⾄磁盘时做异步保存，不影响训练过程，提⾼训练效率，保存速度🚀 🚀 🚀

• ckpt_quant_stage：是否开启Checkpoint压缩，支持配置为O1(INT8)压缩和O2(INT4)压缩⽅案。

运行命令
1 python -u -m paddle.distributed.launch \
2 --gpus "0,1,2,3,4,5,6,7" \
3  run_finetune.py \
4  ./config/llama/sft_argument.json \
5  --unified_Checkpoint 1 \
6  --unified_Checkpoint_config "async_save remove_master_weight" \
7  --ckpt_quant_stage "O1"

更多参数配置选项可参考 https://github.com/PaddlePaddle/PaddleNLP/blob/develop/docs/trainer.md。

总结

PaddleNLP Unified Checkpoint是针对大模型大规模分布式训练设计和实现的开源存储工具，可以自适应感知分布式训练策略的组合变化，还可以通过异步保存策略以及⽆损压缩算法实现快速存储、降低存储空间。

在性能收益方面，Unified Checkpoint通过其独特的异步保存机制和无损压缩算法，实现了存储速度的⼤幅提升和存储空间的显著节省。异步存储技术使得模型训练与保存操作能够并行进行，大幅减少了训练过程中的等待时间，根据测试数据，存储耗时减少最高可达95%，这意味着训练过程更加流畅，训练资源的利⽤率得到了极大提高。同时，⽆损压缩算法的应用使得Checkpoint的存储空间需求大幅降低，最高可节省78.5%的存储空间。

在使用效果方面，Unified Checkpoint的统一存储格式和分布式策略解耦设计，使得模型在不同分布式策略间切换变得轻松自如，无需担⼼因策略变化而导致的恢复训练难题。这⼀特性极大地增强了训练的灵活性和可扩展性，使得⽤户能够根据实际资源情况灵活调整训练策略，而无需担⼼存储格式的兼容性问题。此外，训练和推理模型存储格式的无缝切换，更是为⽤户提供了从训练到部署的一站式解决⽅案，简化了模型应用的复杂度，提高了开发和部署效率。

综上所述，PaddleNLP Unified Checkpoint以其卓越的性能收益和便捷的使⽤效果，为大型语言模型的训练与部署提供了强有力的⽀持。它不仅解决了大规模训练中存储效率和存储空间的瓶颈问题，还通过其灵活性和可扩展性，为用户带来了更加⾼效、便捷的训练体验。

Unified Checkpoint使用文档：https://paddlenlp.readthedocs.io/zh/latest/llm/docs/unified_checkpoint.html 

Paddle 链接：https://github.com/PaddlePaddle/Paddle

PaddleNLP 链接：https://github.com/PaddlePaddle/PaddleNLP

精彩活动预告

为了让您能够迅速且深入地了解飞桨大模型套件Unified Checkpoint技术，我们在官方交流群内为大家提供了核心技术解读视频课程。同时为了帮助大家更好地上手体验Unified Checkpoint对大模型训练流程的显著优化，官方交流群内提供了详细的全流程使用教程文档，参与实战营活动体验任务算力全免费，成功打卡全部任务还将额外获得AI Studio算力会员卡，官方技术人员也会提供技术答疑。机会难得，立即扫描下方二维码预约吧！