Excp & FastPersist：数十倍 LLM Checkpoint 保存加速和压缩

amei2000go

发布于 2024-7-4 08:48

浏览

0收藏

一、背景

在 LLM 的预训练过程中，通常会不断地保存 Checkpoint，主要有以下几个好处：

中断恢复：在大规模训练中，硬件故障和偶然的软件故障是不可避免的。通过保存 Checkpoint，可以在发生异常时恢复训练，避免从头开始。
进度管理：LLM 需要评估的指标很多，但通常不会在模型训练中实时评估，而是通过保存 Checkpoint 进行离线评估。这有助于评估模型是否朝着正确的方向发展。
增量学习：保持 Checkpoint 也为后续的增量学习和迁移学习提供基础，可以从中间阶段继续训练。
模型调优：在需要重新调整超参数或者跳过部分数据（比如出现 spike）时，可以回到之前的 Checkpoint 进行调整。

虽然保存 Checkpoint 有众多好处，但也会带来一些挑战。例如，频繁保存 Checkpoint 可能影响训练速度，也可能带来较大的存储压力。本文中简单介绍最近两个与 Checkpoint 相关的工作：一个是 Checkpoint 压缩，降低存储压力；另一个是加快 Checkpoint 保存速度。

二、引言

2.1 LLM 预训练的挑战

在 OPT [2205.01068] OPT: Open Pre-trained Transformer Language Models 的 57 天预训练中，前期由于各种问题至少手动重启了 35 次任务。为了避免每次需要人工重启，后续加入了自动重启机制，由于硬件故障又触发了 70+ 次，平均每天都要重启一次任务。

在 PaLM [2204.02311] PaLM: Scaling Language Modeling with Pathways 的预训练中，作者遇到了 20 次毛刺需要重启训练。

在 Falcon [2311.16867] The Falcon Series of Open Language Models 中作者也遇到了 9 次毛刺需要重启训练。

在 Bloom [2211.05100] BLOOM: A 176B-Parameter Open-Access Multilingual Language Model 的预训练中作者也频繁遇到各种异常，因此会每 3 小时保存一次 Checkpoint。

在上海 AI lab 和商汤的 [2403.07648] Characterization of Large Language Model Development in the Datacenter，作者也有同样的问题，其每 30 分钟就保存一次 Checkpoint。

在 DeepSeek [2401.02954] DeepSeek LLM: Scaling Open-Source Language Models with Longtermism 中，作者更是每 5 分钟就保存一次 Checkpoint。

如下图 Table 1 所示，阿里 C4 [2406.04594] Boosting Large-scale Parallel Training Efficiency with C4: A Communication-Driven Approach 中，作者统计了一个代表性任务在一个月内遇到的错误。数据显示，由于这些错误，该作业在一个月内失败了 40 次：

Excp & FastPersist：数十倍 LLM Checkpoint 保存加速和压缩-AI.x社区

2.2 Checkpoint 大小

在 LLM 预训练中，为了保证模型异常后能够继续训练，需要保存模型训练的中间状态，也就是 Checkpoint，Checkpoint 中除了模型参数外，还包含模型的优化器状态，以及一些超参（比如 batch size，learning rate 等，通常很小，大小可以忽略）。

以 LLM 经常使用的 Adam 优化器为例，Adam 是一种结合了动量（Momentum）和自适应学习率（如 RMSprop）技术的优化算法，广泛用于深度学习模型的训练。Adam 优化器的核心优势在于能够根据参数的更新历史自适应调整每个参数的学习率，这有助于加快收敛速度并提高训练稳定性。为了实现这一点，Adam 维护了 2 个额外的参数（或“矩”），也就是每个参数都对应 2 个优化器参数：

一阶矩 v：过去梯度的移动平均，提供了平滑过的梯度方向。这有助于增加稳定性，使模型在参数空间中的移动更加平滑。
二阶矩 m：过去梯度平方的指数移动平均，代表过去梯度的变化率，允许每个参数有自适应的学习率。参数更新的规模由其梯度的不确定性决定，有助于调整学习步长，避免在参数空间中过大或过小的步进。

以此可以推测出 LLM 预训练 Checkpoint 的大小：（GPT 175B 模型，Adam 优化器）：

模型参数：175B 参数 * 2 Byte = 350 GB。在常见的混合精度训练中，为了减少累积误差，提高训练稳定性，通常会保留一份FP32 精度的模型参数，如果是这种情况，则对应的模型参数为 700GB 左右。
Adam 优化器状态：通常使用 FP32 存储，175B * 4 Byte * 2 = 1400 GB，对应上述的一阶矩和二阶矩。

综上，LLM 预训练中显存占用大概是 10 倍参数量或 12 倍参数量。

如下图 Table 1 所示为一系列模型对应的 Checkpoint 大小，差不多是 13 倍左右：

Excp & FastPersist：数十倍 LLM Checkpoint 保存加速和压缩-AI.x社区

三、Checkpoint 压缩（ExCP）

3.1 方法

LLM 预训练时间往往比较长，可能会持续几个月。为了应对故障等问题会频繁的保存 Checkpoint，常见的为 1-3 小时保存一次，与此同时，每个 Checkpoint 都很大，可能上百 GB 到 2TB，这会给存储带来极大的压力。为了解决这个问题，常见的方式是定期、滚动地删除更早的 Checkpoint。

在 [2406.11257] ExCP: Extreme LLM Checkpoint Compression via Weight-Momentum Joint Shrinking 中，作者提出了 Checkpoint 压缩的方案 ExCP（Extreme Checkpoint Compression），其利用 Checkpoint 之间的冗余信息来压缩，在 410M 到 7B 参数量的模型上进行评估，410M 模型可以压缩 70x 大小，而基本不损失下游任务上的精度。

ExCP 具体方案如下图 Figure 2 所示，考虑到相近 Checkpoint 的模型权重相似度比较高，而优化器状态会通过滑动平均，相似度不一定特别高，因此主要是利用权重相似性来稀疏化：

左侧：相邻权重相减，求出残差，其残差的稀疏性比较高。
中间：

M2W：首先使用残差权重和二阶矩 m 确定裁剪 Mask Mw。

W2M：然后使用裁剪 Mask Mw 和一阶矩 v 确定裁剪 Mask Mo。

右侧：
使用 Mw 裁剪权重。
使用 Mo 裁剪一阶矩和二阶矩。
压缩 Checkpoint。

Excp & FastPersist：数十倍 LLM Checkpoint 保存加速和压缩-AI.x社区

PS：以上的过程考虑了两个相邻的 Checkpoint，那么针对整个训练中的 Checkpoint 如何处理呢？最简单的想法是针对原始的 Checkpoint，每一个都依赖之前的 Checkpoint 进行压缩。但是这种方式会存在一个问题，从压缩的 Checkpoint 恢复出来的只是残差，依然需要上一个原始 Checkpoint 才能恢复当前的 Checkpoint，然而此时原始的 Checkpoint 可能已经被清理。因此，作者采用的是当前 Checkpoint 与上一个已经稀疏化处理的 Checkpoint 的残差。这样只用保留最原始的 Checkpoint，之后的 Checkpoint 都顺序解压即可恢复所有 Checkpoint。当然，也可以考虑阶段性地保留一些中间的原始 Checkpoint。

3.2 结果

如下图 Table 3 所示，作者使用 Pythia-410M 进行了验证，可以从原始的 4.53GB 压缩到 0.06GB，压缩 70+ 倍；并且在下游任务上的平均指标只从 43.11 下降到 42.93：

Excp & FastPersist：数十倍 LLM Checkpoint 保存加速和压缩-AI.x社区

如下图 Figure 4 所示，在整个训练过程中对 Checkpoint 进行压缩，其 Loss 与未压缩的原始 Loss 非常接近：

Excp & FastPersist：数十倍 LLM Checkpoint 保存加速和压缩-AI.x社区

如下图 Table 4 所示，作者在 PanGu-π-1B 和 PanGu-π-7B 上进行了实验，同样在压缩 20+ 倍的情况下平均精度甚至更高了一点点：

Excp & FastPersist：数十倍 LLM Checkpoint 保存加速和压缩-AI.x社区

四、Checkpoint 快速保存（FastPersist）

4.1 方法

在 [2406.13768] FastPersist: Accelerating Model Checkpointing in Deep Learning 中，Deepspeed 的作者提出了 FastPersist，其主要是为了加快 DL 训练中 Checkpoint 保存的速度，最多可以比基线加快 116x。FastPersist 主要包含 3 点：

加快 Checkpoint 写入（下图 b）：利用节点上的 NVMe 设备的高速写入带宽，并使用新的 I/O 库（比如 Linux 中的 libaio 和 io_uring）来异步的加快写入性能。
并行化写入（下图 c）：在 DP（Data Parallelism）之间分布式写入，充分利用不同 DP 上的硬件，减少写入延迟。
Pipeline 写入（下图 d）：创建专门的 Python 进程来异步的写入。

Excp & FastPersist：数十倍 LLM Checkpoint 保存加速和压缩-AI.x社区

PS：其实之前很多文章中也提到了加快 Checkpoint 保存，避免影响训练速度的工作。主要是分布式保存，异步保存等。

4.2 结果

作者进行了一系列的实验来评估 FastPersist 的效果，结果如下图 Figure 9 所示：

Figure 9a：在 128 个 V100 GPU 上，FastPersist 获得了 28x（gpt 13B）到 116x（gpt 0.7B）的加速。
Figure 9b：随着 DP（Data Parallelism）的增加，保存速度近线性增加，对于 gpt 13b，最高达到了 146GB/s，将近是 8 个节点理论峰值的 80%。
Figure 9c：使用 FastPersist，模型的训练速度也明显提升，并且模型越小提升越明显。从 1.6x（gpt 13B）到 21.8x（gpt 0.7B）。
Figure 9d：进一步绘制了训练速度与 DP 数目的关系。

Excp & FastPersist：数十倍 LLM Checkpoint 保存加速和压缩-AI.x社区