`use_stateful_dataloader` under multi-process training restores the cursor, not the shuffle order — sampler permutation is never serialized

快速结论：在 Accelerate 多进程（DDP/Gloo，num_processes > 1）训练中启用 use_stateful_dataloader=True 且 DataLoader 使用 shuffle=True 时，保存/恢复 dataloader 状态只还原了“消费到第几批（cursor）”，但没有序列化 shuffle 的排列顺序，导致恢复后位置对、顺序错（use_stateful_dataloader under multi-process training restores the cursor, not the shuffle order — sampler permutation is never serialized）。优先排查 checkpoint 是否保存了 epoch 计数（iteration），以及是否启用了 use_seedable_sampler=True。

适用环境：Accelerate 1.13.0、torch 2.13.0、torchdata 0.11.0、Linux；已在多卡 DDP 及 2 进程 CPU/Gloo 上复现（复现脚本不需要 GPU）。

最快修复方案：暂无确认的一步修复方案。该 Issue 已在 f13f7c13（main）上复现并定位，官方在 PR #4238 中通过持久化 _iteration 并在重建 sampler 迭代器前恢复它来修复；如需精确保序恢复，需要 use_seedable_sampler=True。

注意事项：PR #4238 的修复要到对应版本发布后才生效，Issue 本身未给出可直接运行的命令。精确 shuffle-resume 依赖 use_seedable_sampler=True（排列成为 epoch + seed 的纯函数，可跨重启恢复）；普通 RandomSampler 路径无法在进程重启后做到精确恢复，此时方案改为在保存/加载时给出显式警告。

问题场景

用户在多进程（多卡 DDP，或 2 进程 CPU/Gloo）训练中用 Accelerate 的 Accelerator.prepare() 包装了 DataLoader(num_workers=0, shuffle=True)，并通过 DataLoaderConfiguration(use_stateful_dataloader=True) 启用有状态 dataloader。训练中途（epoch 1 消费 3 个 batch 后）用 loader.state_dict() 按 rank 保存 checkpoint，重启后用 loader.load_state_dict() 恢复。恢复后 dataloader 从正确的位置继续，但该位置上的样本顺序与真实续跑不一致。单进程下没有问题，多进程下才触发。

报错原文

use_stateful_dataloader under multi-process training restores the cursor, not the shuffle order — sampler permutation is never serialized

[rank 0] saved dl state keys: ['_index_sampler_state', '_sampler_iter_state', '_sampler_iter_yielded', '_num_yielded', '_IterableDataset_len_called', '_shared_seed', 'fetcher_state', 'dataset_state', '_iterator_finished']

[rank 0] batches consumed before checkpoint: [[39, 18, 10, 13], [63, 27, 2, 62], [25, 6, 37, 50]]
[rank 0] true epoch-1 remainder            : [[55, 16, 59, 0], [22, 46, 17, 57], [23, 45, 19, 51], [21, 26, 40, 36], [20, 5, 41, 32]]
[rank 0] POSITION restored, ORDER lost: epoch-0's order at epoch-1's position

AssertionError: not even the deterministic epoch-0 replay
AssertionError: resumed matched the true continuation -- bug not reproduced

原因分析

已确认的根因（issue 复现于 f13f7c13，官方开出 PR #4238）：分片的有状态路径用 DataLoaderShard.iteration（epoch 计数器，经 set_epoch 传递）作为 shuffle 的种子，而这个计数器没有被写入 checkpoint。因此新进程恢复时 iteration=0，torchdata 的 fast-forward 就按 epoch 0 的排列回放到保存时的 cursor，于是出现“位置对、顺序错”——即用 epoch 0 的顺序填到了 epoch 1 的位置。保存的状态 key 列表里也确实没有 _iteration。

PR #4238 的做法是把 _iteration 持久化，并在重建 sampler 迭代器之前恢复它。同时注意 Issue 中的“either/or”：要做到跨重启精确保序恢复，需要 use_seedable_sampler=True，此时排列是 epoch + seed 的纯函数，可以重建；普通 RandomSampler 路径无法做到精确恢复，因此 PR 在该路径下退化为给出显式警告。

环境排查

Accelerate 版本：Issue 使用 1.13.0；确认是否已包含 PR #4238 的修复。
PyTorch 版本：2.13.0。
torchdata 版本：0.11.0。
操作系统：Linux。
进程配置：是否 num_processes > 1（多卡 DDP 或 CPU/Gloo 多进程）。
DataLoader 配置：是否 shuffle=True、num_workers 取值、是否设置 use_seedable_sampler=True。
checkpoint 内容：state_dict() 的 key 中是否包含 _iteration（旧行为下没有）。
epoch 计数：保存前是否调用了 set_epoch()，其值是否随 checkpoint 一起落盘。

解决步骤

用 Issue 提供的两段脚本在本机复现：先 save 再 resume，观察 resumed 是否等于 ref['cont']。该复现只需 2 个 CPU 进程，无需 GPU（在 GPU 机器上以 CUDA_VISIBLE_DEVICES="" 运行）。
打印 loader.state_dict().keys()，确认列表里是否存在 _iteration。若不存在，说明当前版本仍在用 iteration 作为 shuffle 种子且未持久化，与本 Issue 描述一致。
确认自己的 DataLoader 配置里是否设置了 use_seedable_sampler=True。只有该选项开启时，排列才可能作为 epoch + seed 的纯函数跨重启还原。
可优先尝试：升级/切换到包含 PR #4238 的 Accelerate 版本（该 PR 持久化 _iteration 并在重建 sampler 迭代器前恢复）。升级前请自行核对目标版本的 changelog，Issue 未直接给出具体版本号。
若无法升级，且必须依赖多进程 shuffle 的精确恢复，可优先尝试在训练脚本中对 sampler 设置固定 seed 与可持久化的 epoch，并在保存 checkpoint 时把当前 epoch 一并落盘，恢复时先恢复 epoch 再重建 dataloader。注意此路径依赖 use_seedable_sampler=True 才能做到精确恢复，普通 RandomSampler 下仍无法保证。
如果没有精确保序恢复的硬需求，至少应在保存/加载处记录或提示“多进程下仅恢复 cursor、不恢复 shuffle 顺序”，避免误以为续跑与未中断一致。

验证方法

重新运行 Issue 的复现脚本（save 后 resume），检查断言 resumed == ref['cont'] 是否通过，即恢复后的 batch 流与保存时记录的“真实 epoch 1 剩余部分”完全一致。修复生效后，state_dict() 的 key 中应出现 _iteration，且不再触发 POSITION restored, ORDER lost 的报错路径。若只开启了普通 RandomSampler 路径，则应看到明确的警告输出，而不是静默返回错误顺序。