Save and resume the state of a DataLoader
CONTRIBUTOR
Something that'd be nice to have is "manual update of state". One of the learning from training LLMs is the ability to skip some batches whenever we notice huge spike might be handy.
2023-01-23T11:02:30Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-1400160215
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps5TdL_X
https://api.github.com/repos/huggingface/datasets/issues/comments/1400160215/reactions
2023-01-23T11:02:30Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1400160215
https://avatars.githubusercontent.com/u/24695242?v=4
https://api.github.com/users/thomasw21/events{/privacy}
https://api.github.com/users/thomasw21/followers
https://api.github.com/users/thomasw21/following{/other_user}
https://api.github.com/users/thomasw21/gists{/gist_id}
https://github.com/thomasw21
thomasw21
MDQ6VXNlcjI0Njk1MjQy
https://api.github.com/users/thomasw21/orgs
https://api.github.com/users/thomasw21/received_events
https://api.github.com/users/thomasw21/repos
https://api.github.com/users/thomasw21/starred{/owner}{/repo}
https://api.github.com/users/thomasw21/subscriptions
User
https://api.github.com/users/thomasw21
public
CONTRIBUTOR
Your outline spec is very sound and clear, @lhoestq - thank you!

@thomasw21, indeed that would be a wonderful extra feature. In Megatron-Deepspeed we manually drained the dataloader for the range we wanted. I wasn't very satisfied with the way we did it, since its behavior would change if you were to do multiple range skips. I think it should remember all the ranges it skipped and not just skip the last range - since otherwise the data is inconsistent (but we probably should discuss this in a separate issue not to derail this much bigger one).
2023-01-24T01:39:12Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-1401291735
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps5ThgPX
https://api.github.com/repos/huggingface/datasets/issues/comments/1401291735/reactions
2023-01-24T01:45:48Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1401291735
https://avatars.githubusercontent.com/u/10676103?v=4
https://api.github.com/users/stas00/events{/privacy}
https://api.github.com/users/stas00/followers
https://api.github.com/users/stas00/following{/other_user}
https://api.github.com/users/stas00/gists{/gist_id}
https://github.com/stas00
stas00
MDQ6VXNlcjEwNjc2MTAz
https://api.github.com/users/stas00/orgs
https://api.github.com/users/stas00/received_events
https://api.github.com/users/stas00/repos
https://api.github.com/users/stas00/starred{/owner}{/repo}
https://api.github.com/users/stas00/subscriptions
User
https://api.github.com/users/stas00
public
NONE
Hi there! I think this is a critical issue and have an urgent need for it, in my attempt to train on a super large-scale dataset using `datasets`. It is impossible to resume a time-consuming (like one month) experiment by iterating all seen data again, which could possibly cost several days.

@stas00 @thomasw21 @lhoestq Any updates on this problem after 1 year passed?
2024-01-25T09:13:19Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-1909715719
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps5x0_MH
https://api.github.com/repos/huggingface/datasets/issues/comments/1909715719/reactions
2024-01-25T09:26:30Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1909715719
https://avatars.githubusercontent.com/u/55196500?v=4
https://api.github.com/users/yqy2001/events{/privacy}
https://api.github.com/users/yqy2001/followers
https://api.github.com/users/yqy2001/following{/other_user}
https://api.github.com/users/yqy2001/gists{/gist_id}
https://github.com/yqy2001
yqy2001
MDQ6VXNlcjU1MTk2NTAw
https://api.github.com/users/yqy2001/orgs
https://api.github.com/users/yqy2001/received_events
https://api.github.com/users/yqy2001/repos
https://api.github.com/users/yqy2001/starred{/owner}{/repo}
https://api.github.com/users/yqy2001/subscriptions
User
https://api.github.com/users/yqy2001
public
NONE
any update？
2024-02-01T10:25:33Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-1920992826
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps5ygAY6
https://api.github.com/repos/huggingface/datasets/issues/comments/1920992826/reactions
2024-02-01T10:25:33Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1920992826
https://avatars.githubusercontent.com/u/20511825?v=4
https://api.github.com/users/dancingpipi/events{/privacy}
https://api.github.com/users/dancingpipi/followers
https://api.github.com/users/dancingpipi/following{/other_user}
https://api.github.com/users/dancingpipi/gists{/gist_id}
https://github.com/dancingpipi
dancingpipi
MDQ6VXNlcjIwNTExODI1
https://api.github.com/users/dancingpipi/orgs
https://api.github.com/users/dancingpipi/received_events
https://api.github.com/users/dancingpipi/repos
https://api.github.com/users/dancingpipi/starred{/owner}{/repo}
https://api.github.com/users/dancingpipi/subscriptions
User
https://api.github.com/users/dancingpipi
public
MEMBER
No update so far, I wonder if someone implemented a resumable pytorch Sampler somwhere.

Then regarding resuming a streaming dataset, we'd first like to have an efficient way to skip shards automatically but this is not implemented yet
2024-02-02T09:39:04Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-1923430979
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps5ypTpD
https://api.github.com/repos/huggingface/datasets/issues/comments/1923430979/reactions
2024-02-02T09:39:04Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1923430979
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 a draft here for IterableDataset: https://github.com/huggingface/datasets/pull/6658



```python
"""Requires https://github.com/huggingface/datasets/pull/6658 (WIP)"""
from datasets import load_dataset
from torch.utils.data import DataLoader

ds = load_dataset(..., streaming=True)
# ds = ds.map(tokenize)
# ds = ds.shuffle(seed=42, buffer_size=1000)

# Init the dataset state_dict, or load it from a checkpoint
dataset_state_dict = ds.state_dict()

# Resumable training loop
ds.load_state_dict(dataset_state_dict)
dataloader = DataLoader(ds, batch_size=batch_size)
for step, batch in enumerate(dataloader):
    ...
    if step % save_steps == 0:
        dataset_state_dict = ds.state_dict()
```
2024-02-19T15:50:53Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-1952748801
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps50ZJUB
https://api.github.com/repos/huggingface/datasets/issues/comments/1952748801/reactions
2024-02-21T11:25:23Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1952748801
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
Hi @lhoestq - can you provide more information and how to implement on saving and restoring vanilla DataLoader states with map-style datasets?


2024-02-21T08:33:23Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-1956133316
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps50mDnE
https://api.github.com/repos/huggingface/datasets/issues/comments/1956133316/reactions
2024-02-21T08:33:23Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1956133316
https://avatars.githubusercontent.com/u/45981257?v=4
https://api.github.com/users/jwliu36/events{/privacy}
https://api.github.com/users/jwliu36/followers
https://api.github.com/users/jwliu36/following{/other_user}
https://api.github.com/users/jwliu36/gists{/gist_id}
https://github.com/jwliu36
jwliu36
MDQ6VXNlcjQ1OTgxMjU3
https://api.github.com/users/jwliu36/orgs
https://api.github.com/users/jwliu36/received_events
https://api.github.com/users/jwliu36/repos
https://api.github.com/users/jwliu36/starred{/owner}{/repo}
https://api.github.com/users/jwliu36/subscriptions
User
https://api.github.com/users/jwliu36
public
MEMBER
For now the easiest is probably to use the vanilla DataLoader only for batching and multiprocessing, and implement the resuming logic using a `Dataset` (it has `.select()` to skip examples) and a `dataset_state_dict`:


```python
from datasets import load_dataset
from torch.utils.data import DataLoader

ds = load_dataset(...)
# ds = ds.map(tokenize)
# ds = ds.shuffle(seed=42)

# Init the dataset state_dict, or load it from a checkpoint
dataset_state_dict = {"step": 0}  

# Resumable training loop
start_step = dataset_state_dict["step"]
dataloader = DataLoader(ds.select(range(start_step * batch_size, len(ds))), batch_size=batch_size)
for step, batch in enumerate(dataloader, start=start_step):
    ...
    if step % save_steps == 0:
        dataset_state_dict = {"step": step}
```
2024-02-21T11:17:53Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-1956431084
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps50nMTs
https://api.github.com/repos/huggingface/datasets/issues/comments/1956431084/reactions
2024-02-21T11:17:53Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1956431084
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
Hello, I found a similar implementation online that seems to solve your problem.  https://github.com/facebookresearch/vissl/blob/main/vissl/data/data_helper.py#L93
it looks like we can set_start_iter in StatefulDistributedSampler to implement the stateful  resume requirement we want.


2024-03-19T02:50:43Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2005651101
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps53i86d
https://api.github.com/repos/huggingface/datasets/issues/comments/2005651101/reactions
2024-03-19T02:50:43Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2005651101
https://avatars.githubusercontent.com/u/8856191?v=4
https://api.github.com/users/xgbj/events{/privacy}
https://api.github.com/users/xgbj/followers
https://api.github.com/users/xgbj/following{/other_user}
https://api.github.com/users/xgbj/gists{/gist_id}
https://github.com/xgbj
xgbj
MDQ6VXNlcjg4NTYxOTE=
https://api.github.com/users/xgbj/orgs
https://api.github.com/users/xgbj/received_events
https://api.github.com/users/xgbj/repos
https://api.github.com/users/xgbj/starred{/owner}{/repo}
https://api.github.com/users/xgbj/subscriptions
User
https://api.github.com/users/xgbj
public
NONE
Hi y'all, @lhoestq I wanted to flag that we currently have a StatefulDataLoader in `pytorch/data/torchdata` that has state_dict/load_state_dict methods, which will call a dataset's state_dict/load_state_dict methods but also handle multiprocessing under the hood. Any chance we can collaborate on this and try to get them to work well together? Please have a look here for some basic examples: https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader#saving-and-loading-state 
2024-04-29T22:27:48Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2083794563
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps58NC6D
https://api.github.com/repos/huggingface/datasets/issues/comments/2083794563/reactions
2024-04-29T22:27:48Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2083794563
https://avatars.githubusercontent.com/u/5349063?v=4
https://api.github.com/users/andrewkho/events{/privacy}
https://api.github.com/users/andrewkho/followers
https://api.github.com/users/andrewkho/following{/other_user}
https://api.github.com/users/andrewkho/gists{/gist_id}
https://github.com/andrewkho
andrewkho
MDQ6VXNlcjUzNDkwNjM=
https://api.github.com/users/andrewkho/orgs
https://api.github.com/users/andrewkho/received_events
https://api.github.com/users/andrewkho/repos
https://api.github.com/users/andrewkho/starred{/owner}{/repo}
https://api.github.com/users/andrewkho/subscriptions
User
https://api.github.com/users/andrewkho
public
MEMBER
Fantastic ! This will help pushing our IterableDataset state_dict implementation at https://github.com/huggingface/datasets/pull/6658 :) I'll check if there is anything missing to maker them work together, and add tests and some docs referring to the StatefulDataLoader :)
2024-04-30T09:08:58Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2084782796
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps58Q0LM
https://api.github.com/repos/huggingface/datasets/issues/comments/2084782796/reactions
2024-04-30T09:08:58Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2084782796
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
Ah I just saw this disclaimer in the torchdata README and it feels like people should not rely on it. Should the StatefulDataLoader live elsewhere @andrewkho ?

> ⚠️ As of July 2023, we have paused active development on TorchData and have paused new releases. We have learnt a lot from building it and hearing from users, but also believe we need to re-evaluate the technical design and approach given how much the industry has changed since we began the project. During the rest of 2023 we will be re-evaluating our plans in this space. Please reach out if you suggestions or comments (please use https://github.com/pytorch/data/issues/1196 for feedback).
2024-04-30T12:52:16Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2085248704
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps58Sl7A
https://api.github.com/repos/huggingface/datasets/issues/comments/2085248704/reactions
2024-04-30T12:52:35Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2085248704
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 Good find, we are in the midst of updating this disclaimer as we're re-starting development and regular releases, though our approach will be to iterate on DL V1 (ie StatefulDataLoader) instead of continuing development on datapipes+DLV2. Let's discuss on a call at some point to figure out the best path forward! 
2024-04-30T18:30:54Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2086464910
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps58XO2O
https://api.github.com/repos/huggingface/datasets/issues/comments/2086464910/reactions
2024-04-30T18:30:54Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2086464910
https://avatars.githubusercontent.com/u/5349063?v=4
https://api.github.com/users/andrewkho/events{/privacy}
https://api.github.com/users/andrewkho/followers
https://api.github.com/users/andrewkho/following{/other_user}
https://api.github.com/users/andrewkho/gists{/gist_id}
https://github.com/andrewkho
andrewkho
MDQ6VXNlcjUzNDkwNjM=
https://api.github.com/users/andrewkho/orgs
https://api.github.com/users/andrewkho/received_events
https://api.github.com/users/andrewkho/repos
https://api.github.com/users/andrewkho/starred{/owner}{/repo}
https://api.github.com/users/andrewkho/subscriptions
User
https://api.github.com/users/andrewkho
public
MEMBER
As a heads up, `IterableDataset` state_dict has been added in https://github.com/huggingface/datasets/pull/6658

...and it works out of the box with the `torchdata` `StatefulDataLoader` :)

See the docs at https://huggingface.co/docs/datasets/main/en/use_with_pytorch#checkpoint-and-resume
2024-07-22T11:13:17Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2242704062
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps6FrPK-
https://api.github.com/repos/huggingface/datasets/issues/comments/2242704062/reactions
2024-07-22T11:14:18Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2242704062
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
CONTRIBUTOR
amazing! Thank you, @lhoestq 

does it work with non-iterable dataset as well? the docs only mention iterable dataset
2024-07-24T20:13:57Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2248818962
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps6GCkES
https://api.github.com/repos/huggingface/datasets/issues/comments/2248818962/reactions
2024-07-24T20:13:57Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2248818962
https://avatars.githubusercontent.com/u/10676103?v=4
https://api.github.com/users/stas00/events{/privacy}
https://api.github.com/users/stas00/followers
https://api.github.com/users/stas00/following{/other_user}
https://api.github.com/users/stas00/gists{/gist_id}
https://github.com/stas00
stas00
MDQ6VXNlcjEwNjc2MTAz
https://api.github.com/users/stas00/orgs
https://api.github.com/users/stas00/received_events
https://api.github.com/users/stas00/repos
https://api.github.com/users/stas00/starred{/owner}{/repo}
https://api.github.com/users/stas00/subscriptions
User
https://api.github.com/users/stas00
public
MEMBER
It's for iterable dataset only. For regular dataset I believe the sampler should implement state_dict, but maybe @andrewkho might know best how to resume a regular dataset with torchdata
2024-07-25T12:00:53Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2250156269
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps6GHqjt
https://api.github.com/repos/huggingface/datasets/issues/comments/2250156269/reactions
2024-07-25T12:00:53Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2250156269
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
@stas00 stateful dataloader will save and resume samplers for map style datasets. If no state_dict/load_state_dict is provided by the sampler, it will naively skip samples to fast forward. See here for more details https://github.com/pytorch/data/blob/main/torchdata/stateful_dataloader/README.md 

Hope this helps! 
2024-07-25T15:10:16Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2250629343
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps6GJeDf
https://api.github.com/repos/huggingface/datasets/issues/comments/2250629343/reactions
2024-07-25T15:12:13Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2250629343
https://avatars.githubusercontent.com/u/5349063?v=4
https://api.github.com/users/andrewkho/events{/privacy}
https://api.github.com/users/andrewkho/followers
https://api.github.com/users/andrewkho/following{/other_user}
https://api.github.com/users/andrewkho/gists{/gist_id}
https://github.com/andrewkho
andrewkho
MDQ6VXNlcjUzNDkwNjM=
https://api.github.com/users/andrewkho/orgs
https://api.github.com/users/andrewkho/received_events
https://api.github.com/users/andrewkho/repos
https://api.github.com/users/andrewkho/starred{/owner}{/repo}
https://api.github.com/users/andrewkho/subscriptions
User
https://api.github.com/users/andrewkho
public
CONTRIBUTOR
Thank you very much for clarifying that, Andrew.


2024-07-25T17:47:43Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2251071906
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps6GLKGi
https://api.github.com/repos/huggingface/datasets/issues/comments/2251071906/reactions
2024-07-25T17:47:43Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2251071906
https://avatars.githubusercontent.com/u/10676103?v=4
https://api.github.com/users/stas00/events{/privacy}
https://api.github.com/users/stas00/followers
https://api.github.com/users/stas00/following{/other_user}
https://api.github.com/users/stas00/gists{/gist_id}
https://github.com/stas00
stas00
MDQ6VXNlcjEwNjc2MTAz
https://api.github.com/users/stas00/orgs
https://api.github.com/users/stas00/received_events
https://api.github.com/users/stas00/repos
https://api.github.com/users/stas00/starred{/owner}{/repo}
https://api.github.com/users/stas00/subscriptions
User
https://api.github.com/users/stas00
public
NONE
👋 I am trying to use `HF Streaming Dataset + TorchDDP  + Stateful Dataloader`, to train using multiple nodes and large datasets. 

So far, I have been able to use  HF Streaming Dataset + TorchDDP with Vanilla Datasets. To do so, I implemented a custom iterable to make sure that shards are distributed across the multiple nodes, while letting the `dataset` take care of the multiple workers. The implementation uses `split_dataset_by_node`:

```
import torch
from torch.distributed import get_rank, get_world_size
from torch.utils.data import DataLoader, IterableDataset

class MyIterableDataset(IterableDataset):
    def __init__(self, dataset):
        super().__init__()
        self.dataset = dataset
        self._iterable_by_node = None

    def __iter__(self):
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            world_size = get_world_size()
            process_rank = get_rank()
        else:
            world_size = 1
            process_rank = 0

        if world_size > 1:
            self._iterable_by_node = split_dataset_by_node(
                    self.dataset, rank=process_rank, world_size=world_size
                )
        else:
            self._iterable_by_node = self.dataset

        for example in self._iterable_by_node:
            # Trying with _state_dict, since `.state_dict()` creates a copy
            self._state_dict.update(self._iterable_by_node._state_dict)
            yield example

    def state_dict(self):
        return self._state_dict

    def load_state_dict(self, state):
        pass # Not implemented yet

```
 
This doesn't seem to work with `StatefulDataLoader` though. I can see the state of the worker's dataset being updated in its corresponding workers' processes, but somehow the updates are not propagated back to the main process. I have tried with different variants of the above code without success.  

I confirmed that if I skip the custom class and pass `dataset` directly to the loader as in the [docs](https://huggingface.co/docs/datasets/main/en/use_with_pytorch#checkpoint-and-resume), the StatefulDataLoader sees the updates for each worker. However, if I do this, multiple nodes will see the same examples, which I definitely don't want.

Is there something I am missing? It would be nice if streaming `dataset`s would support by default the multinode training (unless it already does it and I am missing something).



2024-11-25T12:31:42Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2497894859
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps6U4tnL
https://api.github.com/repos/huggingface/datasets/issues/comments/2497894859/reactions
2024-11-25T12:31:42Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2497894859
https://avatars.githubusercontent.com/u/148977796?v=4
https://api.github.com/users/diego-andres-ardila/events{/privacy}
https://api.github.com/users/diego-andres-ardila/followers
https://api.github.com/users/diego-andres-ardila/following{/other_user}
https://api.github.com/users/diego-andres-ardila/gists{/gist_id}
https://github.com/diego-andres-ardila
diego-andres-ardila
U_kgDOCOE4hA
https://api.github.com/users/diego-andres-ardila/orgs
https://api.github.com/users/diego-andres-ardila/received_events
https://api.github.com/users/diego-andres-ardila/repos
https://api.github.com/users/diego-andres-ardila/starred{/owner}{/repo}
https://api.github.com/users/diego-andres-ardila/subscriptions
User
https://api.github.com/users/diego-andres-ardila
public
MEMBER
Hi ! Have you tried using `split_dataset_by_node()` and pass the result to the StatefulDataLoader ?

```python
dataloader = StatefulDataLoader(split_dataset_by_node(dataset, rank=process_rank, world_size=world_size))
```
2024-11-25T13:59:12Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2498097053
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps6U5e-d
https://api.github.com/repos/huggingface/datasets/issues/comments/2498097053/reactions
2024-11-25T13:59:12Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2498097053
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
> Hi ! Have you tried using split_dataset_by_node() and pass the result to the StatefulDataLoader ?

@lhoestq  it took me some time to test, but it works like a charm. Thanks for the pointer. Totally missed this 🤦. 
2024-11-27T01:19:20Z
https://github.com/huggingface/datasets/issues/5454#issuecomment-2502408841
https://api.github.com/repos/huggingface/datasets/issues/5454
IC_kwDODunzps6VJ7qJ
https://api.github.com/repos/huggingface/datasets/issues/comments/2502408841/reactions
2024-11-27T01:19:20Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2502408841
https://avatars.githubusercontent.com/u/148977796?v=4
https://api.github.com/users/diego-andres-ardila/events{/privacy}
https://api.github.com/users/diego-andres-ardila/followers
https://api.github.com/users/diego-andres-ardila/following{/other_user}
https://api.github.com/users/diego-andres-ardila/gists{/gist_id}
https://github.com/diego-andres-ardila
diego-andres-ardila
U_kgDOCOE4hA
https://api.github.com/users/diego-andres-ardila/orgs
https://api.github.com/users/diego-andres-ardila/received_events
https://api.github.com/users/diego-andres-ardila/repos
https://api.github.com/users/diego-andres-ardila/starred{/owner}{/repo}
https://api.github.com/users/diego-andres-ardila/subscriptions
User
https://api.github.com/users/diego-andres-ardila
public
It would be nice when using `datasets` with a PyTorch DataLoader to be able to resume a training from a DataLoader state (e.g. to resume a training that crashed)

What I have in mind (but lmk if you have other ideas or comments):

For map-style datasets, this requires to have a PyTorch Sampler state that can be saved and reloaded per node and worker.

For iterable datasets, this requires to save the state of the dataset iterator, which includes:
- the current shard idx and row position in the current shard
- the epoch number
- the rng state
- the shuffle buffer

Right now you can already resume the data loading of an iterable dataset by using `IterableDataset.skip` but it takes a lot of time because it re-iterates on all the past data until it reaches the resuming point.

cc @stas00 @sgugger 