Repository: huggingface/accelerate

Issue:
`accelerator.prepare(dataloader)` sanity check fails when `batch_sampler` is given and `split_batches` is `True`
### System Info

```Shell
Ubuntu 20.04, Python 3.10, accelerate 0.25.0, torch 2.1.2
```


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

I am attempting to manually pass a `BatchSampler` to `DataLoader` while enabling the `split_batches` option.

```
from accelerate import Accelerator
from torch.utils.data import Dataset, DataLoader, BatchSampler, RandomSampler


class SimpleDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __getitem__(self, index) -> dict:
        return self.data[index]

    def __len__(self):
        return len(self.data)

accelerator = Accelerator(split_batches=True)

dataset = SimpleDataset([0, 1, 2, 3, 4])
batch_sampler = BatchSampler(RandomSampler(dataset), 2, True)
loader = DataLoader(dataset, batch_sampler=batch_sampler)
accelerator.prepare(loader)
```

But it fails at the sanity check of https://github.com/huggingface/accelerate/blob/ceb7c699bc36bdb3bbf32cceaaca2d1ceaf62dae/src/accelerate/data_loader.py#L827

```
Traceback (most recent call last):
  File "test.py", line 20, in <module>
    accelerator.prepare(loader)
  File "/envs/flash/lib/python3.10/site-packages/accelerate/accelerator.py", line 1213, in prepare
    result = tuple(
  File "/envs/flash/lib/python3.10/site-packages/accelerate/accelerator.py", line 1214, in <genexpr>
    self._prepare_one(obj, first_pass=True, device_placement=d) for obj, d in zip(args, device_placement)
  File "/envs/flash/lib/python3.10/site-packages/accelerate/accelerator.py", line 1092, in _prepare_one
    return self.prepare_data_loader(obj, device_placement=device_placement)
  File "/envs/flash/lib/python3.10/site-packages/accelerate/accelerator.py", line 1788, in prepare_data_loader
    prepared_data_loader = prepare_data_loader(
  File "/envs/flash/lib/python3.10/site-packages/accelerate/data_loader.py", line 820, in prepare_data_loader
    if split_batches and dataloader.batch_size > 1 and dataloader.batch_size % num_processes != 0:
TypeError: '>' not supported between instances of 'NoneType' and 'int'
```

The `dataloader.batch_size` is set to `None` in `dataloader.__init__` if `batch_sampler` is given, see https://github.com/pytorch/pytorch/blob/f6dfbffb3bb46ada6fe66b5da4f989f9d4d69b3c/torch/utils/data/dataloader.py#L336
```
        if batch_sampler is not None:
            # auto_collation with custom batch_sampler
            if batch_size != 1 or shuffle or sampler is not None or drop_last:
                raise ValueError('batch_sampler option is mutually exclusive '
                                 'with batch_size, shuffle, sampler, and '
                                 'drop_last')
            batch_size = None
            drop_last = False
```



### Expected behavior

Pass the check. 


Discussion:
I suppose a simple modification could fix this issue:
```
    # Sanity check
    batch_size = dataloader.batch_size if dataloader.batch_size is not None else dataloader.batch_sampler.batch_size
    if split_batches and batch_size > 1 and batch_size % num_processes != 0:
        raise ValueError(
            ...
```

Can I contribute this simple change?
@izhx yes feel free! Just make sure all the dataloader related tests pass when doing so :) 

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
@@ -837,12 +837,26 @@ def prepare_data_loader(
         process_index = state.process_index
 
     # Sanity check
-    batch_size = dataloader.batch_size if dataloader.batch_size is not None else dataloader.batch_sampler.batch_size
-    if split_batches and batch_size > 1 and batch_size % num_processes != 0:
-        raise ValueError(
-            f"To use a `DataLoader` in `split_batches` mode, the batch size ({dataloader.batch_size}) "
-            f"needs to be a round multiple of the number of processes ({num_processes})."
-        )
+    if split_batches:
+        if dataloader.batch_size is not None:
+            batch_size_for_check = dataloader.batch_size
+        else:
+            # For custom batch_sampler
+            if hasattr(dataloader.batch_sampler, "batch_size"):
+                batch_size_for_check = dataloader.batch_sampler.batch_size
+            else:
+                raise ValueError(
+                    "In order to use `split_batches==True` you must have a `batch_size` attribute either in the passed "
+                    "`dataloader` or `dataloader.batch_sampler` objects, and it has to return a natural number. "
+                    "Your `dataloader.batch_size` is None and `dataloader.batch_sampler` "
+                    f"(`{type(dataloader.batch_sampler)}`) does not have the `batch_size` attribute set."
+                )
+
+        if batch_size_for_check > 1 and batch_size_for_check % num_processes != 0:
+            raise ValueError(
+                f"To use a `DataLoader` in `split_batches` mode, the batch size ({dataloader.batch_size}) "
+                f"needs to be a round multiple of the number of processes ({num_processes})."
+            )
 
     new_dataset = dataloader.dataset
     # Iterable dataset doesn't like batch_sampler, but data_loader creates a default one for it


Test patch: