Feature request: Add a streaming_shard operator for early sample-level sharding when file-level sharding is insufficient

[muyihao]

Feature request: Add a streaming_shard operator for early sample-level sharding when file-level sharding is insufficient

Proposed API

It would be useful to provide a first-class streaming sharding operator in datasets, for example:

dataset = datasets.load_dataset(..., streaming=True, split="train")

dataset = dataset.streaming_shard(
    num_shards=world_size,
    index=rank,
)

or:

dataset = datasets.streaming_shard(
    dataset,
    num_shards=world_size,
    index=rank,
)

Then users could write:

dataset = datasets.load_dataset(..., streaming=True, split="train")

dataset = dataset.streaming_shard(
    num_shards=world_size,
    index=rank,
)

dataset = dataset.map(tokenize_fn)
dataset = dataset.map(build_labels_or_loss_mask)
dataset = dataset.filter(...)

This would move rank-aware sample sharding earlier in the streaming pipeline, before expensive transformations.

Why this is needed

Existing file-level sharding works well when the dataset has enough physical shards. However, it is not enough when:

num_physical_files < global_world_size

or when physical files cannot be split.

For example, if we have:

global_world_size = 64
num_physical_files = 4

then file-level assignment alone cannot give each rank a distinct physical shard. Many ranks would either read the same files or receive no file-level shard. In practice, the only reliable way to divide the dataset is sample-level sharding, such as what IterableDatasetShard does.

The problem is that IterableDatasetShard is applied at the training/dataloader layer. This means users cannot easily place the shard boundary before expensive dataset transformations.

Expected behavior

For a streaming dataset yielding:

[0, 1, 2, 3, 4, 5, 6, 7]

with:

num_shards = 2

the result could be:

rank 0 -> [0, 2, 4, 6]
rank 1 -> [1, 3, 5, 7]

This is sample-level logical sharding and does not depend on the number of physical files.

A batch-aware mode may also be useful to align with IterableDatasetShard semantics:

dataset.streaming_shard(
    num_shards=world_size,
    index=rank,
    batch_size=batch_size,
    split_batches=False,
    drop_last=False,
)

Current workaround

The current workaround is to rely on IterableDatasetShard:

dataset = datasets.load_dataset(..., streaming=True, split="train")

dataset = dataset.map(tokenize_fn)
dataset = dataset.map(build_labels_or_loss_mask)

dataset = IterableDatasetShard(
    dataset,
    batch_size=batch_size,
    num_processes=world_size,
    process_index=rank,
)

This is correct from a data partitioning perspective, but inefficient because the sharding happens after the previous dataset transformations.

Desired pattern

The desired pattern is:

dataset = datasets.load_dataset(..., streaming=True, split="train")

dataset = dataset.streaming_shard(
    num_shards=world_size,
    index=rank,
)

dataset = dataset.map(tokenize_fn)
dataset = dataset.map(build_labels_or_loss_mask)

This way, each rank only preprocesses the examples it will actually consume.

Use cases

This would be useful for:

1. Large-scale LLM pretraining with a small number of very large JSONL / Parquet files.
2. SFT pipelines where tokenization, chat template application, and loss-mask construction are expensive.
3. Multimodal training where image/video metadata parsing or remote object loading should not be repeated across ranks.
4. Distributed jobs where the number of source files is smaller than global_world_size.
5. Data sources where physical repartitioning is expensive or controlled by an upstream production pipeline.

Relationship to existing APIs

This feature is related to:

• accelerate.data_loader.IterableDatasetShard, which provides sample-level sharding at the dataloader/training layer.
• Datasets file-level sharding, which works well when there are enough physical source shards.
• to_iterable_dataset(num_shards=...), which can create iterable shards when converting a map-style dataset.

The proposed feature is different because it provides logical sample-level sharding directly in the datasets streaming pipeline, before downstream transformations such as map, filter, tokenization, or loss-mask construction.

Open questions

1. Should the operator exactly match IterableDatasetShard semantics?
2. Should it support both sample-level and batch-aware sharding?
3. Should it support both strided and contiguous logical sharding?
4. How should it interact with shuffle(buffer_size=...) and set_epoch()?
5. Should this be implemented as a method on IterableDataset, or as a utility function under datasets.distributed?
6. Should it expose options such as drop_last, even_batches, and split_batches to match distributed dataloader behavior?

Summary

When physical file-level sharding is impossible or when the number of source files is smaller than the global world size, distributed streaming training has to rely on IterableDatasetShard for sample-level data partitioning. However, this sharding happens too late, after expensive dataset transformations may already have been executed.

A streaming_shard operator in datasets would allow users to apply logical rank-aware sharding earlier in the streaming pipeline, reducing duplicated preprocessing work and improving efficiency for large-scale distributed training.

Motivation

When using datasets.load_dataset(..., streaming=True) for distributed training, Hugging Face Datasets can benefit from physical file-level sharding if the dataset is already split into enough source files.

However, in many real-world training pipelines, this assumption does not hold.

For example:

• the dataset may contain only a few very large JSONL / Parquet files;
• the number of physical files may be smaller than the global distributed world size;
• the source files may be generated by an upstream pipeline and cannot be repartitioned easily;
• the files may be stored on remote storage such as HDFS or object storage, where rewriting or repartitioning is expensive;
• some file formats or data sources do not provide a convenient way to split one physical file into multiple independent streaming shards.

In these cases, file-level sharding is not sufficient. The training job has to fully rely on accelerate.data_loader.IterableDatasetShard or similar dataloader-level logic to split the data across ranks.

This works functionally, but it happens too late in the pipeline.

A common pipeline looks like this:

dataset = datasets.load_dataset(..., streaming=True, split="train")

dataset = dataset.map(tokenize_fn)
dataset = dataset.map(build_labels_or_loss_mask)
dataset = dataset.filter(...)

Then the distributed dataloader wraps the dataset with IterableDatasetShard.

The issue is that every rank may execute the upstream streaming pipeline before the dataloader-level shard decides which examples should be kept by the current rank. For expensive preprocessing steps such as tokenization, chat template application, multimodal metadata loading, filtering, or loss-mask construction, this may cause duplicated CPU work and unnecessary I/O pressure.

Your contribution

I propose adding a streaming_shard operator that can be applied before expensive streaming dataset transformations.

I’m happy to work on this feature if the maintainers think it makes sense.

[lhoestq]

I feel like IterableDatasetShard is not well suited for this and shouldn't be used / should be improved. Instead there is ds.shard() or even split_dataset_by_node() which works when num_physical_files < global_world_size. WDYT ?

[beanscg]

I checked current main, and the requested behavior seems to sit between the two existing streaming split mechanisms:

• src/datasets/iterable_dataset.py:3928-3973 implements IterableDataset.shard(...) by calling self._ex_iterable.shard_data_sources(...), so it primarily operates at the data-source/shard level.
• src/datasets/distributed.py:10-43 routes iterable datasets to _split_by_node_iterable_dataset(...).
• src/datasets/iterable_dataset.py:5341-5348 documents that _split_by_node_iterable_dataset(...) assigns shards evenly when possible, otherwise keeps 1 example out of world_size.
• I did not see an open PR that appears to cover issue Feature request: Add a streaming_shard operator for early sample-level sharding when file-level sharding is insufficient #8253.

That suggests a decision point: either improve the existing fallback sample-skip path into a public early sample-level operator, or document that split_dataset_by_node(...) is the intended sample-level fallback when file shards are fewer than workers. A useful acceptance test would cover num_physical_files < world_size and assert no worker receives an empty stream.

[muyihao]

@lhoestq @beanscg Thanks for the response. I checked the implementation and found that split_dataset_by_node() already falls back to StepExamplesIterable when num_shards is not divisible by world_size.

However, I think there are still two gaps for this use case:

1. The sample-level split is only used as a fallback. If num_shards % world_size == 0, the implementation always prefers data-source-level sharding, and users cannot explicitly request sample-level sharding. This can be problematic when physical files are highly imbalanced in size or number of examples, where a more balanced sample-level split would be preferred even though the number of shards is divisible by world_size.

2. The fallback seems to address distributed rank/node-level splitting, while DataLoader workers are still split through shard_data_sources(). Therefore, whether multiple DataLoader workers can be effectively used still depends on having enough physical shards/data sources. When the number of physical files is smaller than num_workers, some workers may receive no shard and stop. If sample-level splitting were also available at the worker level, then even a single streaming file could be processed concurrently by multiple DataLoader workers.

So maybe the missing piece is not only the fallback itself, but an explicit/public way to request sample-level sharding, with clear semantics for both distributed ranks and DataLoader workers.

[muyihao]

Should split_dataset_by_node support explicitly choosing sample-level sharding when data-source-level sharding is not suitable? WDYT?

[lhoestq]

yes adding an argument for explicit splitting strategy makes sense imo

[muyihao]

@lhoestq I think a streaming_shard semantic is still useful.

split_dataset_by_node mainly addresses distributed rank-level assignment and works well when the input dataset has enough physical shards. DataLoader workers are another level of parallelism and also require enough shards to be effective.

However, another case is a streaming dataset backed by a single very large file. Since it only has one physical shard, increasing num_workers may not help much. Resharding beforehand is possible, but it adds extra preprocessing, may not support all formats, and can still be imbalanced at the file level.

Would a streaming_shard-like API be useful for this case, by doing row-level/sample-level sharding during streaming and before expensive transforms like tokenization or loss-mask construction. WDYT?