Repository: huggingface/accelerate

Issue:
stronger gather implement
`gather` is crucial in evaluating model. current implementation assumes each tensor across devices has same size, which is very unlikely to happen in the real case. calling `pad_across_processes` results to additional 0s in the output, leading to incorrect results. 

```python
def _gpu_gather_one(tensor, same_size = True):
    if tensor.ndim == 0:
        tensor = tensor.clone()[None]

    if not same_size:
        size = torch.tensor(tensor.shape[0])[None]
        size_list = [torch.zeros(1) for _ in range(torch.distributed.get_world_size())]
        torch.distributed.all_gather(size_list, size)

    output_tensors = [torch.zeros_like(tensor) for _ in range(torch.distributed.get_world_size())]
    torch.distributed.all_gather(output_tensors, tensor)

    if not same_ssize:
        output_tensors = [i[:j] for i, j in zip(output_tensors, size_list)]
    return torch.cat(output_tensors, dim=0)
```

one another workaround is to convert results into `list` and use `gather_object` to gather. However, `gather_object` is also not provided in accelerate.


Discussion:
Hi there, thanks for the feedback! We'll look into adding support for `gather_object` as it's impossible to automatically return a tensor from tensors on different devices that are not of the same sizes.

But out of curiosity, in what type of problems do you have tensors of different sizes on different GPUs without having padding already in those tensors? Do they have a dimension of 1? Usually those arises in problems with different shapes where padding is already implemented.
> Hi there, thanks for the feedback! We'll look into adding support for `gather_object` as it's impossible to automatically return a tensor from tensors on different devices that are not of the same sizes.

I'm more than happy to submit a patch for this~

> But out of curiosity, in what type of problems do you have tensors of different sizes on different GPUs without having padding already in those tensors? Do they have a dimension of 1? Usually those arises in problems with different shapes where padding is already implemented.

In my case, I only have inference and groundtruth synced across devices to calculate score such as accuracy, recall, Pearson correlation, etc. and they are all 1-d tensor. 

Resolving patch:
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
--- a/src/accelerate/utils.py
+++ b/src/accelerate/utils.py
@@ -18,7 +18,7 @@
 from collections.abc import Mapping
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Union
 
 import numpy as np
 import torch
@@ -350,6 +350,39 @@ def gather(tensor):
         return tensor
 
 
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(AcceleratorState().num_processes)]
+        torch.distributed.all_gather_object(output_objects, object)
+        return output_objects
+
+    return recursively_apply(_gpu_gather_object_one, object)
+
+
+_cpu_gather_object = _gpu_gather_object
+
+
+def gather_object(object: Any):
+    """
+    Recursively gather object in a nested list/tuple/dictionary of objects from all devices.
+
+    Args:
+        object (nested list/tuple/dictionary of picklable object):
+            The data to gather.
+
+    Returns:
+        The same data structure as :obj:`object` with all the objects sent to every device.
+    """
+    if AcceleratorState().distributed_type == DistributedType.TPU:
+        raise NotImplementedError("gather objects in TPU is not supported")
+    elif AcceleratorState().distributed_type == DistributedType.MULTI_GPU:
+        return _gpu_gather_object(object)
+    elif AcceleratorState().distributed_type == DistributedType.MULTI_CPU:
+        return _cpu_gather_object(object)
+    else:
+        return object
+
+
 def _gpu_broadcast(data, src=0):
     def _gpu_broadcast_one(tensor, src=0):
         torch.distributed.broadcast(tensor, src=src)


Test patch: