Accelerator incorrectly gather batches if features do not have batch dim
Author: qubvel
Pull request: no
Labels: name: trainer

name: bug

name: Accelerate

### System Info

- `transformers` version: 4.40.0.dev0
- Platform: Linux-5.15.0-1056-aws-x86_64-with-glibc2.31
- Python version: 3.10.9
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.2
- Accelerate version: 0.29.3
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: NO
        - mixed_precision: bf16
        - use_cpu: False
        - debug: False
        - num_processes: 1
        - machine_rank: 0
        - num_machines: 1
        - gpu_ids: all
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: False
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 1.13.0+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@muellerzr 

### Information

- [X] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

In certain tasks, like object detection (and potentially others) some features within a batch are not concatenated along the batch dimension but are instead represented as a list of dictionaries. This approach is used because each image might contain a different number of objects, hence the feature `boxes` in object detection is not combined across the batch.

During the evaluation phase, there might be instances where the last batch is smaller than the others. For example, having only 5 or even just 1 sample, whereas other batches contain 8 samples. When the `accelerator.gather_for_metrics` function is applied to the last batch, features that do not have a batch dimension are incorrectly sliced along the `0` dimension according to the number of samples in that batch. For example, if the last batch contains only one sample, then all applicable tensors are sliced as `tensor = tensor[:1]`. This slicing can lead to inaccuracies because it assumes the presence of a batch dimension where there is none.

https://github.com/huggingface/accelerate/blob/abc86c0e35171558b071bbbf4bd377d9675d2bfe/src/accelerate/accelerator.py#L2290

This happens in the `Trainer` evaluation loop and could be reproduced with clear `accelerate` script:

```python
import torch
import random
from torch.utils.data import DataLoader, Dataset
from accelerate import Accelerator


class Dataset(Dataset):
    def __init__(self, size):
        self.size = size

    def __len__(self):
        return self.size

    def __getitem__(self, idx):
        return {"boxes": torch.rand(random.randint(7, 15), 4)}


dataloader = DataLoader(Dataset(5), batch_size=8, collate_fn=lambda x: x)
accelerator = Accelerator()
dataloader = accelerator.prepare(dataloader)

for i, batch in enumerate(dataloader):
    
    #skip until the last batch
    if i != len(dataloader) - 1:
        continue
    
    print("Batch before gather:", [x["boxes"].shape for x in batch])
    gathered_batch = accelerator.gather_for_metrics(batch)
    print("Batch after gather: ", [x["boxes"].shape for x in gathered_batch])

# Batch before gather: [torch.Size([11, 4]), torch.Size([12, 4]), torch.Size([15, 4]), torch.Size([8, 4]), torch.Size([9, 4])]
# Batch after gather:  [torch.Size([5, 4]), torch.Size([5, 4]), torch.Size([5, 4]), torch.Size([5, 4]), torch.Size([5, 4])]
```

### Expected behavior

Expected to have the correct boxes after `gather_for_metrics`.

Probably, before slicing we have to check that 
1) 0 dim is equal to batch size
2) and the rest part is duplicate or padded

Comments:
Hi @qubvel, thanks for the detailed report. I'm working on a solution for that in the above PR !