PyTorch DistributedSampler 每轮 shuffle 的复现配置

PyTorch DistributedSampler 要做到每轮重新打乱、重跑同一轮又能对齐索引，核心配置是：同一进程组使用相同的 sampler seed，各 rank 使用相同 epoch，并在创建该轮 DataLoader 迭代器之前调用 sampler.set_epoch(epoch)。固定 seed 不意味着所有 epoch 都应重复同一种顺序。

官方文档地址：https://docs.pytorch.org/docs/stable/data.html

下面用一个原创的双进程样本索引场景说明配置。本文讨论现有接口的使用调整，不涉及某次新版本发布；代码是供读者复现的示例，没有在本文写作环境中执行，也不提供虚构的训练日志。

每轮应变化的是 epoch，固定的是实验约定

假设训练任务有两个 rank，样本列表与顺序固定。第一次迁移到分布式采样时，只设置 shuffle=True，接着反复迭代同一个 loader，可能发现每轮样本索引没有变化。排查重点应放在 sampler 的 epoch，而不是反复改变全局随机种子。

当前官方实现为 sampler 建立独立的随机数生成器，以 seed 加 epoch 初始化索引排列。set_epoch 更新这个 epoch 值；如果每轮都传零，排列的输入条件就一直相同。这也解释了为什么“使用同一实验种子”和“每轮换一种排列”可以同时成立。

为什么各 rank 应共用 sampler seed

共同的 seed 与 epoch 定义同一份索引排列，rank 和 num_replicas 定义各自取得的部分。若为了让不同进程“更随机”而把 sampler seed 写成 base_seed + rank，各 rank 就会从不同排列中取自己的片段，无法保持原本统一划分的关系。这里说的是 sampler seed，不是在规定模型和数据增强必须采用同一种播种策略。

数据集还需要保持长度稳定、相同索引对应相同样本，所有 rank 的数据顺序也应一致。重跑时只保留同一个 seed，却换了文件清单排列，无法期待索引对应的样本仍然相同。

旧循环只需把 epoch 传到正确位置

下面的函数可接入已经初始化进程组的训练程序。dataset、rank、world_size 由调用方提供，consume_batch 接收每个 batch 并完成调用方的训练处理。它不是独立的 DDP 启动脚本。

from torch.utils.data import DataLoader, DistributedSampler

def consume_epochs(dataset, rank, world_size, start_epoch,
                   end_epoch, consume_batch, base_seed=31415):
    # 所有rank使用相同base_seed；rank仅负责索引划分。
    sampler = DistributedSampler(
        dataset, num_replicas=world_size, rank=rank,
        shuffle=True, seed=base_seed, drop_last=False,
    )
    loader = DataLoader(
        dataset, batch_size=16, sampler=sampler,
        shuffle=False, num_workers=0, drop_last=False,
    )
    for epoch in range(start_epoch, end_epoch):
        # 必须先设置该轮epoch，再让for循环创建loader迭代器。
        sampler.set_epoch(epoch)
        for batch in loader:
            # 异常交给外层训练协调器处理，不在单个rank静默吞错。
            consume_batch(epoch, batch)

DataLoader 已使用 sampler 时，这里明确关闭 loader 自己的 shuffle；索引乱序由 DistributedSampler 负责。不要提前执行 iterator = iter(loader)，随后才调用 set_epoch。即便采样索引暂未取出，这样也会让代码依赖迭代器内部的延迟执行细节，维护时容易出错。

用小样本对照索引，限定复现结论

下面用可被两个 rank 均分的 24 条样本定义一个索引快照函数。显式指定进程数和 rank，因此这个采样示例不依赖真实进程组，也不进行模型训练。只对小数据调用 list(sampler)，不要把大训练集的全部索引复制到日志里。

from torch.utils.data import DistributedSampler

def index_snapshot(epoch, rank):
    # 固定样本清单与共同seed，隔离观察采样顺序。
    sampler = DistributedSampler(
        list(range(24)), num_replicas=2, rank=rank,
        shuffle=True, seed=31415, drop_last=False,
    )
    sampler.set_epoch(epoch)
    return list(sampler)

# 同一环境、同一条件下重建相同rank的同轮索引。
first = index_snapshot(epoch=7, rank=0)
repeat = index_snapshot(epoch=7, rank=0)
assert first == repeat

# 本例样本数可均分，可检查不同rank覆盖互斥且完整。
other_rank = index_snapshot(epoch=7, rank=1)
assert set(first).isdisjoint(other_rank)
assert set(first) | set(other_rank) == set(range(24))

# 比较相邻轮排列；不把有限排列“必不相同”作为通用断言。
next_epoch = index_snapshot(epoch=8, rank=0)
print(first, next_epoch)

当样本数不能被进程数整除时，sampler 的 drop_last=False 会补充索引，不能继续断言各 rank 永不重复；True 则丢弃不能均分的尾部。DataLoader 的同名参数控制每个 rank 最后不足 batch_size 的批次，是另一个层次。

真实实验除采样索引外，还有数据增强随机状态、模型计算与运行环境。PyTorch 官方不承诺跨版本、跨平台或 CPU 与 GPU 之间完全复现。本文的配置目标是对齐既定环境中的索引顺序，不是用一组 sampler 参数保证最终损失完全一致。

恢复训练与常见疑问

从第八轮恢复，要重新从零传 epoch 吗？如果恢复点在完整 epoch 结束后，应恢复真实的下一轮编号，并保留 seed、数据清单与进程规模等记录。把恢复后的循环重新从零计数，会回到较早的采样排列。

只保存 epoch 能从某个 batch 精确继续吗？不能。set_epoch 不保存当前轮已消费的位置；轮内恢复还涉及 batch 游标、预取及随机状态。若只实现轮边界恢复，应在检查点说明这一限制。

shuffle=False 时 set_epoch 会改变顺序吗？它不会开启随机打乱；需要乱序就显式设置 sampler 的 shuffle=True。排查时先分清采样器配置与 loader 配置，避免修改错对象。

进一步阅读：PyTorch 复现性说明；DistributedSampler 官方实现。