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

## 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:

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

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

or:

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

Then users could write:

```python
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:

```text
num_physical_files < global_world_size
```

or when physical files cannot be split.

For example, if we have:

```text
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:

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

with:

```python
num_shards = 2
```

the result could be:

```python
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:

```python
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`:

```python
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:

```python
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:

```python
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.

Labels: name: enhancement

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 ?

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 #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.


@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.


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

yes adding an argument for explicit splitting strategy makes sense imo

@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?

Subject: Proposal for Efficient Early-Stage Streaming Sharding via "Logical Index-based Filtering"


Hi @lhoestq and @huggingface/datasets team, 
I’ve been closely following the discussion in #8253 and #8268 regarding distributed streaming bottlenecks. As a CS student (BCA 4th semester) deeply interested in distributed systems and AI architecture, I’ve been analyzing the overhead caused by late-stage sharding.

I would like to propose a refined implementation strategy—an "Early-Stage Logical Sharding Gate"—that balances performance with the existing library architecture.

The Core Problem
In the current streaming pipeline, sharding via IterableDatasetShard occurs after heavy transformations (tokenization, mapping). For large-scale training where num_physical_files < world_size, this results in redundant CPU computation and I/O pressure, as all workers process data only to discard it later.

Proposed Solution: The "Gatekeeper" Operator
Instead of just a fallback mechanism, I propose a first-class StreamingShardOperator that intercepts the stream before expensive operations:

Logical Indexing: Maintain a deterministic global_index counter within the stream.

Early Gatekeeping: Filter samples using global_index % world_size == rank.

Zero-Copy Vectorization: For batched input (pyarrow.Table), use Table.filter() with pre-computed masks. This leverages PyArrow's C++ backend, ensuring high-throughput and "Zero-Copy" slicing.

State-Dict Awareness: By storing global_index in the state_dict, we ensure full determinism across restarts and epochs.

Why this approach?
Performance: It reduces the transformation workload by (world_size - 1) / world_size.

Non-intrusive: By implementing this as a modular operator (e.g., in a new streaming_shard.py), we keep the core IterableDataset clean and avoid breaking changes.

Alignment: This is compatible with PyTorch DistributedSampler logic, providing a familiar API for distributed training users.

Next Steps
I am currently working on a Proof-of-Concept (PoC) to benchmark this operator against the existing fallback mechanisms. I believe this implementation would significantly benefit the community, especially for those training LLMs on limited file structures.

I am eager to contribute this as a PR. Could you please let me know if this architectural direction aligns with your roadmap?



khankamil67130 (CS Student, BCA 4th Semester)