随着大模型工程技术的迅猛进步，提升大模型训练效率已成为推动其发展的关键要素。训练效率 = 训练吞吐 × 训练有效率 × 收敛效率，其中，训练有效率的保障离不开灵活且强大的模型恢复机制。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 根据参数路由表读取参数文件中所需的张量部分即可，无需进行跨卡或者跨机发送。原地加载的情况已经基本可以适配大多数的训练场景，在确保不同机器为共享存储的情况下，则可以灵活地进行不同分布式策略之间的切换。
- 动态加载的情况，主要指训练进程所需要的参数文件不在当前机器，需要进行跨卡或者跨机之间的张量发送。这种情况往往发生在机器为非共享存储，且需要切换分布式策略的情况。为了解决这一问题，我们实现了自适应动态加载机制，通过构造了张量动态路由表，用于确认哪些 worker 需要负责 Tensor 的发送，哪些进程需要负责接收 Tensor。在构造完成动态路由表之后，即可进行不同卡之间的 Send、Recv 异步通信。在每个 worker 获取到所需要的参数之后，即可完成 Checkpoint 的加载，实现自适应动态加载，进而实现不同分布式策略、不同机器之间的切换。
实现上述加载操作的核心，本质上是依赖张量路由表的实现。在原地加载场景下，每个进程可以直接访问如图的参数路由表，根据实际需要惰性加载数据。可以看到，参数路由表包含了参数名称和参数文件。
而在动态加载情况下，参数路由表进一步升级，需要包含张量发送表和张量接收表，以指明哪些 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 通常包含两个关键阶段：
显存拷贝操作（D2H Copy）：将 GPU 显存中的参数拷贝到主机内存。
模型参数持久化（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：加载和存储的优化配置:
{
"skip_save_model_weight": "当 master_weights 存在时，跳过保存模型权重，与 remove_master_weight 不兼容开启。",
"master_weight_compatible": [
"仅当 optimizer 需要 master_weights 时，才进行加载;",
"如果 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 以其卓越的性能收益和便捷的使用效果，为大型语言模型的训练与部署提供了强有力的支持。它不仅解决了大规模训练中存储效率和存储空间的瓶颈问题，还通过其灵活性和可扩展性，为用户带来了更加高效、便捷的训练体验。
