Repository: huggingface/accelerate

Issue:
Error in prepared DataLoader with BatchSampler
### System Info

```Shell
accelerate: 0.12.0
OS: Linux 5.4.188+ (Colab)
Python: 3.7.13
numpy: 1.21.6
torch: 1.12.1+cu113
config: 1 CPU
```


### 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

MRE : https://colab.research.google.com/drive/17krCJCF_nWtNFSiMBo3oz12l7eX1bBZ6

First of all, thanks for this library and the great docs and examples that comes with it 😄!

I am using a custom torch Dataset that contains a Hugging Face Dataset (pyarrow) instance. Therefore, as indicated in the Datasets docs (https://huggingface.co/docs/datasets/v2.4.0/en/use_with_pytorch#use-a-batchsampler), I tried to use a BatchSampler to reduce the number of queries. However, I have not been able yet to make it work yet with accelerate.

I tried many different possibilities, one of which works one CPU or one GPU, but gets stuck when using distributed training.

Thanks for your help!


Discussion:
You are passing your batch sampler using the `sampler` argument, not the `batch_sampler` argument when creating your `DataLoader`. That may be the source of the issue.
 Thanks @sgugger for your answer. This confused me too at first, but it is apparently the way it should be done, as shown in the [HuggingFace docs](https://huggingface.co/docs/datasets/v2.4.0/en/use_with_pytorch#use-a-batchsampler) and in this [StackOverflow answer](https://stackoverflow.com/a/61458670).
I believe this a typo in both cases.
Thanks for your suggestion. I tried it and added a cell in the MRE above where I pass the batch sampler to the `batch_sampler` argument.

Sadly,

```py
sampler = BatchSampler(RandomSampler(dataset), batch_size=4, drop_last=False)
loader = DataLoader(dataset, batch_sampler=sampler)
```

produces the same results as the "basic" version :

```py
loader = DataLoader(dataset, batch_size=4, shuffle=True, drop_last=False)
```

which accesses the dataset one index at a time (not in bulk to reduce overhead). I believe this is why the HF docs uses the `sampler` argument.
Oh ok. I did not know about this use of a sampler as a batch sampler. It is indeed not supported by Accelerate. Will have a look on how to add support when I have a bit of time.

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
@@ -622,6 +622,7 @@ def prepare_data_loader(
     new_dataset = dataloader.dataset
     # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it
     new_batch_sampler = dataloader.batch_sampler if not isinstance(new_dataset, IterableDataset) else None
+    sampler_is_batch_sampler = False
     generator = getattr(dataloader, "generator", None)
     # No change if no multiprocess
     if num_processes != 1 and not dispatch_batches:
@@ -645,8 +646,10 @@ def prepare_data_loader(
                     generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
             elif getattr(dataloader.batch_sampler, "generator", None) is not None:
                 generator = dataloader.batch_sampler.generator
+            sampler_is_batch_sampler = isinstance(dataloader.sampler, BatchSampler)
+            batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler
             new_batch_sampler = BatchSamplerShard(
-                dataloader.batch_sampler,
+                batch_sampler,
                 num_processes=num_processes,
                 process_index=process_index,
                 split_batches=split_batches,
@@ -684,6 +687,15 @@ def prepare_data_loader(
             _drop_last=dataloader.drop_last,
             **kwargs,
         )
+    elif sampler_is_batch_sampler:
+        dataloader = DataLoaderShard(
+            new_dataset,
+            device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
+            sampler=new_batch_sampler,
+            rng_types=rng_types,
+            generator=generator,
+            **kwargs,
+        )
     else:
         dataloader = DataLoaderShard(
             new_dataset,


Test patch: