`datasets` Iterable Dataset sharding support

问题场景

用户在使用 HuggingFace datasets 通过 load_dataset(..., streaming=True) 创建 IterableDataset，并希望借助其自带的 shard、shuffle、state_dict 在多节点训练中让各节点读取不同数据、便于随机化和恢复。但当把该 dataloader 交给 Accelerate 的 prepare 后，Accelerate 默认假设不同 rank 上的数据集实例相同，并由 IterableDatasetShard 做超额采样后再分配子批次，导致与数据集自带分片逻辑冲突。

报错原文

datasets Iterable Dataset sharding support

该 Issue 本身没有贴出 Python 异常堆栈；评论指出更典型的失效方式是静默数据丢量：

No error and no hang - the epoch just covers a fraction of the data, and the loss curve looks normal.

原因分析

可能原因在于 prepare_data_loader 对 IterableDataset 有两种处理路径：

• 当是 datasets.IterableDataset 且 n_shards >= num_processes 且 split_batches=False 时，走 .shard() 路径，不会触发额外的超额采样；
• 其他情况——分片数少于进程数、split_batches=True，或普通 torch.utils.data.IterableDataset——会落到 IterableDatasetShard，每个进程都遍历完整数据集后只保留 1 / num_processes 的元素。

如果数据集本身已按 source 分片（各进程各自打开自己的文件），再用 prepare 会被二次分片。评论中在 main 上实测：8 个 shard、每 shard 10 条记录、num_processes=4、batch_size=2 时，自分片数据集再 prepare 后每进程仅 6 条、全局唯一记录只有 24/80，而 shard 读取次数仍为 8。

环境排查

• 确认 accelerate 版本，以及是否启用了 split_batches。
• 确认 num_processes 与数据集实际分片数 n_shards 的关系：n_shards 是否大于等于进程数。
• 确认数据集类型是 datasets.IterableDataset 还是普通 torch.utils.data.IterableDataset，这决定是否走 .shard() 路径。
• 确认数据集是否已经按 source 自行分片（每个进程打开自己的文件），否则容易发生二次分片。
• 确认 prepare() 后 dataloader 的 len() 是否异常下降，可作为数据被截断的排查信号。
• Issue 未记录 Python、CUDA、PyTorch、显卡型号，暂无法据此做版本级排查。

解决步骤

1. 先定位当前路径：检查是否为 datasets.IterableDataset，并核对 n_shards >= num_processes 与 split_batches=False 是否同时满足；不满足时会进入 IterableDatasetShard。
2. 若数据集已自带分片，先尝试不 prepare 该 dataloader（Issue 正文提出的临时方案，可优先尝试），并对每个进程自行调用 .shard() 获取对应数据。
3. 可优先尝试评论中提到的规避方式：使用 .shard()、开启 dispatch_batches=True、或按 source 分片并跳过 prepare；注意跳过 prepare 时需自行接管原本由 prepare 提供的能力（如设备放置、批次分发等）。
4. 不要依赖“训练无报错”作为通过标准；对照上一步的规避方案，重新核对每进程样本数与全局唯一样本数。
5. 若希望从框架侧解决，Issue 讨论方向是新增 DataLoaderConfiguration 开关，或让 prepare_data_loader 尊重数据集上的属性以跳过已分片数据集；评论中尚未确认最终采用哪种形式，且“Completed already !”未附带具体合并 PR 或版本号，暂不能作为已验证的升级方案。

验证方法

统计各进程实际消费的记录数、全局唯一记录数，以及底层 shard 的读取次数。以评论中的基准为例：8 个 shard、每 shard 10 条、num_processes=4、batch_size=2 时，正确结果应接近每进程 20 条、全局 80/80、shard 读取 8 次；若出现每进程仅 6 条、全局 24/80，说明仍在二次分片。同时检查 prepare() 后 dataloader 的 len() 是否异常下降。

参考来源

huggingface/accelerate #3547

huggingface/accelerate #2859（Issue 正文提及的相关 Issue，目标是给 dataloader 增加 load_state_dict）

评论中提及的文档 PR：#4144（仅记录当前行为与规避方式，不改变行为）。

AI 工具推荐

想把多个 AI 模型放在一个入口？

GamsGo AI 集成 ChatGPT、DeepSeek、Gemini、Claude、Midjourney、Veo 等常用模型，适合写作、绘图、视频和日常 AI 工作流。

了解 GamsGo AI

推广链接：通过此链接购买，我可能获得佣金，不影响你的价格。