Repository: huggingface/accelerate

Issue:
Mixed type batches in data loader
Hello! Thanks for this cool library first of all!

When upgrading to accelerate==0.5.x (from 0.4) I'm receiving the following Exception when iterating through an IterableDataset (prepared by accelerate):
```
  File "/home/martin/miniconda/envs/accelerate/lib/python3.8/site-packages/tqdm/std.py", line 1180, in __iter__
    for obj in iterable:
  File "/home/martin/miniconda/envs/accelerate/lib/python3.8/site-packages/accelerate/data_loader.py", line 351, in __iter__
    batch = concatenate(batches, dim=0)
  File "/home/martin/miniconda/envs/accelerate/lib/python3.8/site-packages/accelerate/utils.py", line 465, in concatenate
    return type(data[0])(**{k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
  File "/home/martin/miniconda/envs/accelerate/lib/python3.8/site-packages/accelerate/utils.py", line 465, in <dictcomp>
    return type(data[0])(**{k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
  File "/home/martin/miniconda/envs/accelerate/lib/python3.8/site-packages/accelerate/utils.py", line 467, in concatenate
    raise TypeError(f"Can only concatenate tensors but got {type(data[0])}")
TypeError: Can only concatenate tensors but got <class 'transformers.tokenization_utils_base.BatchEncoding'>
```

Does accelerate expect the batches to be pure tensors? In previous versions I was able to use mixed type batches. Torch tensors got correctly assigned to their respective devices and other Python types were (correctly) ignored.

I would like to use the prepare function for handling the data splitting in distributed settings.

Thanks for your help! 

A few notes:
* I'm running this on a TPU VM `torch==1.9.1`, therefore trying to use the XLA devices with `torch_xla-1.9-cp38`
* Python 3.8.12, `accelerate==0.5.1`, `transformers==4.11.0`
* I'm using `IterableDataset` (unfortunately unavoidable for now)

Here a script to reproduce:
```python
from accelerate import Accelerator                                                                                                                                                    
from torch.utils.data import DataLoader, IterableDataset                                                                                                                              
import logging                                                                                                                                                                        
from tqdm import tqdm                                                                                                                                                                 
from transformers import AutoTokenizer, AutoModelForSequenceClassification, DataCollatorForTokenClassification                                                                        
                                                                                                                                                                                      
# logging                                                                                                                                                                             
logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)-5.5s] [%(name)-12.12s]: %(message)s')                                                                       
logger = logging.getLogger(__name__)

BATCH_SIZE = 8
NUM_WORKERS = 0
MODEL_NAME = 'bert-base-cased'

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
data_collator = DataCollatorForTokenClassification(tokenizer, padding='max_length')


class MyDataset(IterableDataset):
    def __init__(self):
        super(MyDataset).__init__()
        self.n = int(1e6)

    def __len__(self):
        return self.n

    def __iter__(self):
        for i in range(self.n):
            example = tokenizer('hello world', return_tensors='pt', truncation=True)
            # squeeze all of the examples - we will use collator to pad/truncate and combine into full batch
            example = {k: v.squeeze(0) for k, v in example.items()}
            yield example


def collate_fn(batch):
    return data_collator(batch)


def main():
    accelerator = Accelerator(fp16=False, cpu=False)
    ds = MyDataset()
    data_loader = DataLoader(ds, batch_size=BATCH_SIZE, collate_fn=collate_fn, num_workers=NUM_WORKERS)
    data_loader = accelerator.prepare(data_loader)
    num_batches = len(data_loader)
    pbar = tqdm(data_loader, unit='batch', total=num_batches)

    for j, batch in enumerate(pbar):
        print(f'Process {accelerator.local_process_index}: {batch=} {batch.device=}')


if __name__ == "__main__":
    main()
```


Discussion:
I think this is the same error as in #176, which I fixed yesterday? Would you mind trying on the main branch with a source install to check if you still have the problem? I will do a patch release once you confirm it has been solved.
Nope, same behaviour when using commit 120b82bfcec998c3545615579889cc7b360bc315
Ah yes, this is linked to the central dataloader feature, it's not a duplicate of the issue I was mentioning. I see where the problem is and will try to push a fix soon.

In the meantime, a workaround is to convert your example to a regular dict in your dataset before yielding it:
```
example = {k: v for k, v in example.items()}
```
the problem comes from the fact it is a `BatchEncoding`, which is an instance of a `UserDict`, which for some reason is not considered as a `dict` by Python.
converting to `dict` works - but devices are not being set automatically in this case
Mmm, looking more into it, the `collate_fn` also returns a `BatchEncoding`, which is maybe why you get the probem. On my side, replacing it with:
```
def collate_fn(batch):
    return {k: v for k, v in data_collator(batch).items()}
```
solved the issue and I see batches on the right devices.

Resolving patch:
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -15,6 +15,7 @@
 import importlib
 import os
 import random
+from collections import UserDict
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import List, Optional, Union
@@ -162,6 +163,15 @@ def recursively_apply(func, data, *args, test_type=is_torch_tensor, error_on_oth
                 for o in data
             ),
         )
+    elif isinstance(data, UserDict):
+        return type(data)(
+            {
+                k: recursively_apply(
+                    func, v, *args, test_type=test_type, error_on_other_type=error_on_other_type, **kwargs
+                )
+                for k, v in data.items()
+            }
+        )
     elif isinstance(data, dict):
         return type(data)(
             **{
@@ -300,7 +310,7 @@ def extract_model_from_parallel(model):
 def _tpu_gather(tensor, name="gather tensor"):
     if isinstance(tensor, (list, tuple)):
         return honor_type(tensor, (_tpu_gather(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
-    elif isinstance(tensor, dict):
+    elif isinstance(tensor, (dict, UserDict)):
         return type(tensor)({k: _tpu_gather(v, name=f"{name}_{k}") for k, v in tensor.items()})
     elif not isinstance(tensor, torch.Tensor):
         raise TypeError(f"Can't gather the values of type {type(tensor)}, only of nested list/tuple/dicts of tensors.")
@@ -355,7 +365,7 @@ def _gpu_broadcast_one(tensor, src=0):
 def _tpu_broadcast(tensor, src=0, name="broadcast tensor"):
     if isinstance(tensor, (list, tuple)):
         return honor_type(tensor, (_tpu_broadcast(t, name=f"{name}_{i}") for i, t in enumerate(tensor)))
-    elif isinstance(tensor, dict):
+    elif isinstance(tensor, (dict, UserDict)):
         return type(tensor)({k: _tpu_broadcast(v, name=f"{name}_{k}") for k, v in tensor.items()})
     return xm.mesh_reduce(name, tensor, lambda x: x[src])
 
@@ -438,7 +448,7 @@ def find_batch_size(data):
     """
     if isinstance(data, (tuple, list)):
         return find_batch_size(data[0])
-    elif isinstance(data, dict):
+    elif isinstance(data, (dict, UserDict)):
         for k in data.keys():
             return find_batch_size(data[k])
     elif not isinstance(data, torch.Tensor):
@@ -461,6 +471,8 @@ def concatenate(data, dim=0):
     """
     if isinstance(data[0], (tuple, list)):
         return honor_type(data[0], (concatenate([d[i] for d in data], dim=dim) for i in range(len(data[0]))))
+    elif isinstance(data[0], UserDict):
+        return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
     elif isinstance(data[0], dict):
         return type(data[0])(**{k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
     elif not isinstance(data[0], torch.Tensor):


Test patch:
diff --git a/tests/test_utils.py b/tests/test_utils.py
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import unittest
-from collections import namedtuple
+from collections import UserDict, namedtuple
 
 import torch
 
@@ -54,3 +54,11 @@ def test_send_to_device(self):
         self.assertTrue(torch.equal(result3.b[0].cpu(), tensor))
         self.assertTrue(torch.equal(result3.b[1].cpu(), tensor))
         self.assertEqual(result3.c, 1)
+
+        result4 = send_to_device(UserDict({"a": tensor, "b": [tensor, tensor], "c": 1}), device)
+        self.assertIsInstance(result4, UserDict)
+        self.assertTrue(torch.equal(result4["a"].cpu(), tensor))
+        self.assertIsInstance(result4["b"], list)
+        self.assertTrue(torch.equal(result4["b"][0].cpu(), tensor))
+        self.assertTrue(torch.equal(result4["b"][1].cpu(), tensor))
+        self.assertEqual(result4["c"], 1)