IterableDataset 只有一个 shard 时怎么继续切分数据

Hugging Face Datasets 的 IterableDataset 在流式读取大数据时可以通过 shard 分配工作；但如果 reshard() 之后 num_shards 仍然是 1，就不应该继续把希望寄托在物理分片上。此时可以把数据看成一条有顺序的流，先确定顺序，再用 skip(offset) 和 take(limit) 划出逻辑窗口。

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

单 shard 的处理口诀是：能增加物理 shard 就用分片；增加不了就先 shuffle，再用 skip 定位起点、用 take 限定长度。切片之后不要再改变流的顺序，并为短流和重复遍历留出边界判断。

先判断 shard 还能不能继续增加

这里要区分两种“切分”。物理 shard 是 IterableDataset 底层可分配给读取者的分片数量；逻辑窗口则是从同一条可迭代数据流中，按位置取出一段记录。前者适合让多个读取者各拿一组底层文件，后者适合底层只有一个文件或格式无法继续重分片的情况。

官方文档给出的边界是：如果数据在 reshard() 之后仍然只有一个 shard，就使用 take() 和 skip() 做进一步切分。这个结论并不意味着每个窗口都拥有一个独立的底层文件，它只说明调用方可以把消费范围明确下来。

切片前先固定流的顺序

IterableDataset.shuffle() 不只是改变样本顺序，它还可能改变底层 shard 的访问顺序。为了让窗口划分有清晰语义，应当在 skip 或 take 之前完成 shuffle。官方文档还特别提醒，take 和 skip 会锁定后续的 shard 顺序；完成切片后再调用 shuffle，调用链就不再是同一个可自由重排的流。

from datasets import load_dataset

# streaming=True 返回可迭代数据集，避免入口一次性加载全部样本
dataset = load_dataset(
    "HuggingFaceFW/fineweb",
    split="train",
    streaming=True,
)

# 先决定随机顺序，再创建逻辑窗口；seed 让同一配置可复现
shuffled = dataset.shuffle(seed=42, buffer_size=10_000)

# 如果 reshard 后仍只有一个底层 shard，就保留这条流做 skip/take
print("num_shards:", shuffled.num_shards)

这里的 buffer_size 是流式 shuffle 的采样缓冲区大小，不是 shard 数量，也不是窗口长度。它越大，局部随机化的范围通常越宽，但会增加内存占用；它不能替代后面的 skip 和 take。

用 skip 和 take 划出不重叠窗口

假设每个消费方需要 limit 条记录，第 worker_id 个窗口从 worker_id * limit 开始，就可以把起点和长度显式写出来。窗口的关系是 offset = worker_id * limit，然后先跳过 offset 条，再读取最多 limit 条。

def make_window(dataset, worker_id, window_size):
    # 每个窗口使用固定长度，offset 决定它在逻辑流中的起点
    offset = worker_id * window_size

    # skip 定位起点，take 限制本窗口最多消费多少条记录
    return dataset.skip(offset).take(window_size)


# 先创建窗口，不在窗口创建后再次调用 shuffle
window_a = make_window(shuffled, worker_id=0, window_size=1_000)
window_b = make_window(shuffled, worker_id=1, window_size=1_000)

# 两个窗口的逻辑范围分别是 [0, 1000) 与 [1000, 2000)
for row in window_a:
    consume(row)
for row in window_b:
    consume(row)

这个写法表达的是“逻辑区间不重叠”，不是“底层读取成本已经被均摊”。如果底层只能提供一个 shard，不同窗口可能仍然需要从流的开头扫描到各自的 offset；想减少重复扫描，应优先准备更多底层文件或让数据格式支持更细的 reshard，而不是只改窗口公式。

短流、重复遍历和 worker 分配要单独处理

take(window_size) 的语义是“最多取这么多”，数据不足时不会凭空补齐。最后一个窗口可能短于固定大小，因此下游不能把窗口长度无条件当成真实样本数。需要固定 batch 大小时，应在消费层处理不足一批的情况，而不是把它误判成 shard 切分失败。

还要留意流式数据的可重复遍历成本。为多个 worker 分别构造 skip(offset) 视图，能够清晰表达不重叠范围，但每个视图都可能重新走过前面的记录。若数据源是网络流或压缩文件，重复跳过的代价会比较明显。此时应把窗口分配和数据格式一起设计，必要时先落成多个可独立读取的文件。

如果使用 PyTorch DataLoader，可以优先利用 IterableDataset 的 shard 分配能力；只有在官方文档所说的单 shard 边界下，才把 skip/take 作为逻辑切片工具。不要同时让 DataLoader worker 和业务代码各自按 offset 切一次，否则很容易出现重复跳过或窗口重叠。

一段紧凑的窗口工厂

把顺序固定、窗口参数和消费函数分开，后续更容易替换数据源。下面的工厂不运行检查，也不假设流一定足够长；它只负责把窗口边界写成明确对象。

def window_plan(dataset, count, window_size):
    # 返回逻辑窗口，调用方决定何时真正迭代数据
    for worker_id in range(count):
        offset = worker_id * window_size
        yield {
            "worker_id": worker_id,
            "offset": offset,
            "limit": window_size,
            "dataset": dataset.skip(offset).take(window_size),
        }


# shuffle 应在传入工厂前完成，避免窗口建立后改变顺序
plans = window_plan(shuffled, count=2, window_size=1_000)
for plan in plans:
    print(plan["worker_id"], plan["offset"], plan["limit"])

如果你要的是底层文件级并行，优先检查 num_shards 和 reshard() 的结果；如果你要的是单流上的确定范围，使用 skip/take 并把 offset、limit、短流策略记录在任务配置中。两者解决的问题不同，不要用一个指标替代另一个。

常见问题

为什么 reshard 后 num_shards 还是 1？

这通常意味着当前数据格式或底层文件结构无法继续提供更多独立 shard。此时按照官方边界使用 skip 和 take，或者在数据准备阶段生成更多可独立读取的文件。

skip 和 take 的顺序能交换吗？

如果目标是从原始流的某个位置取固定长度，应先 skip(offset) 再 take(limit)。先 take 再 skip 作用于已经截短的视图，得到的范围就不是同一个窗口定义。

shuffle 应该放在哪里？

放在创建窗口之前。因为 take 和 skip 会锁定后续的 shard 顺序，切片之后再 shuffle 可能触发限制，也会让窗口含义变得不清楚。

归纳起来，单 shard 不是“不能切分”，而是不能继续依赖物理 shard 做切分：先确定流式顺序，再用 offset、skip 和 take 建立逻辑窗口，最后单独评估短流和重复扫描成本。这样既符合 Hugging Face Datasets 的官方边界，也能让窗口分配规则在代码和配置中保持可读。