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.
Would you like to give a try to the PR linked above? I believe it should solve your issue.
I tried it and it works partially. On more than 1 GPU, it no longer crashes or hangs at the end!

However, the `RandomSampler` inside the `BatchSampler` yields duplicated indices and `gather_for_metrics` does not remove the extra elements as it is supposed to.

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
@@ -364,10 +364,11 @@ def __iter__(self):
 
     @property
     def total_batch_size(self):
+        batch_sampler = self.sampler if isinstance(self.sampler, BatchSampler) else self.batch_sampler
         return (
-            self.batch_sampler.batch_size
-            if self.batch_sampler.split_batches
-            else (self.batch_sampler.batch_size * self.batch_sampler.num_processes)
+            batch_sampler.batch_size
+            if batch_sampler.split_batches
+            else (batch_sampler.batch_size * batch_sampler.num_processes)
         )
 
     @property
@@ -639,14 +640,17 @@ def prepare_data_loader(
             )
         else:
             # New batch sampler for the current process.
-            if hasattr(dataloader.sampler, "generator"):
-                if dataloader.sampler.generator is None:
-                    dataloader.sampler.generator = torch.Generator()
-                    generator = dataloader.sampler.generator
-                    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)
+            if sampler_is_batch_sampler:
+                sampler = dataloader.sampler.sampler
+            else:
+                sampler = dataloader.batch_sampler.sampler
+            if hasattr(sampler, "generator"):
+                if sampler.generator is None:
+                    sampler.generator = torch.Generator()
+                generator = sampler.generator
+                generator.manual_seed(int(torch.empty((), dtype=torch.int64).random_().item()))
+
             batch_sampler = dataloader.sampler if sampler_is_batch_sampler else dataloader.batch_sampler
             new_batch_sampler = BatchSamplerShard(
                 batch_sampler,
@@ -692,6 +696,7 @@ def prepare_data_loader(
             new_dataset,
             device=device if put_on_device and state.distributed_type != DistributedType.TPU else None,
             sampler=new_batch_sampler,
+            batch_size=getattr(dataloader, "batch_size", _PYTORCH_DATALOADER_KWARGS["batch_size"]),
             rng_types=rng_types,
             generator=generator,
             **kwargs,


Test patch: