Repository: huggingface/accelerate

Issue:
Cant guarantee the reproducibility when using DataLoader shuffle
### System Info

```Shell
- `Accelerate` version: 0.28.0
- Platform: Windows-10-10.0.22000-SP0
- Python version: 3.11.5
- Numpy version: 1.26.4
- PyTorch version (GPU?): 2.2.1+cu118 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- System RAM: 15.37 GB
- GPU type: GeForce GTX 1650 Ti
- `Accelerate` default config:
        Not found
```


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

```python
import torch
from accelerate import Accelerator

accelerator = Accelerator()
dataset = range(5)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)
dataloader = accelerator.prepare(dataloader)

for idx, batch in enumerate(dataloader):
    print(f'Batch {idx}: {batch}')
    if idx == 2:
        accelerator.save_state('ckpt')
        print('Checkpoint saved.')
    
print('Resumed from chechpoint')
accelerator.load_state('ckpt')
dataloader = accelerator.skip_first_batches(dataloader, 3)
for idx, batch in enumerate(dataloader, 3):
    print(f'Batch {idx}: {batch}')
```
And the results were:
```
Batch 0: tensor([2], device='cuda:0')
Batch 1: tensor([3], device='cuda:0')
Batch 2: tensor([0], device='cuda:0')
Checkpoint saved.
Batch 3: tensor([1], device='cuda:0')
Batch 4: tensor([4], device='cuda:0')
Resumed from chechpoint
Batch 3: tensor([1], device='cuda:0')
Batch 4: tensor([2], device='cuda:0')
```

### Expected behavior

After saving the ckpt, the original dataloader yields samples `[1], [4]`. However, when resuming from the ckpt and skipping first 3 samples as instructed by [this page](https://huggingface.co/docs/accelerate/en/usage_guides/checkpoint), the recovered dataloader yields `[1], [2]`. 

This indicates that the accelerator does not properly record and recover randomness, which may lead to the model being trained on identical samples within a single epoch. As a result, some samples may never be encountered during the entire epoch if checkpoints are set in the middle of an epoch.


Discussion:
Set `use_seedable_sampler=True` in a `DataLoaderConfiguration` + pass it in to your `Accelerator` and this will work. 
> Set `use_seedable_sampler=True` in a `DataLoaderConfiguration` + pass it in to your `Accelerator` and this will work.

Thank you for your kind reply. I am not sure if I am using it rightly, but the results still differ:

```python
import torch
from accelerate import Accelerator, DataLoaderConfiguration

accelerator = Accelerator(
    dataloader_config=DataLoaderConfiguration(
        use_seedable_sampler=True,
    )
)
dataset = range(5)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)
dataloader = accelerator.prepare(dataloader)

for idx, batch in enumerate(dataloader):
    print(f'Batch {idx}: {batch}')
    if idx == 2:
        accelerator.save_state('ckpt')
        print('Checkpoint saved.')
    
print('Resumed from chechpoint')
accelerator.load_state('ckpt')
dataloader = accelerator.skip_first_batches(dataloader, 3)
for idx, batch in enumerate(dataloader, 3):
    print(f'Batch {idx}: {batch}')
```
```
Batch 0: tensor([2], device='cuda:0')
Batch 1: tensor([3], device='cuda:0')
Batch 2: tensor([4], device='cuda:0')
Checkpoint saved.
Batch 3: tensor([0], device='cuda:0')
Batch 4: tensor([1], device='cuda:0')
Resumed from chechpoint
Batch 3: tensor([4], device='cuda:0')
Batch 4: tensor([0], device='cuda:0')
```
Make sure you’re setting your seed as well via `accelerate.utils.set_seed`
> Make sure you’re setting your seed as well via `accelerate.utils.set_seed`

Thank you again. After following the advice, the problem still exists 😢
```python
import torch
from accelerate import Accelerator, DataLoaderConfiguration
from accelerate.utils import set_seed

set_seed(42)

accelerator = Accelerator(
    dataloader_config=DataLoaderConfiguration(
        use_seedable_sampler=True,
    ),
)
dataset = range(5)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)
dataloader = accelerator.prepare(dataloader)

for idx, batch in enumerate(dataloader):
    print(f'Batch {idx}: {batch}')
    if idx == 2:
        accelerator.save_state('ckpt')
        print('Checkpoint saved.')
    
print('Resumed from chechpoint')
accelerator.load_state('ckpt')
dataloader = accelerator.skip_first_batches(dataloader, 3)
for idx, batch in enumerate(dataloader, 3):
    print(f'Batch {idx}: {batch}')
```
```
Batch 0: tensor([2], device='cuda:0')
Batch 1: tensor([4], device='cuda:0')
Batch 2: tensor([3], device='cuda:0')
Checkpoint saved.
Batch 3: tensor([0], device='cuda:0')
Batch 4: tensor([1], device='cuda:0')
Resumed from chechpoint
Batch 3: tensor([2], device='cuda:0')
Batch 4: tensor([0], device='cuda:0')
```
Oh wait, I totally see the issue here. `save_state` has nothing to do with the `DataLoaders`, we don't save that. SO not sure what the goal is here by doing that. It'll work if you recreate your dataloader instead:

```python
import torch
from accelerate.utils import set_seed
from accelerate import Accelerator, DataLoaderConfiguration

config = DataLoaderConfiguration(use_seedable_sampler=True)

accelerator = Accelerator(dataloader_config=config)
dataset = range(5)
set_seed(5)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)
dataloader = accelerator.prepare(dataloader)

for idx, batch in enumerate(dataloader):
    print(f'Batch {idx}: {batch}')
    if idx == 2:
        accelerator.save_state('ckpt')
        print('Checkpoint saved.')
    
print('Resumed from chechpoint')
accelerator.load_state('ckpt')
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)
dataloader = accelerator.prepare(dataloader)
dataloader = accelerator.skip_first_batches(dataloader, 3)
for idx, batch in enumerate(dataloader, 3):
    print(f'Batch {idx}: {batch}')
```
Got it! I will try this tomorrow. Thank you so much!
@muellerzr Hi! Sorry for bothering you again. I've extended the previous code to a multi-process version. Interestingly, while it works well in a single-process setup, it encounters issues in the multi-process configuration:
```python
import torch
from accelerate import Accelerator, DataLoaderConfiguration
from accelerate.utils import set_seed
from accelerate import notebook_launcher

def main():
    for use_ckpt in [False, True]:
        set_seed(42)
        accelerator = Accelerator(
            dataloader_config=DataLoaderConfiguration(use_seedable_sampler=True)
        )
        dataset = range(10)
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)
        dataloader = accelerator.prepare(dataloader)
        
        if use_ckpt:
            print(f'{accelerator.process_index}: Resumed from checkpoint')
            accelerator.load_state('ckpt')
            dataloader = accelerator.skip_first_batches(dataloader, 3)
            for idx, batch in enumerate(dataloader, 3):
                print(f'{accelerator.process_index}: Batch {idx}: {batch}')
        else:
            for idx, batch in enumerate(dataloader):
                if idx == 2:
                    accelerator.save_state('ckpt')
                    print(f'{accelerator.process_index}: Checkpoint saved.')
                elif idx > 2:
                    print(f'{accelerator.process_index}: Batch {idx}: {batch}')
    
notebook_launcher(main, num_processes=1)
```
When setting `num_processes=1` for `notebook_launcher`, the result is:
```
0: Checkpoint saved.
0: Batch 3: 8
0: Batch 4: 4
0: Batch 5: 5
0: Batch 6: 0
0: Batch 7: 9
0: Batch 8: 3
0: Batch 9: 7

0: Resumed from checkpoint
0: Batch 3: 8
0: Batch 4: 4
0: Batch 5: 5
0: Batch 6: 0
0: Batch 7: 9
0: Batch 8: 3
0: Batch 9: 7
```
Which is expected. However, when setting `num_processes=2`, the result turns to:
```
0: Checkpoint saved.
0: Batch 3: 6
0: Batch 4: 2

0: Resumed from checkpoint
0: Batch 3: 2
0: Batch 4: 1

1: Checkpoint saved.
1: Batch 3: 3
1: Batch 4: 7

1: Resumed from checkpoint
1: Batch 3: 0
1: Batch 4: 8
```
We can see that the resumed results for processes 0 and 1 are `[2, 1], [0, 8]`, but the expected are `[6, 2], [3, 7]`. More interestingly, when inserting `set_seed(42)` right after `accelerator.load_state('ckpt')`, the results become expected:
```python
...
if use_ckpt:
            print(f'{accelerator.process_index}: Resumed from checkpoint')
            accelerator.load_state('ckpt')
            set_seed(42)
            ...
```
```
1: Checkpoint saved.
1: Batch 3: 3
1: Batch 4: 7

1: Resumed from checkpoint
1: Batch 3: 3
1: Batch 4: 7

0: Checkpoint saved.
0: Batch 3: 6
0: Batch 4: 2

0: Resumed from checkpoint
0: Batch 3: 6
0: Batch 4: 2
```

Resolving patch:
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -971,6 +971,8 @@ def prepare_data_loader(
             dataloader.sampler.sampler = sampler
         else:
             dataloader.batch_sampler.sampler = sampler
+            if hasattr(dataloader.batch_sampler, "batch_sampler"):
+                dataloader.batch_sampler.batch_sampler.sampler = sampler
     if state.distributed_type == DistributedType.XLA:
         return MpDeviceLoaderWrapper(dataloader, device)
     return dataloader


Test patch:
diff --git a/src/accelerate/test_utils/scripts/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
--- a/src/accelerate/test_utils/scripts/test_script.py
+++ b/src/accelerate/test_utils/scripts/test_script.py
@@ -379,6 +379,25 @@ def check_seedable_sampler():
     assert torch.allclose(original_items, new_items), "Did not obtain the same items with the same seed and epoch."
 
 
+def check_seedable_sampler_in_batch_sampler_shard():
+    set_seed(42)
+
+    config = DataLoaderConfiguration(use_seedable_sampler=True)
+    accelerator = Accelerator(dataloader_config=config)
+    assert accelerator.num_processes > 1, "This test requires more than one process."
+
+    dataloader = DataLoader(list(range(10)), batch_size=1, shuffle=True)
+    prepared_data_loader = prepare_data_loader(
+        dataloader=dataloader,
+        use_seedable_sampler=True,
+    )
+
+    target_sampler = prepared_data_loader.batch_sampler.batch_sampler.sampler
+    assert isinstance(
+        target_sampler, SeedableRandomSampler
+    ), "Sampler in BatchSamplerShard is not SeedableRandomSampler."
+
+
 def mock_training(length, batch_size, generator, use_seedable_sampler=False):
     set_seed(42)
     generator.manual_seed(42)
@@ -742,6 +761,9 @@ def main():
         custom_sampler_check()
         check_seedable_sampler()
 
+    if state.num_processes > 1:
+        check_seedable_sampler_in_batch_sampler_shard()
+
     # Trainings are not exactly the same in DeepSpeed and CPU mode
     if state.distributed_type == DistributedType.DEEPSPEED:
         return