大模型训练效率太低，存储空间太大？有一个新的开源技术能优化这个流程！

为应对这一挑战，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.1Checkpoint存储格式

在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灵活扩展的效果。其中最核心的模块是参数路由表的更新，路由表包含了模型参数名称、文件路径、机器结点三个关键信息。

在分布式训练中训练节点在加载(Load)模型⽂件阶段，通过路由表获取模型参数所在节点位置以及参数文件名称，通过网络传输获取模型参数；当分布式策略发生发变化模型参数的分布也随之变化，模型参数在保存(Save)模型文件阶段，各个节点需要更新参数的所在节点以及模型所在文件信息到全局路由表。

在Unified Checkpoint的设计原则中，需要确保每个参数被保存时需要是一个完整的张量，而不是被多个节点切分保存，这样可以做到存储下来的Checkpoint文件与分布式策略解耦。

如果同一个参数被多个节点保存，那在分布式策略发生变化时参数的切分将会极为复杂。上图中我们通过Gather的通信将参数进行融合，然后再将参数信息同步到路由表中。

同时为了加快模型参数的保存，每个worker进程都会并发保存部分的完整张量，提升保存效率，下面是参数融合的示意图。

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

• 原地加载的情况，主要指当前训练进程所需要的参数文件均在其所在机器，此时主要依赖Safetensors格式的惰性加载特点，实现零冗余加载。每个worker 各自读取参数文件中所需的张量部分即可，无需进行跨卡或者跨机发送。原地加载的情况已经基本可以适配大多数的训练场景，在确保不同机器为共享存储的情况下，则可以灵活地进行不同分布式策略之间的切换。
• 动态加载的情况，主要指训练进程所需要的参数文件不在当前机器，需要进行跨卡或者跨机之间的张量发送。这种情况往往发生在机器为飞共享存储，且需要切换分布式策略的情况。为了解决这一问题，我们实现了自适应动态加载机制。通过构造了张量动态路由表，用于确认哪些worker需要负责Tensor的发送，哪些进程需要负责接收Tensor。在构造完成动态路由表之后，即可进行不同卡之间的Send、Recv异步通信。在每个worker获取到所需要的参数之后，即可完成Checkpoint的加载，实现自适应动态加载，进而实现不同分布式策略、不同机器之间的切换。

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.2Checkpoint压缩最高可节省 78.5 %存储空间

在大模型训练过程中，为了能够在训练中断后从最新的Checkpoint快速恢复，我们需要不断地保存模型训练的Checkpoint。理论上，保存的Checkpoint频次越⾼，恢复训练时所需重训的step就越少。

但是大模型的 Checkpoint会占用大量的磁盘空间，从⽽限制了能够保存的Checkpoint数量。那么，我们是否可以对模型的 Checkpoint进行压缩以节省空间呢？当我们使⽤AdamW优化器和混合精度策略来训练一个参数量为Φ的⼤模型时，模型参数和优化器状态的总占⽤预计会达到14Φ的大小。

72B参数的模型，会占用接近1TB的磁盘空间。从下图展示的Checkpoint各权重参数所占⽤的空间量可以看出，FP32 格式的模型参数、动量1和动量2占据了绝大部分的磁盘空间。因此，这部分是我们需要重点优化的对象。

2.2.1AdamW优化器更新方式

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

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

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

θt = θt−1 − γ ∗ Adam_ratio

2.2.2Checkpoint压缩方案

通过大量的推导和实验我们在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)

其中bnt = 2quant_bit−1 − 1

将优化器的一阶动量（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：加载和存储的优化配置:

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

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

python -u -m paddle.distributed.launch \
      --gpus "0,1,2,3,4,5,6,7" \
      run_finetune.py \
      ./config/llama/sft_argument.json \
      --unified_Checkpoint 1 \
      --unified_Checkpoint_config "async_save remove_master_weight" \
      --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