Datasets >= 5 ruins sharding when shuffle() is called

Hugging Face's Datasets 5.0.0 changed IterableDataset.shuffle() to fill the shuffle buffer from multiple input shards, which collapses an 8-shard Parquet stream into 1 logical shard and leaves only 1 effective DataLoader worker, down from 4 in v4.8.5, according to a reproduction by a Hugging Face forum user. The user recommends dataset.reshard() or the max_buffer_input_shards=1 parameter to restore pre-v5 behavior, while noting all 4096 rows were returned without duplicates in every tested configuration.

The shuffle change in v5 appears to be intentional, but whether this particular behavior is intended is much more questionable: @lhoestq https://discuss.huggingface.co/u/lhoestq I can reproduce the 8 shards - shuffle - 1 shard - one effective DataLoader worker behavior independently. My current read is: IterableDataset.shuffle so that the shuffle buffer can be fed from multiple input shards at once; num shards For an 8-file Parquet stream, the most practical routes I would try are: Route A: keep the v5 multi-shard shuffle, but create more logical shards first when the Parquet files contain multiple row groups. dataset = dataset.reshard dataset = dataset.shuffle seed=42, buffer size=... or, if preserving the pre-v5 behavior is more important: Route B: documented compatibility path for the old shuffle behavior dataset = dataset.shuffle seed=42, buffer size=..., max buffer input shards=1, There is also a useful middle ground for your specific 8-shard / 4-worker case: dataset = dataset.shuffle seed=42, buffer size=..., max buffer input shards=2, In my small reproduction, that retained 4 logical shards and all 4 DataLoader workers , while still allowing the shuffle buffer to draw from more than one input shard at a time. I would probably try reshard first for Parquet if it gives you enough logical shards, because that preserves the new v5 cross-shard mixing behavior. max buffer input shards=1 is the clean compatibility option if you specifically want the old semantics. A controlled reproduction I ran looked like this: | Datasets version / configuration | num shards after shuffle | DataLoader workers that actually yielded examples | |---|---|---| 4.8.5, default shuffle | 8 | 4 | 5.0.0, default shuffle | 1 | 1 | 5.0.1, default shuffle | 1 | 1 | current main 5.0.2.dev0 when tested , default | 1 | 1 | 5.x, max buffer input shards=1 | 8 | 4 | 5.x, max buffer input shards=2 | 4 | 4 | 5.x, max buffer input shards=4 | 2 | 2 | 5.x, reshard then default shuffle | 64 → 6 | 4 | The test dataset was eight local Parquet files, 512 rows each, with eight row groups per file. Every tested configuration still returned all 4096 unique rows with zero duplicates , so in that small case the thing that changed was the execution/sharding topology, not dataset coverage. I would not read much into throughput numbers from such a small local-file test; active worker count and correctness are the useful observations here. Why this happens in v5So, for your original questions: Why the change from Datasets 4 to 5? To improve streaming shuffle quality by filling the shuffle buffer from multiple input shards instead of effectively processing one input shard at a time. That change is intentional and documented in the 5.0 release https://github.com/huggingface/datasets/releases/tag/5.0.0 and 8194 https://github.com/huggingface/datasets/pull/8194 . Is the 8 → 1 behavior itself a bug? The 8 - 1 logical-shard result follows from the current implementation and is reproducible. What is much less clear is whether losing DataLoader worker parallelism as a consequence was an intended part of that API change. Given the current worker documentation, I would consider this worth an upstream clarification/issue rather than assuming it is expected. Do you need to change how the streaming dataset is instantiated? Probably not fundamentally. For Parquet, I would first try: dataset = dataset.reshard dataset = dataset.shuffle ... If you need exact pre-v5 behavior: dataset = dataset.shuffle ..., max buffer input shards=1, And for eight source shards with four workers, max buffer input shards=2 is also a reasonable low-cost experiment if you want to keep some of the new cross-shard mixing without collapsing below four logical shards.