Repository: huggingface/accelerate
Pull request: add support of gather_object
fixes #235

Problem statement:
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.

Hints from issue 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.

Commit message: add support of gather_object

fixes #235

Patch to review:
diff --git a/src/accelerate/utils.py b/src/accelerate/utils.py
index 99199728248..8a027312b2a 100644
--- 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
@@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather_object(output_objects, object)
+        return [i for j in output_objects for i in j]
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
+
+
+_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 :obj:`Any`):
+            The data to gather.
+
+    Returns:
+        The same data structure as :obj:`Any` with all objects sent to the proper 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)


Code review comments:
diff_hunk: @@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather_object(output_objects, object)
+        return [i for j in output_objects for i in j]
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
+
+
+_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 :obj:`Any`):
+            The data to gather.
+
+    Returns:
+        The same data structure as :obj:`Any` with all objects sent to the proper device.
path: src/accelerate/utils.py
text: @user1:
```suggestion
        The same data structure as :obj:`object` with all the objects sent to every device.
```
I think we need to explain where the list arrives (since each tensor is replaced by a list of tensors normally).

diff_hunk: @@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
path: src/accelerate/utils.py
text: @user1:
Here we should use the `state.num_processes` instead.

diff_hunk: @@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather_object(output_objects, object)
+        return [i for j in output_objects for i in j]
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
path: src/accelerate/utils.py
text: @user1:
`torch.all_gather_object` works for all picklable objects, so we can skip the `error_on_other_type=True` I believe.

diff_hunk: @@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather_object(output_objects, object)
+        return [i for j in output_objects for i in j]
path: src/accelerate/utils.py
text: @user1:
I don't understand this line. This presupposes that every element in `output_objects` is iterable, but `object` is supposed to be a primitive object at this point (tensor or other if we remove `error_on_type=True`).

diff_hunk: @@ -345,6 +345,40 @@ def gather(tensor):
         return tensor
 
 
+
+def _gpu_gather_object(object: Any):
+    def _gpu_gather_object_one(object: Any):
+        output_objects = [None for _ in range(torch.distributed.get_world_size())]
+        torch.distributed.all_gather_object(output_objects, object)
+        return [i for j in output_objects for i in j]
+
+    return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
+
+
+_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 :obj:`Any`):
path: src/accelerate/utils.py
text: @user1:
```suggestion
        object (nested list/tuple/dictionary of any object):
```