Datasets shard 怎样切分数据集做并行处理

用 Hugging Face Datasets 做并行预处理时，shard() 适合把一个数据集按 worker 数量拆成互不重叠的子集。最小写法是：

from datasets import load_dataset

# 先加载数据，再让第 1 个 worker 领取 4 份中的第 0 份
dataset = load_dataset("rotten_tomatoes", split="train")
worker_dataset = dataset.shard(num_shards=4, index=0)
print(len(worker_dataset))

官方地址：https://huggingface.co/docs/datasets/

这里的 num_shards 是总分片数，index 是当前 worker 要领取的分片编号，从 0 开始。真正容易出错的地方不在调用形式，而在于：你想要连续的数据区间，还是想让每个 worker 交错拿到记录。

先看清楚并行任务到底需要哪种分片

假设原始数据按索引排列为 0, 1, 2, 3, 4, 5, 6, 7, 8，调用 shard(num_shards=3, index=i) 时，默认的 contiguous=True 会尽量切成连续区间。不能整除的余数会优先分给前面的分片，因此结果可能是：

• index=0：0, 1, 2
• index=1：3, 4, 5
• index=2：6, 7, 8

如果设置 contiguous=False，则按“记录索引除以分片数的余数”选择记录：

• index=0：0, 3, 6
• index=1：1, 4, 7
• index=2：2, 5, 8

连续分片更适合按文件区间、时间区间或原始顺序处理；交错分片则能让每个 worker 更均匀地看到整个数据集的索引分布。shard() 只选择数据子集，不会自动创建进程，也不会替你启动并行框架。

用 num_shards 和 index 给每个 worker 分配任务

在多进程或多机任务里，通常把总 worker 数作为 num_shards，把当前 worker 的编号作为 index。同一批任务必须使用同一个总数，并保证每个 worker 的编号唯一：

from datasets import load_dataset

dataset = load_dataset("rotten_tomatoes", split="train")

num_workers = 3
worker_id = 1

# 每个 worker 只领取自己的分片；编号从 0 到 num_workers - 1
part = dataset.shard(
    num_shards=num_workers,
    index=worker_id,
    contiguous=True,
)

print({
    "worker_id": worker_id,
    "rows": len(part),
})

当记录数不能被 num_shards 整除时，前面的分片多一条，后面的分片少一条。不要用“每个分片一定相同长度”去写调度逻辑；如果业务要求严格等量，需要先处理余数或接受少量不均衡。

把 shard 放到 shuffle 前面，边界更容易复查

官方文档建议尽早调用 shard，尤其不要先对完整数据集随机化，再把随机后的结果分给 worker。推荐的结构是“先确定边界，再在分片内部随机化”：

from datasets import load_dataset

dataset = load_dataset("rotten_tomatoes", split="train")
num_workers = 3
worker_id = 1

# 先切出稳定边界，避免 worker 之间共享同一个随机化结果
part = dataset.shard(num_shards=num_workers, index=worker_id)

# 再在当前 worker 的分片内打乱；seed 让本 worker 的顺序可复现
part = part.shuffle(seed=2026)

for row in part.select(range(min(3, len(part)))):
    print(row["text"])

这不意味着任何随机化都必须禁止，而是要先决定“谁负责哪一段数据”，再决定“这一段内部怎样取样”。这样排查重复记录、遗漏记录和复现差异时，边界更清楚。

遇到空分片或越界时，先检查这三个条件

如果程序报错、某个 worker 没有记录，或不同 worker 的结果看起来不完整，可以按下面顺序检查。

1. 检查 index 范围。index 必须是非负整数，并且小于 num_shards。把 worker 编号从 1 开始传入时，最后一个编号很容易越界。
2. 检查分片数量是否过大。普通 Dataset 不应把 num_shards 设得大于数据集记录数；流式 IterableDataset 还要考虑底层数据本身可用的 shard 数量。
3. 检查分片对象类型。普通数据集可以用 len() 检查记录数，流式数据集不应假设能直接得到完整长度，应按迭代消费和底层 shard 数来安排任务。

def get_worker_part(dataset, num_workers, worker_id):
    # 入口处统一检查编号，避免把 1-based worker id 直接当成 index
    if num_workers = num_workers:
        raise ValueError("worker_id 必须位于 [0, num_workers) 内")

    # contiguous=True 保持原始顺序中的连续区间
    return dataset.shard(
        num_shards=num_workers,
        index=worker_id,
        contiguous=True,
    )

不要为了让每个 worker “都有数据”而盲目增大底层数据集或重复重试。先确认总任务数是否合理，再决定降低 worker 数、换成交错分片，或让空任务正常退出。

怎样确认切分结果符合预期

对普通 Dataset，可以用三个分片重新拼接，观察记录顺序是否恢复。连续分片最适合这种检查：

from datasets import concatenate_datasets

parts = [
    dataset.shard(num_shards=3, index=i, contiguous=True)
    for i in range(3)
]

# 连续分片按 index 顺序拼接，可恢复原数据的记录顺序
recombined = concatenate_datasets(parts)
print(len(dataset), len(recombined))
print(dataset[0] == recombined[0])

这段代码只是在内存中检查分片关系，不代表实际并行任务已经完成。生产任务还应记录每个 worker 的 num_shards、index、样本计数和随机种子，后续出现重复或缺失时才有可追踪的线索。

几个容易混淆的问题

shard 会复制数据吗？

它返回的是选定分片的数据集视图或数据集对象，不应把它理解成把整份原始数据复制到每个 worker。实际缓存和索引行为取决于数据集类型与后续操作。

contiguous=False 能保证每个分片长度完全一样吗？

不能。记录数不能整除时，余数仍然会让分片长度相差一条左右；它改变的是选取模式，不是消除余数。

为什么不直接用 select 做并行切分？

select 适合明确给出索引集合的场景；shard 直接表达“按总分片数和当前编号领取一份”，更容易和 worker 调度参数保持一致。

总结

Datasets 的 shard 可以概括成三句话：num_shards 决定总份数，index 决定当前 worker 领取哪一份，contiguous 决定按连续区间还是交错索引挑选。并行处理时先固定分片边界，再在分片内部 shuffle，最后记录每个 worker 的边界和计数，排查重复、遗漏与顺序问题会简单很多。