huggingface / huggingface/accelerate

`use_stateful_dataloader` under multi-process training restores the cursor, not the shuffle order — sampler permutation is never serialized

Open
#4,195 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
9.9k
Forks
1.5k
Avg merge
5d 2h
Merged PRs (30d)
27

Description

### System Info

```Shell
- accelerate `1.13.0`, torch `2.13.0`, torchdata `0.11.0`, Linux
- Reproduced on multi-GPU DDP and on 2-process CPU/Gloo (reproducer below needs no GPU)
```

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [x] My own task or dataset (give details below)

### Reproduction

```bash
python -m torch.distributed.run --standalone --nproc_per_node=2 issue_repro.py save /tmp/sfd_ckpt
python -m torch.distributed.run --standalone --nproc_per_node=2 issue_repro.py resume /tmp/sfd_ckpt
# on a GPU machine, prefix with: CUDA_VISIBLE_DEVICES=""
```

```python
# issue_repro.py — DDP + use_stateful_dataloader + shuffle=True:
# resume restores the right POSITION but the wrong ORDER.
import os
import sys

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator, DataLoaderConfiguration

MODE, OUT = sys.argv[1], sys.argv[2]
RANK = int(os.environ["RANK"])
SEED, N, BS, SKIP = 1234, 64, 4, 3 # 16 batches/epoch; checkpoint after 3 batches of epoch 1

torch.manual_seed(SEED) # identical in both phases -> reconstruction is deterministic
acc = Accelerator(cpu=True, dataloader_config=DataLoaderConfiguration(use_stateful_dataloader=True))
loader = acc.prepare(DataLoader(TensorDataset(torch.arange(N)), batch_size=BS, shuffle=True, num_workers=0))

def as_list(batch):
return batch[0].tolist() # TensorDataset yields a 1-tuple

if MODE == "save":
os.makedirs(OUT, exist_ok=True)
epoch0 = [as_list(b) for b in loader] # consume full epoch 0
it = iter(loader) # epoch 1
consumed = [as_list(next(it)) for _ in range(SKIP)]
state = loader.state_dict() # per-rank save, as in #3080/#4165
torch.save(state, f"{OUT}/dl_rank{RANK}.bin")
torch.save(
{"epoch0": epoch0, "consumed": consumed, "cont": [as_list(b) for b in it]},
f"{OUT}/ref_rank{RANK}.pt",
)
print(f"[rank {RANK}] saved dl state keys: {list(state.keys())}")
else:
ref = torch.load(f"{OUT}/ref_rank{RANK}.pt", weights_only=False)
loader.load_state_dict(torch.load(f"{OUT}/dl_rank{RANK}.bin", weights_only=False))
resumed = [as_list(b) for b in loader]
print(f"[rank {RANK}] batches consumed before checkpoint: {ref['consumed']}")
print(f"[rank {RANK}] true epoch-1 remainder : {ref['cont']}")
print(f"[rank {RANK}] resumed stream : {resumed}")
print(f"[rank {RANK}] epoch-0 remainder (for reference) : {ref['epoch0'][SKIP:]}")
assert resumed == ref["epoch0"][SKIP:], "not even the deterministic epoch-0 replay"
assert resumed != ref["cont"], "resumed matched the true continuation -- bug not reproduced"
print(f"[rank {RANK}] POSITION restored, ORDER lost: epoch-0's order at epoch-1's position")
```

### Expected behavior

Either:

- exact resume under multi-process as well (preferred — matching torchdata's
native semantics: the restored stream equals the recorded continuation), or
- an explicit warning at save/load time that only the cursor is being restored
when `num_processes > 1` and the sampler chain is shuffled.

Observed output (rank 0; rank 1 analogous):

```
# save
[rank 0] saved dl state keys: ['_index_sampler_state', '_sampler_iter_state', '_sampler_iter_yielded', '_num_yielded', '_IterableDataset_len_called', '_shared_seed', 'fetcher_state', 'dataset_state', '_iterator_finished']

# resume
[rank 0] batches consumed before checkpoint: [[39, 18, 10, 13], [63, 27, 2, 62], [25, 6, 37, 50]]
[rank 0] true epoch-1 remainder : [[55, 16, 59, 0], [22, 46, 17, 57], [23, 45, 19, 51], [21, 26, 40, 36], [20, 5, 41, 32]]
[rank 0] resumed stream : [[15, 41, 52, 61], [5, 50, 18, 14], [31, 53, 36, 47], [35, 9, 23, 58], [22, 7, 51, 55]]
[rank 0] epoch-0 remainder (for reference) : [[15, 41, 52, 61], [5, 50, 18, 14], [31, 53, 36, 47], [35, 9, 23, 58], [22, 7, 51, 55]]
[rank 0] POSITION restored, ORDER lost: epoch-0's order at epoch-1's position
```

Contributor guide

Open the contributing guide

Research direction

Start by running issue_repro.py with the two torch.distributed.run commands and compare the saved continuation with the resumed stream. Trace use_stateful_dataloader through DataLoaderConfiguration and the loader state_dict/load_state_dict calls, focusing on sampler state under multiple processes. Done means the resumed stream matches ref['cont'], or save/load emits the specified warning for shuffled multi-process samplers.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
data, distributed-systems
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
58/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.