`IterableDataset` drops samples when resuming from a checkpoint
MEMBER
Thanks for reporting ! I fixed the issue using RebatchedArrowExamplesIterable before the formatted iterable
2025-05-06T14:04:04Z
https://github.com/huggingface/datasets/issues/7538#issuecomment-2854703906
https://api.github.com/repos/huggingface/datasets/issues/7538
IC_kwDODunzps6qJ1Mi
https://api.github.com/repos/huggingface/datasets/issues/comments/2854703906/reactions
2025-05-06T14:04:04Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2854703906
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
When resuming from a checkpoint, `IterableDataset` will drop samples if `num_shards % world_size == 0` and the underlying example supports `iter_arrow` and needs to be formatted. 

In that case, the `FormattedExamplesIterable` fetches a batch of samples from the child iterable's `iter_arrow` and yields them one by one (after formatting). However, the child increments the `shard_example_idx` counter (in its `iter_arrow`) before returning the batch for the whole batch size, which leads to a portion of samples being skipped if the iteration (of the parent iterable) is stopped mid-batch. 

Perhaps one way to avoid this would be by signalling the child iterable which samples (within the chunk) are processed by the parent and which are not, so that it can adjust the `shard_example_idx` counter accordingly. This would also mean the chunk needs to be sliced when resuming, but this is straightforward to implement.

The following is a minimal reproducer of the bug:
```python
from datasets import Dataset
from datasets.distributed import split_dataset_by_node

ds = Dataset.from_dict({"n": list(range(24))})
ds = ds.to_iterable_dataset(num_shards=4)

world_size = 4
rank = 0    
ds_rank = split_dataset_by_node(ds, rank, world_size)

it = iter(ds_rank)
examples = []
for idx, example in enumerate(it):
    examples.append(example)
    if idx == 2:
        state_dict = ds_rank.state_dict()
        break

ds_rank.load_state_dict(state_dict)
it_resumed = iter(ds_rank)
examples_resumed = examples[:]

for example in it:
    examples.append(example)

for example in it_resumed:
    examples_resumed.append(example)

print("ORIGINAL ITER EXAMPLES:", examples)
print("RESUMED ITER EXAMPLES:", examples_resumed)
```