检查点

检查点#

本文档描述 AReaL 的检查点系统，用于在分布式 RL 训练期间保存模型以进行评估和容错恢复。

概述#

AReaL 提供两种互补的检查点机制：

机制

用途

格式

包含优化器/数据加载器状态

Saver

导出模型用于评估或发布

HuggingFace

否

RecoverHandler

故障后恢复训练

DCP（分布式检查点）

是

两种机制在训练期间自动调用，可分别通过 config.saver 和 config.recover 进行配置。

检查点格式#

HuggingFace 格式#

由 Saver 用于模型导出：

  • 标准 HuggingFace 模型格式（safetensors + config.json）

  • 兼容 transformers.AutoModel.from_pretrained()

  • 可上传到 HuggingFace Hub

  • 不包含优化器状态

DCP 格式（分布式检查点）#

由 RecoverHandler 用于容错：

  • 后端原生分布式检查点格式（torch.distributed.checkpoint 或 Megatron 分布式检查点）

  • 跨所有 rank 分片以实现高效的并行 I/O

  • 包含模型权重、优化器状态、RNG 状态等

  • 后端特定：检查点仅与相同并行配置兼容

  • 在 single-controller 模式下写入不可变 generation，并通过 LATEST 指针原子发布

  • SPMD 模式继续使用原有的覆盖写入布局

以 with_optim=True 保存的 FSDP 检查点还包含 lr_scheduler.pt，用于在加载优化器时恢复学习率调度进度。 这适用于 DCP 和 HF 训练检查点。恢复训练时，应保持学习率调度配置和训练总步数不变。缺少该文件的旧检查点仍可加载， 但会输出警告并保留当前调度器状态，无法精确恢复原来的学习率轨迹。仅加载模型权重（with_optim=False）时不会恢复调度器状态。

架构#

PPOTrainer.train()
│
├── 训练循环
│   ├── Rollout、计算值、PPO 更新...
│   │
│   ├── _save_hf()                          # HuggingFace 导出
│   │   └── Saver.save()
│   │       └── engine.save(weight_format="hf")
│   │
│   └── _save_recover_checkpoint()          # 容错
│       └── RecoverHandler.dump()
│           └── engine.save(weight_format="dcp", with_optim=True)
│
└── 重启时
    └── RecoverHandler.load()
        ├── 恢复 dataloader、saver、evaluator 状态
        └── engine.load(weight_format="dcp", with_optim=True)

Saver：HuggingFace 模型导出#

Saver 定期以 HuggingFace 格式导出模型权重，用于评估或部署。

保存模式#

mode 参数控制检查点的写入方式：

模式

行为

auto

对 Archon 引擎使用异步，其他使用同步（默认）。对所有引擎都是零配置最优选择。

sync

始终同步 dcp.save()。

async

始终使用进程内异步，带有固定内存暂存。仅适用于 Archon 引擎；其他引擎回退到同步并发出警告。额外的 CPU 固定内存与每 rank 模型分片大小成正比（例如，70B 模型在 8 GPU 上每 rank 约 17.5 GB）。

使用默认的 auto 模式，Archon 引擎用户自动获得异步检查点保存 —— 训练循环仅在检查点暂存到固定 CPU 内存时阻塞，实际磁盘 I/O 在后台进程中发生。

配置#

通过 config.saver 配置：

参数

类型

默认值

描述

mode

str

"auto"

保存模式（见上文）。

freq_epochs

int | None

None

每 N 个 epoch 保存一次。None 禁用。

freq_steps

int | None

None

每 N 步保存一次。None 禁用。

freq_secs

int | None

None

每 N 秒保存一次。None 禁用。

配置示例：

saver:
  freq_epochs: 1      # 每个 epoch 结束时保存
  freq_steps: null    # 禁用
  freq_secs: null      # 禁用
  # mode 默认为 "auto" - Archon 用户自动获得异步

当满足 epoch/步数/时间任一条件时触发保存。

输出位置#

检查点保存到：

{fileroot}/checkpoints/{user}/{experiment_name}/{trial_name}/default/
└── epoch{E}epochstep{S}globalstep{G}/
    ├── config.json
    ├── model.safetensors（或 model-00001-of-00002.safetensors 等）
    ├── tokenizer.json
    └── ...

使用方式#

使用标准 HuggingFace API 加载保存的检查点：

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "/path/to/checkpoint/epoch0epochstep99globalstep99"
)
tokenizer = AutoTokenizer.from_pretrained(
    "/path/to/checkpoint/epoch0epochstep99globalstep99"
)

RecoverHandler：容错#

RecoverHandler 通过保存完整训练状态实现故障后恢复训练。

配置#

通过 config.recover 配置：

参数

类型

默认值

描述

mode

str

“disabled”

恢复模式：”on”/”auto” 或 “off”/”disabled”

freq_epochs

int | None

None

每 N 个 epoch 检查点

freq_steps

int | None

None

每 N 步检查点

freq_secs

int | None

None

每 N 秒检查点

retries

int

3

启用恢复时的恢复重试次数

恢复模式#

模式

行为

on 或 auto

如果存在有效检查点则自动恢复

off 或 disabled

不进行检检查点或恢复

启用恢复（on/auto）时，系统将：

  1. 定期保存恢复检查点（模型权重、优化器状态、dataloader 位置）

  2. 重启时自动从最后一个有效检查点恢复

  3. 失败时重试最多 retries 次

配置示例：

recover:
  mode: on            # 或 "auto"（向后兼容）
  freq_steps: 100     # 每 100 步检查点
  retries: 3

保存的内容#

RecoverHandler 保存完整训练状态：

组件

内容

模型权重

DCP 格式，跨 rank 分片

优化器状态

动量、方差（Adam）、学习率调度器

RNG 状态

Python、NumPy、PyTorch、CUDA 随机状态

Dataloader 状态

数据集中的当前位置

训练进度

Epoch、step、global_step 计数器

辅助状态

Saver、Evaluator、StatsLogger 状态

输出位置#

在 single-controller 模式下，恢复检查点保存为不可变 generation：

{fileroot}/checkpoints/{user}/{experiment_name}/{trial_name}/
├── LATEST                      # 原子 JSON 指针，指向一个完整 generation
└── checkpoint_generations/
    └── generation_step00000123/
        ├── manifest/           # Dataloader 和训练进度元数据
        └── payloads/
            ├── default/        # 模型 + 优化器（DCP 格式）
            └── critic/         # 保存 critic 时存在

只有所有 payload 都持久化后才更新指针。当一个 generation 包含多个训练引擎时，RecoverHandler 会先等待除最后一个异步引擎以外的所有保存完成，再调度最后一个异步引擎。发布操作只追加到该引擎的 MCore finalize callback，因此仍可让一次后台保存与训练重叠，同时不会暴露只写完 actor 或 critic 的 generation。 进程崩溃可能在磁盘上留下未发布的 generation，但恢复仍使用 LATEST 选中的 generation。指针成功切换后才删除上一个已发布 generation。

为保持兼容性，SPMD trainer 继续使用原有的 recover_info/ 和 <engine>/recover_checkpoint/ 布局。恢复逻辑可以读取两种布局。

恢复过程#

恢复训练时：

  1. RecoverHandler.load() 恢复所有保存的状态（如有）

  2. 训练从 last_step_info.next().global_step 继续

  3. 推理引擎权重同步以匹配恢复的状态

最佳实践#

频率指南#

场景

推荐设置

长时间训练

freq_epochs: 1 或 freq_steps: 1000

不可预测的时间

freq_secs: 7200

不稳定的集群

freq_steps: 100 配合 recover.mode: on

磁盘空间有限

较低频率，依赖最终检查点

调试

freq_steps: 1 以快速迭代

磁盘空间注意事项#

  • Saver：每次保存创建新目录。高频率会消耗大量空间。

  • RecoverHandler：通常保留一个已发布 generation，以及可能存在的正在写入或崩溃遗留 generation。SPMD 模式继续覆盖原有路径。

恢复技巧#

  1. 验证检查点有效性：读取 LATEST，再检查其选中 generation 的 manifest/step_info.json

  2. 需要相同配置：DCP 检查点需要相同的并行配置、实验名称和试次名称

  3. 每个 trial 仅一个 writer：不要让多个检查点 writer 并发写入同一实验和 trial 目录

  4. 干净重启：同时删除 LATEST 和 checkpoint_generations/。对于旧版 SPMD 检查点， 删除 recover_info/ 和每个引擎的 recover_checkpoint/