The shuffle change in v5 appears to be intentional, but whether this particular behavior is intended is much more questionable:
@lhoestq
I can reproduce the 8 shards -> shuffle() -> 1 shard -> one effective DataLoader worker behavior independently.
My current read is:
- the new shuffle algorithm itself is intentional: Datasets 5.0 changed
IterableDataset.shuffle() so that the shuffle buffer can be fed from multiple input shards at once;
- but the resulting reduction in the dataset’s logical
num_shards, and therefore in the number of DataLoader workers that can actually receive work, looks much more like a cross-layer regression/design mismatch — or at minimum a behavioral change that is not obvious from the current DataLoader-facing documentation.
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 v5
1. The shuffle change itself is definitely intentional
The Datasets 5.0.0 release notes explicitly describe the new streaming shuffle as a breaking change:
- the shuffle buffer now uses multiple input shards;
- the default is
max_buffer_input_shards=10;
- the old mechanism is available with
max_buffer_input_shards=1.
The implementation came from PR #8194, “Use multiple input shards for shuffle buffer”.
That PR solves a real shuffle-quality problem. With the old implementation, one shard could dominate the buffer for a long time, so data that is correlated or clustered within physical shards may remain poorly mixed even though a sample-level shuffle buffer is being used. The motivation is also connected to #8015.
So I would not characterize “Datasets 5 mixes several shards into the buffer” as the bug. That part appears quite deliberate.
The questionable part is what happens to the sharding level afterwards.
2. Why eight source shards can become one logical shard
With the new algorithm, an IterableDataset with multiple input shards is internally reorganized into multiple child iterables that are consumed together.
Conceptually, with eight source shards and the default maximum of ten:
8 source shards
|
| max_buffer_input_shards = 10
| actual number used here = min(8, 10) = 8
v
8 input streams are interleaved together
|
v
one multi-source logical iterable
The relevant implementation is in src/datasets/iterable_dataset.py.
The important detail is that the multi-source iterable’s sharding level is based on the sharding level of its children. For this simple eight-file case, each of those eight children has one shard, so the resulting logical sharding level becomes one.
That explains the otherwise surprising observation:
print(dataset.num_shards)
# 8
dataset = dataset.shuffle(...)
print(dataset.num_shards)
# 1
In other words, after the transformation, num_shards no longer means “how many physical Parquet files did I originally give this dataset?”. It is describing the current iterable graph’s logical sharding level.
That distinction matters because num_shards is also used elsewhere for execution scheduling.
3. Why that disables the other DataLoader workers
The PyTorch integration also uses that logical num_shards when deciding how many workers can be given independent shard subsets.
If you request:
DataLoader(dataset, num_workers=4)
but the post-shuffle dataset has:
dataset.num_shards == 1
Datasets warns that there are too many workers and stops the unused ones.
In my reproduction, v5 emits the equivalent of:
Too many dataloader workers: 4
(max is dataset.num_shards=1).
Stopping 3 dataloader workers.
and only worker 0 actually yields examples.
This is not merely a cosmetic num_shards display change: it directly changes how many DataLoader processes participate.
That worker/shard boundary has been a tricky part of IterableDataset for some time. Issue #6594, “IterableDataset sharding logic needs improvement”, discusses broader performance and consistency problems around splitting iterable datasets between distributed ranks and DataLoader workers.
There is also #7999, which reports the same general “Too many dataloader workers / num_shards limits workers” symptom in another setup. I would treat that as a related case rather than evidence of the same root cause, because its surrounding pipeline is different.
Why I think the worker collapse is questionable even though the shuffle change is documented
The v5 release notes clearly document this:
old shuffle
->
new multi-input-shard shuffle
What I do not see documented as an intended consequence is:
increase shuffle mixing width
->
reduce logical num_shards
->
reduce DataLoader worker capacity
Those are arguably two different concerns:
- statistical mixing — how many source shards contribute to the shuffle buffer;
- execution parallelism — how many DataLoader processes can independently read portions of the stream.
Right now max_buffer_input_shards influences both.
That coupling is especially noticeable because the current documentation still describes the worker model in terms of the original shuffled shard list.
For example, the current streaming guide shows:
iterable_dataset = dataset.to_iterable_dataset(num_shards=64)
iterable_dataset = iterable_dataset.shuffle(buffer_size=10_000)
dataloader = torch.utils.data.DataLoader(
iterable_dataset,
num_workers=4,
)
and describes this as assigning 64 / 4 = 16 shuffled shards to each worker.
Likewise, the current loading.mdx has an example with 32 generator shards:
shards = [f"data{i}.txt" for i in range(32)]
ds = IterableDataset.from_generator(gen, gen_kwargs={"shards": shards})
ds = ds.shuffle(seed=42, buffer_size=10_000)
dataloader = DataLoader(
ds.with_format("torch"),
num_workers=4,
)
with the comment that each worker gets 32 / 4 = 8 shards.
And the PyTorch integration guide says that for a streaming IterableDataset, each worker is given a subset of the list of shards.
Those descriptions make your expectation quite reasonable.
They may simply need updating for the new v5 logical-shard semantics, or the implementation may need to preserve a separate worker-partitioning dimension after multi-shard shuffle. I do not think the public information currently lets us decide which of those the maintainers intend.
That is why I would phrase this as “reproducible and worth reporting” rather than “definitely an accidental bug”.
Which workaround would I use?
I see three reasonable paths, depending on what you care about.
A. Parquet with useful row groups: reshard() first
This looks like the most natural v5-style path to me:
dataset = load_dataset(
...,
split="train",
streaming=True,
)
print(dataset.num_shards)
dataset = dataset.reshard()
print(dataset.num_shards)
dataset = dataset.shuffle(
seed=42,
buffer_size=...,
)
loader = DataLoader(
dataset,
num_workers=4,
...
)
IterableDataset.reshard() can increase the logical sharding level without requiring more physical files.
This is particularly relevant to Parquet. PR #8193 fixed Parquet resharding so that it can reshard per row group instead of per file, and the current streaming documentation explicitly recommends reshard() when you need more sharding parallelism.
In my test:
8 Parquet files
8 row groups/file
|
| reshard()
v
64 logical shards
|
| default v5 shuffle
v
6 logical shards
|
| DataLoader(num_workers=4)
v
4 active workers
So this can preserve both:
- the new cross-shard shuffle behavior;
- enough logical shards for worker parallelism.
Whether it is faster in your real workload depends on the files, storage backend, row-group sizes, decoding cost, network latency, etc., so I would benchmark it rather than assume more workers automatically means more throughput.
B. Exact pre-v5-style shuffle behavior: max_buffer_input_shards=1
If your priority is simply “make v5 behave like v4 did here”, this is the documented compatibility switch:
dataset = dataset.shuffle(
seed=42,
buffer_size=...,
max_buffer_input_shards=1,
)
For the eight-shard reproduction:
before shuffle: 8
after shuffle: 8
active workers: 4 / 4
The caveat is that this intentionally gives up the new multi-input-shard mixing that v5 introduced.
That may not matter if your Parquet shards are already independently/randomly constructed, but it can matter if records within each shard are strongly correlated.
C. Compromise: use a smaller value instead of 1 or the default 10
For your exact topology, this was interesting:
dataset = dataset.shuffle(
seed=42,
buffer_size=...,
max_buffer_input_shards=2,
)
produced:
8 input shards
->
4 logical shards
->
4 active DataLoader workers
while max_buffer_input_shards=4 produced:
8 input shards
->
2 logical shards
->
2 active workers
So max_buffer_input_shards is effectively acting as both:
- a shuffle-mixing control;
- an indirect worker-parallelism control.
I would treat 2 as an experimentally useful compromise for an 8-shard / 4-worker setup, not as a generally correct formula or an official recommendation.
The exact logical sharding level can depend on the iterable composition, and more complicated interleaved datasets do not have to behave like this simple list-of-Parquet-files case.
A minimal sanity check before changing the pipeline
If you want to tell whether a workaround is actually doing what you want, I would separate four questions instead of looking only at the warning.
1. Did the logical sharding level change?
print("before:", dataset.num_shards)
shuffled = dataset.shuffle(...)
print("after:", shuffled.num_shards)
2. Are all requested DataLoader workers actually producing examples?
num_workers=4 only says PyTorch launched/requested four worker processes. It does not guarantee all four receive useful HF shards.
A simple diagnostic wrapper can tag the worker that yielded each example:
from torch.utils.data import IterableDataset, get_worker_info
class WorkerTaggedDataset(IterableDataset):
def __init__(self, dataset):
self.dataset = dataset
def __iter__(self):
info = get_worker_info()
worker_id = -1 if info is None else info.id
for example in self.dataset:
example = dict(example)
example["_worker_id"] = worker_id
yield example
Then inspect which _worker_id values actually appear.
3. Did coverage remain correct?
If your rows have a stable ID, check at least:
number yielded
number of unique IDs
missing IDs
duplicate IDs
In my synthetic test all configurations yielded:
4096 total
4096 unique
0 duplicates
which was useful because it separated the parallelism problem from a data-loss problem.
That distinction becomes even more important if additional sharding is performed elsewhere in a training stack.
4. Did throughput actually improve?
Do this last.
#8194 also uses threads to fetch initial examples from multiple shards, so:
1 DataLoader worker
does not necessarily mean:
only one I/O operation can ever happen
And conversely, restoring four worker processes does not guarantee four times the throughput.
For real performance comparisons, use your actual remote/local storage path and measure batches/sec or examples/sec after warm-up.
Two things I would avoid
Manually overwriting num_shards
I would not do something like:
dataset.num_shards = 8
or otherwise try to make the metadata report the number of physical files.
The value is not just descriptive metadata. The iterable implementation uses its sharding topology to decide which data sources workers receive.
If the internal graph really has one independently shardable stream but the public value is forced to eight, the scheduler and the iterable would disagree about what can actually be partitioned.
Changing the pipeline (reshard() or shuffle configuration) is much safer than changing the reported number.
Assuming max_buffer_input_shards=1 is a free fix
It restores the old topology, but the v5 shuffle change was made for a reason.
The new algorithm improves mixing when individual source shards contain locally correlated data. Returning to one input shard at a time may reintroduce exactly that issue.
So I would choose between reshard(), 1, and an intermediate value according to your data layout rather than treating the worker warning as the only metric.
Why this may be worth a small upstream issue
If nobody has already filed this exact interaction, I think it would be reasonable to report it upstream with a very small reproducer.
Not because the new shuffle itself is wrong, but because the following chain is surprising:
physical / source sharding
|
v
shuffle mixing topology
|
v
logical num_shards
|
v
maximum useful DataLoader workers
The interesting design question is whether the shuffle mixing width should determine the worker partitioning width.
A useful issue could stay neutral and just show the observed contract change:
Datasets 4.8.5
8 source shards
shuffle()
num_shards = 8
4/4 workers active
Datasets 5.0.x
8 source shards
shuffle() with defaults
num_shards = 1
1/4 workers active
Datasets 5.0.x
max_buffer_input_shards=2
num_shards = 4
4/4 workers active
Datasets 5.0.x
reshard() -> shuffle()
enough logical shards remain
4/4 workers active
and link:
There is also #7999, which is not the same root cause as far as I can tell, but shows another real-world case where the effective num_shards/DataLoader-worker relationship was surprising.
One possible regression-test angle would be to verify not only that shuffled examples are sufficiently mixed, but also how many DataLoader workers actually receive data after the shuffle transform. That would distinguish “shuffle output looks correct” from “the transform preserved the expected worker-level parallelism”.
The broader design point
I think the cleanest way to think about this is to keep four concepts separate:
1. Storage topology
physical Parquet files / row groups
2. Logical partitioning
how many independently assignable pieces the IterableDataset exposes
3. Statistical mixing
shard shuffle / multi-source interleave / sample buffer shuffle
4. Execution parallelism
DataLoader workers / I/O threads / distributed ranks
Datasets has good reasons to change (3): the old shuffle could produce poor mixing for clustered source shards.
Your example exposes a possible unintended coupling between (3) and (4):
increase the number of input shards mixed together
->
reduce logical num_shards
->
reduce the number of DataLoader workers
That is why I would avoid framing this as “the v5 shuffle algorithm is bad”.
The new shuffle solves a real problem.
The more specific question is:
should the internal grouping used to improve shuffle quality also reduce the dataset’s externally useful worker-partitioning level?
That seems like the part worth clarifying.
For now, reshard() is a nice way to separate the concerns for Parquet: create enough logical pieces first, then let the v5 shuffle mix several of them together.
And if you specifically need the old semantics, max_buffer_input_shards=1 is explicitly supported by the v5 release.
So, 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 and #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.