DataLoader shuffle sequence replays from epoch 0 after resuming from a checkpoint

快速结论：使用 Accelerate 的 save_state()/load_state() 恢复训练时，DataLoader 的 shuffle 序列会从 epoch 0 重新开始，导致模型重复训练相同顺序的数据。优先排查 DataLoaderShard.iteration 或 sampler 的 generator 状态是否在 checkpoint 中被保存和恢复。

适用环境：Accelerate 1.13.0，Python 3.12.3，PyTorch 2.11.0+cu128，CUDA 12.8，NVIDIA H100 80GB，Linux（Amazon Linux 2023），使用 FSDP 分布式训练（num_processes=2，mixed_precision=bf16）。

最快修复方案：暂无确认的一步修复方案。社区 Workaround：在 load_state() 后手动调用 train_dataloader.set_epoch(first_epoch)，其中 first_epoch = global_step // steps_per_epoch，将 DataLoader 的 epoch 指针同步到正确的恢复位置。

注意事项：该 Workaround 仅恢复 epoch 计数，不恢复 shuffle 的随机序列本身；建议开启 use_seedable_sampler=True 确保每个 epoch 内的 shuffle 是确定性可复现的。当前 Accelerate 主分支上该 bug 仍未修复，没有合并的 PR。

问题场景

在使用 Hugging Face Accelerate 进行分布式训练时，用户通过 accelerator.save_state() 保存 checkpoint，随后用 accelerator.load_state() 恢复训练。恢复后，DataLoader 产生的 shuffle 顺序与 epoch 0、1、2… 完全相同，而不是从正确 epoch 继续获得新的 shuffle。这意味着模型在重复学习它已经见过的数据顺序。该问题在单进程和多进程（FSDP）场景下均可复现。

报错原文

DataLoader shuffle sequence replays from epoch 0 after resuming from a checkpoint

原因分析

可能原因有两类，取决于 sampler 的配置：

Bug 1（use_seedable_sampler=True）：DataLoaderShard.iteration 在 __init__ 中被初始化为 0，并在每个完整 epoch 结束后递增。每个 epoch 开始时，__iter__ 调用 set_epoch(self.iteration)，进而设置 SeedableRandomSampler 的 epoch，sampler 用 initial_seed + epoch 计算随机种子。但 save_state() 没有保存 iteration，load_state() 也不会恢复它，导致恢复后从 0 重新计数，shuffle 序列从头重放。

Bug 2（use_seedable_sampler=False 默认值 + 多 GPU）：在 prepare_data_loader() 中，当 num_processes > 1 且 sampler 的 generator 为 None 时，Accelerate 会创建私有 torch.Generator 并分配给它。该 generator 用全局 torch RNG 确定性播种。但 checkpoint 代码只为 IterableDatasetShard 保存 sampler，map-style 数据集的私有 generator 不会被保存或恢复。恢复时 prepare() 用相同初始种子重新创建 generator，重复了相同的 shuffle 序列。单进程模式碰巧正常，因为 RandomSampler 会回退到全局 torch RNG，而它会被 load_state 恢复。

环境排查

确认 Accelerate 版本是否为 1.13.0 或更新版本（主分支仍存在该 bug）。
确认 Python 版本（报告中为 3.12.3）。
确认 PyTorch 版本及 CUDA 版本（报告中为 PyTorch 2.11.0+cu128，CUDA 12.8）。
确认 GPU 类型与数量（报告中为 NVIDIA H100 80GB × 2，FSDP）。
检查 DataLoader 的 sampler 配置：是否设置 use_seedable_sampler，数据集是 map-style 还是 IterableDataset。
确认训练脚本是否在恢复后正确恢复了 optimizer、scheduler 和 RNG 状态。

解决步骤

临时修复：在 accelerator.load_state() 之后，手动设置 DataLoader 的 epoch 指针：
```
first_epoch = global_step // steps_per_epoch
train_dataloader.set_epoch(first_epoch)
```
其中 global_step 从 checkpoint 中读取，steps_per_epoch 为每个 epoch 的训练步数。
推荐配置：开启 use_seedable_sampler=True，确保每个 epoch 内的 shuffle 是确定性的，不依赖 generator 的运行时状态。
检查 DataLoader 类型：确认你的数据集是 map-style 还是 IterableDataset。如果是 map-style + 多 GPU，Bug 2 的私有 generator 不会被 checkpoint 保存，这是已知缺陷。
验证复现：用最小复现脚本测试 load_state() 后前几个 batch 的数据顺序是否与原始训练 epoch 0 相同，确认是否命中此 bug。
跟踪上游修复：该 issue 已被 stale-bot 标记，但并未关闭为已修复；社区创建了 computefence 预检工具（pip install computefence）来检测此类静默失败，可以在大规模训练前使用。

验证方法

在恢复训练后，打印前几个 batch 的样本索引或数据内容，与原始训练 epoch 0 的前几个 batch 对比。如果完全一致，说明 shuffle 序列确实被重放了。应用上述临时修复后，再对比确认新序列与原始训练的 epoch 1（或对应恢复位置）一致。更彻底的方法是：记录每个 epoch 第一个 batch 的样本 ID，训练中断前和恢复后应能一一对应。