Repository: huggingface/accelerate

Issue:
Doesn't work with new PyTorch TensorDict
### System Info

```Shell
`Accelerate` version: 0.26.1
- Platform: Linux-5.19.0-45-generic-x86_64-with-glibc2.31
- Python version: 3.9.16
- Numpy version: 1.23.4
- PyTorch version (GPU?): 2.2.0+cu121 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- System RAM: 88.38 GB
- GPU type: NVIDIA A100-SXM4-80GB
- `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

1. Use a TensorDict as the dataset passed into a PyTorch dataloader. For example, see the PyTorch example here: https://pytorch.org/tensordict/tutorials/data_fashion.html
2. Use accelerate to `accelerator.prepare` the dataloader.
3. Try to load batches from the prepared dataloader.

When I did this, I get the following error: 
```
ValueError: batch size was not specified when creating the TensorDict instance and it could not be retrieved from source.
```

This happens when Accelerate is putting the tensor on the right device. It seems to lose track of its batch_size attribute in the process. When I got rid of accelerator.prepare on my dataset (which is fine for now cause I'm training on 1 GPU) this error went away.

### Expected behavior

Expected behavior: it successfully returns batches from the dataloader.


Discussion:
Hi, thanks for flagging this.
I'll have a look at it on the tensordict side to see if there's anything we can do.

We were initially reluctant to have the empty batch-size assumed by default but it wouldn't be a big deal to allow it if that unblocks things on accelerate end.
@vmoens Wow how did you even find this issue 😄 Really appreciate this! I just started using TensorDicts and it's such a cool and useful abstraction, I was able to delete a lot of code. Excited to see it get battle-hardened and integrated into more things like Accelerate!
Would accelerate folks be happy using pytree for [`send_to_device`](https://github.com/huggingface/accelerate/blob/b3d2111708a480ddae60c6e1fb80f360326a0dc8/src/accelerate/utils/operations.py#L144)?


Ref: https://github.com/pytorch/pytorch/blob/main/torch/utils/_pytree.py

That would be much more generic than the current implementation no?

Happy to give it a shot if you'd like!
We can see what that API would wind up looking like, certainly! If it's more generic than that would definitely be a win. We just need to make sure it can work with the external/additional device supports like `xpu`
Let me make a quick prototype for you later today!
Alright so here's my take on this feature. I see three options to support tensordict:
- Make it possible to rebuild tensordict without passing the batch-size as in https://github.com/pytorch/tensordict/pull/674. The caveat is that you will lose all the metadata such as batch-size or dim names. That isn't ideal
- Second solution: rely on pytree as in #2437. That will work amazingly well from the next release of pytorch and on but there are 2 caveats here: (1) the function we use in pytree is quite recent (not even in the stable release) and (2) accelerate requires torch>=1.10, at a time where I think PyTree wasn't a thing.
- Third solution: directly implement this for tensordict. TensorDict has its own `tensordict.to` and excluding keys from the op would be very easy with `tensordict.exclude`. But that would mean taking tensordict as a dependency.
- Last option: I see that `send_to_device` looks for a `to` method within the leaves of the structure it receives and attempts to send the tensor to device using that. We could do that earlier on (ie directly on the dict, such as what is done with tensordict).
Currently we have 
```python
    if isinstance(tensor, (tuple, list)):
        ...
    elif isinstance(tensor, Mapping):
        ...
    elif hasattr(tensor, "to"):
        ...
    else:
        return tensor
```
and just moving things around like this 
```python
    if hasattr(tensor, "to"):
        ...
    elif isinstance(tensor, (tuple, list)):
        ...
    elif isinstance(tensor, Mapping):
        ...
    else:
        return tensor
```
would solve the tensordict issue without any new dependency etc.
I'm not a fan of looking for a `to` attribute in the first place but it seems it's working ok in your case so I think that could be the easiest way forward. 
Wdyt?
Note: only options 2 and 3 work with `skip_keys` and tensordict -- in option 4 you won't be able to exclude a key from the tensordict because for `send_to_device` the item is just a tensor.

Resolving patch:
diff --git a/src/accelerate/utils/operations.py b/src/accelerate/utils/operations.py
--- a/src/accelerate/utils/operations.py
+++ b/src/accelerate/utils/operations.py
@@ -145,22 +145,7 @@ def send_to_device(tensor, device, non_blocking=False, skip_keys=None):
     Returns:
         The same data structure as `tensor` with all tensors sent to the proper device.
     """
-    if isinstance(tensor, (tuple, list)):
-        return honor_type(
-            tensor, (send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys) for t in tensor)
-        )
-    elif isinstance(tensor, Mapping):
-        if isinstance(skip_keys, str):
-            skip_keys = [skip_keys]
-        elif skip_keys is None:
-            skip_keys = []
-        return type(tensor)(
-            {
-                k: t if k in skip_keys else send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys)
-                for k, t in tensor.items()
-            }
-        )
-    elif is_torch_tensor(tensor) or hasattr(tensor, "to"):
+    if is_torch_tensor(tensor) or hasattr(tensor, "to"):
         # `torch.Tensor.to("npu")` could not find context when called for the first time (see this [issue](https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue)).
         if device == "npu":
             device = "npu:0"
@@ -188,6 +173,21 @@ def send_to_device(tensor, device, non_blocking=False, skip_keys=None):
             return tensor.to(device, non_blocking=non_blocking)
         except TypeError:  # .to() doesn't accept non_blocking as kwarg
             return tensor.to(device)
+    elif isinstance(tensor, (tuple, list)):
+        return honor_type(
+            tensor, (send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys) for t in tensor)
+        )
+    elif isinstance(tensor, Mapping):
+        if isinstance(skip_keys, str):
+            skip_keys = [skip_keys]
+        elif skip_keys is None:
+            skip_keys = []
+        return type(tensor)(
+            {
+                k: t if k in skip_keys else send_to_device(t, device, non_blocking=non_blocking, skip_keys=skip_keys)
+                for k, t in tensor.items()
+            }
+        )
     else:
         return tensor
 


Test patch: