随着大模型工程技术的迅猛进步，提升大模型训练效率已成为推动其发展的关键要素。训练效率 = 训练吞吐 × 训练有效率 × 收敛效率，其中，训练有效率的保障离不开灵活且强大的模型恢复机制。据悉，Meta的万卡集群在训练Llama3.1时，平均每3小时便遭遇一次故障，这凸显了完善的大模型Checkpoint保存与恢复功能对于提升训练效率的重要性。
为应对这一挑战，PaddleNLP大语言模型套件针对大规模训练场景，研发并开源了大模型统一存储技术——Unified Checkpoint。该技术集训推一体、异步保存、快速恢复、无损压缩等多重优势于一身，显著优化了大模型训练流程。以下是其主要功能亮点：
工业级的大模型训练根据机器数量的变化，会涉及多种分布式策略的灵活组合和调整。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方法无法将训练保存模型参数直接拿来预测。
Unified Checkpoint设计了存储与分布式策略解耦的方案，做到Checkpoint灵活扩展的效果。其中最核心的模块是参数路由表的更新，路由表包含了模型参数名称、文件路径、机器结点三个关键信息。在分布式训练中训练节点在加载（Load）模型文件阶段，通过路由表获取模型参数所在节点位置以及参数文件名称，通过网络传输获取模型参数；当分布式策略发生发变化模型参数的分布也随之变化，模型参数在保存（Save）模型文件阶段，各个节点需要更新参数的所在节点以及模型所在文件信息到全局路由表。
在Unified Checkpoint的设计原则中，需要确保每个参数被保存时需要是一个完整的张量，而不是被多个节点切分保存，这样可以做到存储下来的Checkpoint文件与分布式策略解耦。如果同一个参数被多个节点保存，那在分布式策略发生变化时参数的切分将会极为复杂。上图中我们通过Gather的通信将参数进行融合，然后再将参数信息同步到路由表中。同时为了加快模型参数的保存，每个worker进程都会并发保存部分的完整张量，提升保存效率，下面是参数融合的示意图。
在Checkpoint加载过程中，我们可以具体区分为两种情况，分别表示原地加载和动态加载。
原地加载的情况，主要指当前训练进程所需要的参数文件均在其所在机器，此时主要依赖Safetensors格式的惰性加载特点，实现零冗余加载。每个worker各自读取参数文件中所需的张量部分即可，无需进行跨卡或者跨机发送。原地加载的情况已经基本可以适配大多数的训练场景，在确保不同机器为共享存储的情况下，则可以灵活地进行不同分布式策略之间的切换。
动态加载的情况，主要指训练进程所需要的参数文件不在当前机器，需要进行跨卡或者跨机之间的张量发送。这种情况往往发生在机器为非共享存储，且需要切换分布式策略的情况。为了解决这一问题，我们实现了自适应动态加载机制。通过构造了张量动态路由表，用于确认哪些worker需要负责Tensor的发送，哪些进程需要负责接收Tensor。在构造完成动态路由表之后，即可进行不同卡之间的Send、Recv异步通信。在每个worker获取到所需要的参数之后，即可完成Checkpoint的加载，实现自适应动态加载，进而实现不同分布式策略、不同机器之间的切换。
训练和推理模型存储格式无缝切换
我们将模型的Checkpoint通过Unified Checkpoint技术转换为Safetensors格式进行保存，这种统一的存储方案具备高度的灵活性和兼容性，可以无缝衔接到模型推理的不同环节，无论是在动态图直接推理，还是在将动态图转换为静态图的过程中，均能实现平滑切换。这一特性消除了用户手动处理参数合并或切分的繁琐操作，不仅简化了使用流程，还显著提升了开发和部署效率，为大模型下游应用提供了更加便捷和高效的支持。
Checkpoint存储优化
存储最快加速95%
空间最高可节省78.5%
除了统一训练推理存储格式，支持灵活切换不同的分布式策略以外，我们还针对Checkpoint存储速度和存储空间进行了优化。一般的Checkpoint存储方式往往面临两大挑战：
存储耗时长，影响训练效率：对于拥有百亿乃至千亿参数的大型模型，每次Checkpoint的保存可能需要几分钟甚至十几分钟。若设置较短的保存间隔，会频繁中断训练进程，不仅浪费训练资源，还会显著降低整体训练效率。
存储空间庞大，增加存储成本：以飞桨大模型训练框架为例，一个原始的Checkpoint所占用的存储空间可能是模型权重本身所占空间的7倍。对于70B的模型，这意味着需要接近1TB的磁盘空间来存储一个Checkpoint。
为了解决上述两个问题，我们分别进行了针对性解决和优化，最终使得Unified Checkpoint的存储效率相比原先最大可加速95%，存储空间最高可节省78.5%。
异步存储大幅减少存储耗时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%。
Checkpoint压缩
最高可节省78.5%存储空间
在大模型训练过程中，为了能够在训练中断后从最新的Checkpoint快速恢复，我们需要不断地保存模型训练的Checkpoint。理论上，保存的Checkpoint频次越高，恢复训练时所需重训的step就越少。但是大模型的Checkpoint会占用大量的磁盘空间，从而限制了能够保存的Checkpoint数量。那么，我们是否可以对模型的Checkpoint进行压缩以节省空间呢？当我们使用AdamW优化器和混合精度策略来训练一个参数量为Φ的大模型时，模型参数和优化器状态的总占用预计会达到14Φ的大小。72B参数的模型，会占用接近1TB的磁盘空间。从下图展示的Checkpoint各权重参数所占用的空间量可以看出，FP32 格式的模型参数、动量1和动量2占据了绝大部分的磁盘空间。因此，这部分是我们需要重点优化的对象。
AdamW优化器更新方式
大模型训练常用的优化器为Adam优化器，其中会维护一二阶优化器动量权重：momentum和variance，每次对模型权重的更新需要先根据一二阶动量计算出Adam更新率：
最终以下公式更新模型权重：
Checkpoint压缩方案
通过大量的推导和实验我们在PaddleNLP套件提供了两种Checkpoint量化的方法，即Checkpoint O1、 O2方法，分别对优化器权重参数进行INT8、INT4的优化器压缩。
1. Checkpoint O1 压缩和恢复方法
将优化器的一阶动量（momentum）从 FP32 量化为 INT8 数据格式。最终Checkpoint中只需要存储 INT8 类型的一阶动量与其对应的量化 scales。
实验中发现，独立量化一阶动量和二阶动量，恢复后计算出来的Adam更新率会因量化损失大幅振荡。因此本方法不直接量化二阶动量，而采取直接量化 Adam更新率中的部分：
这里采取非对称均匀量化法将其从FP32量化为INT8数据格式。量化完成后，原本FP32数据类型的一二阶动量量化为INT8数据类型的格式保存，存储空间变为原来的 25%。
· 恢复步骤：
在Checkpoint恢复阶段，首先将INT8数据类型的一阶动量和Adam更新率根据各自的scales反量化为对应的FP32格式。接着需要通过FP32的Adam更新率重新计算二阶动量，使用公式：
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的使用。
1. 总开关unified_checkpoint控制是否使用Unified Checkpoint存储工具。
2. unified_checkpoint_config 包含了相关的配置项，支持如下选项:
3. 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使用文档
Paddle 链接
PaddleNLP 链接
https://github.com/PaddlePaddle/PaddleNLP
精彩课程预告
为了帮助您迅速且深入地了解PaddleNLP大语言模型套件研发开源的大模型统一存储技术——Unified Checkpoint，百度高级研发工程师将在12月24日（周四）20：00，为您详细解读Unified Checkpoint训推一体、异步保存、快速恢复、无损压缩等优势技术。机会难得，立即扫描下方二维码预约吧！
