Epochs shortened after resuming mid-epoch with Iterable dataset+StatefulDataloader(persistent_workers=True)
MEMBER
Thanks for reporting ! Maybe we should store the epoch in the state_dict, and then when the dataset is iterated on again after setting a new epoch it should restart from scratch instead of resuming ? wdyt ?
2025-03-13T11:21:08Z
https://github.com/huggingface/datasets/issues/7447#issuecomment-2720906619
https://api.github.com/repos/huggingface/datasets/issues/7447
IC_kwDODunzps6iLb17
https://api.github.com/repos/huggingface/datasets/issues/comments/2720906619/reactions
2025-03-13T11:21:08Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2720906619
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
NONE
But why does this only happen when `persistent_workers=True`? I would expect it to work correctly even without storing the epoch number in the state_dict of the iterable dataset. 
2025-03-13T22:22:49Z
https://github.com/huggingface/datasets/issues/7447#issuecomment-2722813794
https://api.github.com/repos/huggingface/datasets/issues/7447
IC_kwDODunzps6iStdi
https://api.github.com/repos/huggingface/datasets/issues/comments/2722813794/reactions
2025-03-13T22:22:49Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2722813794
https://avatars.githubusercontent.com/u/4356534?v=4
https://api.github.com/users/dhruvdcoder/events{/privacy}
https://api.github.com/users/dhruvdcoder/followers
https://api.github.com/users/dhruvdcoder/following{/other_user}
https://api.github.com/users/dhruvdcoder/gists{/gist_id}
https://github.com/dhruvdcoder
dhruvdcoder
MDQ6VXNlcjQzNTY1MzQ=
https://api.github.com/users/dhruvdcoder/orgs
https://api.github.com/users/dhruvdcoder/received_events
https://api.github.com/users/dhruvdcoder/repos
https://api.github.com/users/dhruvdcoder/starred{/owner}{/repo}
https://api.github.com/users/dhruvdcoder/subscriptions
User
https://api.github.com/users/dhruvdcoder
public
MEMBER
I think persistent_workers=False simply ignores the dataset state_dict when it starts a new epoch, that's why the issue doesn't appear in that case
2025-03-14T09:58:40Z
https://github.com/huggingface/datasets/issues/7447#issuecomment-2724207156
https://api.github.com/repos/huggingface/datasets/issues/7447
IC_kwDODunzps6iYBo0
https://api.github.com/repos/huggingface/datasets/issues/comments/2724207156/reactions
2025-03-14T09:58:40Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2724207156
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
MEMBER
I opened https://github.com/huggingface/datasets/pull/7451 to fix the issue, let me know if it works for you
2025-03-14T10:32:07Z
https://github.com/huggingface/datasets/issues/7447#issuecomment-2724274405
https://api.github.com/repos/huggingface/datasets/issues/7447
IC_kwDODunzps6iYSDl
https://api.github.com/repos/huggingface/datasets/issues/comments/2724274405/reactions
2025-03-14T10:32:07Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2724274405
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
MEMBER
I just released `datasets` 3.4 that includes the fix :)

PS: in your script you probably want to set the epoch like this, otherwise it's still set to 0 after the first epoch:

```diff
    if state_dict is None:
-       ds.set_epoch(epoch)
        epoch += 1
+       ds.set_epoch(epoch)
```
2025-03-14T17:25:44Z
https://github.com/huggingface/datasets/issues/7447#issuecomment-2725326404
https://api.github.com/repos/huggingface/datasets/issues/7447
IC_kwDODunzps6icS5E
https://api.github.com/repos/huggingface/datasets/issues/comments/2725326404/reactions
2025-03-14T17:26:58Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2725326404
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
NONE
@lhoestq 
If I understand correctly, the issue was:
when training saves a checkpoint of dataloader in epoch 1, the resumed training only consumes partial data in epoch 2, 3, etc.

However, with the fix we are facing the issue that:
when training saves a checkpoint of dataloader in epoch 2, the resumed training starts from scratch instead of consuming remaining partial data in epoch 2.

This makes training inconsistent between resuming from a checkpoint vs. original training if continued without a checkpoint.
2025-07-09T23:04:57Z
https://github.com/huggingface/datasets/issues/7447#issuecomment-3054375914
https://api.github.com/repos/huggingface/datasets/issues/7447
IC_kwDODunzps62DhPq
https://api.github.com/repos/huggingface/datasets/issues/comments/3054375914/reactions
2025-07-09T23:04:57Z
https://api.github.com/repos/huggingface/datasets/issues/comments/3054375914
https://avatars.githubusercontent.com/u/150487191?v=4
https://api.github.com/users/tianyu-l/events{/privacy}
https://api.github.com/users/tianyu-l/followers
https://api.github.com/users/tianyu-l/following{/other_user}
https://api.github.com/users/tianyu-l/gists{/gist_id}
https://github.com/tianyu-l
tianyu-l
U_kgDOCPhAlw
https://api.github.com/users/tianyu-l/orgs
https://api.github.com/users/tianyu-l/received_events
https://api.github.com/users/tianyu-l/repos
https://api.github.com/users/tianyu-l/starred{/owner}{/repo}
https://api.github.com/users/tianyu-l/subscriptions
User
https://api.github.com/users/tianyu-l
public
### Describe the bug

When `torchdata.stateful_dataloader.StatefulDataloader(persistent_workers=True)` the epochs after resuming only iterate through the examples that were left in the epoch when the training was interrupted. For example, in the script below training is interrupted on step 124 (epoch 1) when 3 batches are left. Then after resuming, the rest of epochs (2 and 3) only iterate through these 3 batches.

### Steps to reproduce the bug

Run the following script with and with PERSISTENT_WORKERS=true.

```python
# !/usr/bin/env python3
# torch==2.5.1
# datasets==3.3.2
# torchdata>=0.9.0
import datasets
import pprint
from torchdata.stateful_dataloader import StatefulDataLoader

import os

PERSISTENT_WORKERS = (
    os.environ.get("PERSISTENT_WORKERS", "False").lower() == "true"
)

# PERSISTENT_WORKERS = True  # Incorrect resume


# ds = datasets.load_from_disk("dataset").to_iterable_dataset(num_shards=4)
def generator():
    for i in range(128):
        yield {"x": i}


ds = datasets.Dataset.from_generator(
    generator, features=datasets.Features({"x": datasets.Value("int32")})
).to_iterable_dataset(num_shards=4)

dl = StatefulDataLoader(
    ds, batch_size=2, num_workers=2, persistent_workers=PERSISTENT_WORKERS
)
global_step = 0
epoch = 0
ds_state_dict = None
state_dict = None
resumed = False
while True:
    if epoch >= 3:
        break
    if state_dict is not None:
        dl.load_state_dict(state_dict)
        state_dict = None
        ds_state_dict = None
        resumed = True
        print("resumed")
    for i, batch in enumerate(dl):
        print(f"epoch: {epoch}, global_step: {global_step}, batch: {batch}")
        global_step += 1  # consume datapoint
        # simulate error
        if global_step == 124 and not resumed:
            ds_state_dict = ds.state_dict()
            state_dict = dl.state_dict()
            print("checkpoint")
            print("ds_state_dict")
            pprint.pprint(ds_state_dict)
            print("dl_state_dict")
            pprint.pprint(state_dict)
            break

    if state_dict is None:
        ds.set_epoch(epoch)
        epoch += 1
```

The script checkpoints when there are three batches left in the second epoch. After resuming, only the last three batches are repeated in the rest of the epochs.

If it helps, following are the two state_dicts for the dataloader save at the same step with the two settings. The left one is for `PERSISTENT_WORKERS=False`

![Image](https://github.com/user-attachments/assets/c97d6502-d7bd-4ef4-ae2d-66fe1a9732b1)

### Expected behavior

All the elements in the dataset should be iterated through in the epochs following the one where we resumed. The expected behavior can be seen by setting `PERSISTENT_WORKERS=False`. 

### Environment info

torch==2.5.1
datasets==3.3.2
torchdata>=0.9.0