Accelerate恢复训练时如何保存随机数状态

Accelerate恢复训练时，应使用accelerator.save_state()保存断点，再在同一训练脚本、相同对象准备完成后调用accelerator.load_state()。这组接口包含训练所需的随机数状态；只保存模型参数或重新执行set_seed，不能替代完整训练断点。

官方文档入口：https://huggingface.co/docs/accelerate/index

设想一个训练任务需要让操作者选择“从头开始”或“继续上次训练”。对操作者而言，继续应当对应一个明确断点，而不是重新走一遍随机序列。下面以完整epoch结束处保存为例，把这个选择落实到代码。示例限定单进程CPU、顺序读取的固定数据、无多worker增强和无梯度累积；它是未执行的教学代码，不是已经完成的训练实验。

为什么相同seed不等于恢复到暂停位置

seed决定随机序列的起点，状态则描述随机数发生器当前的位置。模型初始化、Dropout和随机增强都可能消耗随机数。训练已经进行一段时间之后重新设同一个seed，相当于把发生器放回起点，而不是放到断点处。

因此，“从头开始”可以使用固定seed建立初始条件；“继续训练”则需要恢复保存的状态。恢复后不要再调用set_seed，也不要为了打印调试样本先抽一批随机数，否则训练开始前就改变了该状态。

save_state负责什么，业务代码还负责什么

官方checkpoint指南将模型、优化器、随机数发生器和混合精度缩放器视为训练状态。自定义对象需要提供state_dict和load_state_dict并登记。恢复对象的结构及登记方式应与保存时一致，不能把目录中的任意文件都当成自动加载的内容。

Accelerate当前官方checkpoint实现保存Python random、NumPy全局随机状态和PyTorch CPU随机状态，并按可用设备保存相应状态，例如CUDA。这里的“NumPy全局状态”不意味着每一个default_rng实例都被自动识别。

把继续训练做成明确的入口选择

下面脚本使用可选的--resume参数表示操作者选择的断点目录。Progress保存下一次应开始的epoch。数据由arange生成，避免重新构造数据时抽取随机样本；网络保留Dropout，以体现训练期间存在随机操作。

import argparse
from pathlib import Path
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from accelerate import Accelerator
from accelerate.utils import set_seed

class Progress:
    def __init__(self):
        self.next_epoch = 0

    def state_dict(self):
        # 保存业务进度，而不是把epoch当成随机状态。
        return {"next_epoch": self.next_epoch}

    def load_state_dict(self, state):
        self.next_epoch = int(state["next_epoch"])

parser = argparse.ArgumentParser()
parser.add_argument("--resume", default=None)
parser.add_argument("--epochs", type=int, default=4)
parser.add_argument("--output", default="./training-states")
args = parser.parse_args()

# 示例只针对单进程CPU；继续模式稍后会覆盖初始化期间的随机状态。
accelerator = Accelerator(cpu=True, mixed_precision="no")
set_seed(31415)
x = torch.arange(256, dtype=torch.float32).reshape(64, 4) / 256
y = x.sum(dim=1, keepdim=True)
loader = DataLoader(TensorDataset(x, y), batch_size=8, shuffle=False, num_workers=0)
model = nn.Sequential(nn.Linear(4, 12), nn.ReLU(), nn.Dropout(0.2), nn.Linear(12, 1))
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
model, optimizer, loader = accelerator.prepare(model, optimizer, loader)

progress = Progress()
accelerator.register_for_checkpointing(progress)
if args.resume:
    # 初始化和prepare必须在恢复前完成；只加载自己生成的可信断点。
    accelerator.load_state(args.resume)
    accelerator.print(f"继续训练，起始epoch={progress.next_epoch}")
else:
    accelerator.print("从头开始训练")

# 加载后不重新设置seed，不插入随机抽样日志。
for epoch in range(progress.next_epoch, args.epochs):
    model.train()
    for inputs, targets in loader:
        optimizer.zero_grad()
        loss = nn.functional.mse_loss(model(inputs), targets)
        accelerator.backward(loss)
        optimizer.step()
    progress.next_epoch = epoch + 1
    checkpoint = Path(args.output) / f"epoch-{progress.next_epoch:03d}"
    if checkpoint.exists():
        # 不覆盖已有断点；改用新的输出根目录继续训练。
        raise FileExistsError(f"断点目录已存在：{checkpoint}")
    accelerator.save_state(str(checkpoint))
    accelerator.print(f"断点保存完成：{checkpoint}")

这是一个独立的最小训练例子，不包含学习率调度器。已有项目若使用调度器，应按项目的prepare或登记方式一起恢复，不能只借用其中的RNG部分。恢复目录不存在、对象不匹配或加载报错时，应让继续操作失败并给出错误，而不是悄悄退回从头训练。

# 先训练到第二个完整epoch，生成独立断点目录。
python train.py --epochs 2 --output ./training-states
# 继续到第四个epoch；总目标epochs不是“再训练四轮”。
python train.py --resume ./training-states/epoch-002 --epochs 4 --output ./continued-states

代码打印的是状态提示，不是精确复现证明。本文没有展示损失曲线、随机输出或性能测量结果。接口以当前官方文档为依据；项目应记录实际安装的Accelerate、PyTorch与NumPy版本，保存和恢复使用相同依赖环境，不把跨版本兼容当成承诺。

自建NumPy生成器要保存自己的状态

如果训练逻辑使用np.random.default_rng，生成器对象维护独立状态。以下封装可加入前面的单进程脚本，在load_state之前创建并登记，随后业务代码通过noise.rng取随机数。保存和恢复时必须使用相同登记集合与顺序。

import copy
import numpy as np

class LocalNoise:
    def __init__(self, seed):
        self.rng = np.random.default_rng(seed)

    def state_dict(self):
        # 显式保存该实例的位生成器状态，避免误用全局随机状态。
        return {"bit_generator": copy.deepcopy(self.rng.bit_generator.state)}

    def load_state_dict(self, state):
        self.rng.bit_generator.state = state["bit_generator"]

# 插在accelerator.load_state之前，保存与恢复两种入口都要登记。
noise = LocalNoise(2718)
accelerator.register_for_checkpointing(noise)

NumPy的Generator文档区分了独立生成器与全局实例，BitGenerator.state提供状态读写入口。此封装示例不是对任意生成器类型的兼容转换。多进程训练中，如果每个rank的自建生成器状态不同，还需要明确逐rank保存设计，不能将这段单进程封装直接视作所有进程状态的完整备份。多进程save_state/load_state也不应仅放入主进程条件分支；各训练进程应按相同控制逻辑参与，保留对应状态。

数据位置和额外开销不能藏在继续按钮后面

完整epoch边界保存省去了中途批次定位，这也是本文示例选它的原因。如果需要epoch中途暂停，还要登记已完成批次数，并依据数据加载器策略恢复读取位置。官方提供skip_first_batches跳过前面批次，但跳过批次不是通用的采样器、worker随机增强和预取队列状态恢复方案。

界面或命令提示中应区分“指定了断点”“完成加载”和“已经开始继续训练”。错误目录、损坏状态或依赖不匹配不能显示成恢复成功。可读状态文字应包含断点标识与业务进度，不只用颜色表达状态，也不要把未知剩余时长显示为准确倒计时。

保存完整状态需要磁盘空间和写入时间。checkpoint频率应根据可接受的重算量与写入成本选择，不能宣称频繁保存没有开销。真正要求连续训练与恢复训练一致时，还要控制硬件、精度、算子确定性、数据顺序和训练拓扑；保存RNG只是必要组成之一，不构成逐位一致保证。

常见问题

rng_types是不是随机状态保存开关？

不是。Accelerator API将它定义为prepared DataLoader迭代时的随机数同步配置；训练断点保存和加载仍使用save_state/load_state。

只调用save_model，恢复Dropout序列就够了吗？

不够。模型导出与训练状态保存目的不同，继续训练还需要优化器、RNG以及业务进度等材料。

恢复后能先运行随机增强预览吗？

如果预览使用训练同一发生器，会消耗随机状态。应与训练发生器分离，或在预览结束后明确恢复原状态，不能忽略这次消耗。

换GPU数量后还算精确恢复吗？

不能直接承诺。进程对应关系、采样分片和计算行为都可能变化，应作为新配置处理，而不是凭断点存在宣称完全一致。

继续训练的核心不是记住最初seed，而是把训练对象、随机发生器状态和业务进度一起放回暂停时的位置。先用完整epoch边界建立可靠的保存入口，再为独立生成器和中途数据位置补充各自的状态设计。