Repository: huggingface/accelerate

Issue:
Bugs with IterableDataset
Hi there,

I notice few bugs when using IterableDataset and accelerate DataLoaderDispatcher together.
First, notice that in prepare_data_loader
https://github.com/huggingface/accelerate/blob/31fa5e0ce39b0412eaec537f09bb0e547285adda/src/accelerate/data_loader.py#L401
Some attributes, including drop_last, of the dataloder are set back to default. While this is fine for map style dataset, it breaks the behavior for iterable dataset. The batch_size attribute is restored in the following line, I think drop_last should also be restored here.
https://github.com/huggingface/accelerate/blob/31fa5e0ce39b0412eaec537f09bb0e547285adda/src/accelerate/data_loader.py#L536

Meanwhile, in DataLoaderDispatcher, when drop_last is set to False, the data loader will dispatch whatever is left in the temporary batches list. However, it's possible that the length of batches cannot be evenly dispatch to all processes, which means the dispatch batch_size is 0. In this case, each process will get empty batch.
https://github.com/huggingface/accelerate/blob/31fa5e0ce39b0412eaec537f09bb0e547285adda/src/accelerate/data_loader.py#L368

I have some minimum code to demonstrate this.
```
from accelerate import Accelerator
from accelerate.data_loader import DataLoaderDispatcher
import torch.utils.data
import math


class MyIterableDataset(torch.utils.data.IterableDataset):
    def __init__(self) -> None:
        super().__init__()
        self.data = list(range(100))

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            return iter(self.data)
        else:
            per_worker = int(math.ceil(len(self.data) / float(worker_info.num_workers)))
            start = worker_info.id * per_worker
            end = min(start + per_worker, len(self.data))
            return iter(self.data[start: end])


if __name__ == "__main__":
    acc = Accelerator()

    ds = MyIterableDataset()
    loader = torch.utils.data.DataLoader(ds, batch_size=3, num_workers=2, drop_last=True)
    loader = acc.prepare(loader)

    # note that loader.drop_last is set to False by :meth:`prepare`
    acc.print(loader.split_batches, loader.drop_last, loader.batch_size)

    for i, batch in enumerate(loader):
        print(acc.process_index, i, batch)
        acc.wait_for_everyone()

    acc.print("Start testing loader2 ...")

    loader2 = DataLoaderDispatcher(ds, batch_size=3, num_workers=2, drop_last=True)
    for i, batch in enumerate(loader2):
        print(acc.process_index, i, batch)
        acc.wait_for_everyone()
```

I run the above script with command `accelerate launch --multi_gpu --num_processes=3 --num_machines=1 tests/bin/iterable_dataset_and_accelerate.py`

The result looks like the following:
```
False False 3
0 0 tensor([0, 1, 2], device='cuda:0')
2 0 tensor([3, 4, 5], device='cuda:2')
1 0 tensor([50, 51, 52], device='cuda:1')
0 1 tensor([53, 54, 55], device='cuda:0')
2 1 tensor([56, 57, 58], device='cuda:2')
1 1 tensor([6, 7, 8], device='cuda:1')
0 2 tensor([ 9, 10, 11], device='cuda:0')
2 2 tensor([12, 13, 14], device='cuda:2')
1 2 tensor([59, 60, 61], device='cuda:1')
0 3 tensor([62, 63, 64], device='cuda:0')
2 3 tensor([65, 66, 67], device='cuda:2')
1 3 tensor([15, 16, 17], device='cuda:1')
0 4 tensor([18, 19, 20], device='cuda:0')
2 4 tensor([21, 22, 23], device='cuda:2')
1 4 tensor([68, 69, 70], device='cuda:1')
0 5 tensor([71, 72, 73], device='cuda:0')
2 5 tensor([74, 75, 76], device='cuda:2')
1 5 tensor([24, 25, 26], device='cuda:1')
0 6 tensor([27, 28, 29], device='cuda:0')
2 6 tensor([30, 31, 32], device='cuda:2')
1 6 tensor([77, 78, 79], device='cuda:1')
0 7 tensor([80, 81, 82], device='cuda:0')
2 7 tensor([83, 84, 85], device='cuda:2')
1 7 tensor([33, 34, 35], device='cuda:1')
0 8 tensor([36, 37, 38], device='cuda:0')
2 8 tensor([39, 40, 41], device='cuda:2')
1 8 tensor([86, 87, 88], device='cuda:1')
0 9 tensor([89, 90, 91], device='cuda:0')
2 9 tensor([92, 93, 94], device='cuda:2')
1 9 tensor([42, 43, 44], device='cuda:1')
0 10 tensor([45, 46], device='cuda:0')
2 10 tensor([96, 97], device='cuda:2')
1 10 tensor([47, 95], device='cuda:1')
2 11 tensor([], device='cuda:2', dtype=torch.int64)
0 11 tensor([], device='cuda:0', dtype=torch.int64)
1 11 tensor([], device='cuda:1', dtype=torch.int64)
Start testing loader2 ...
2 0 tensor([3, 4, 5], device='cuda:2')
1 0 tensor([50, 51, 52], device='cuda:1')
0 0 tensor([0, 1, 2], device='cuda:0')
0 1 tensor([53, 54, 55], device='cuda:0')
2 1 tensor([56, 57, 58], device='cuda:2')
1 1 tensor([6, 7, 8], device='cuda:1')
0 2 tensor([ 9, 10, 11], device='cuda:0')
2 2 tensor([12, 13, 14], device='cuda:2')
1 2 tensor([59, 60, 61], device='cuda:1')
0 3 tensor([62, 63, 64], device='cuda:0')
2 3 tensor([65, 66, 67], device='cuda:2')
1 3 tensor([15, 16, 17], device='cuda:1')
0 4 tensor([18, 19, 20], device='cuda:0')
2 4 tensor([21, 22, 23], device='cuda:2')
1 4 tensor([68, 69, 70], device='cuda:1')
0 5 tensor([71, 72, 73], device='cuda:0')
2 5 tensor([74, 75, 76], device='cuda:2')
1 5 tensor([24, 25, 26], device='cuda:1')
0 6 tensor([27, 28, 29], device='cuda:0')
2 6 tensor([30, 31, 32], device='cuda:2')
1 6 tensor([77, 78, 79], device='cuda:1')
0 7 tensor([80, 81, 82], device='cuda:0')
2 7 tensor([83, 84, 85], device='cuda:2')
1 7 tensor([33, 34, 35], device='cuda:1')
0 8 tensor([36, 37, 38], device='cuda:0')
2 8 tensor([39, 40, 41], device='cuda:2')
1 8 tensor([86, 87, 88], device='cuda:1')
0 9 tensor([89, 90, 91], device='cuda:0')
2 9 tensor([92, 93, 94], device='cuda:2')
1 9 tensor([42, 43, 44], device='cuda:1')
```

**Hope some one can fix this soon.**


Discussion:
Thanks for reporting and providing a reproducer. I hope to have some time to dig into this and fix early next week :-) 

Resolving patch:
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
--- a/src/accelerate/data_loader.py
+++ b/src/accelerate/data_loader.py
@@ -335,6 +335,7 @@ def __iter__(self):
             # We only iterate through the DataLoader on process 0.
             main_iterator = super().__iter__()
         stop_iteration = False
+        first_batch = None
         while not stop_iteration:
             # On process 0, we gather the batch to dispatch.
             if state.process_index == 0:
@@ -382,7 +383,18 @@ def __iter__(self):
             # Broadcast the batch before splitting it.
             batch = broadcast(batch, from_process=0)
 
-            batch_size = find_batch_size(batch) // state.num_processes
+            if not self.drop_last and first_batch is None:
+                # We keep at least num processes elements of the first batch to be able to complete the last batch
+                first_batch = slice_tensors(batch, slice(0, state.num_processes))
+
+            observed_batch_size = find_batch_size(batch)
+            batch_size = observed_batch_size // state.num_processes
+
+            if not self.drop_last and stop_iteration and observed_batch_size % state.num_processes != 0:
+                # If the last batch is not complete, let's add the first batch to it.
+                batch = concatenate([batch, first_batch], dim=0)
+                batch_size += 1
+
             data_slice = slice(state.process_index * batch_size, (state.process_index + 1) * batch_size)
 
             if state.distributed_type == DistributedType.TPU:
@@ -534,6 +546,7 @@ def prepare_data_loader(
 
     # Need to provide batch_size as batch_sampler is None for Iterable dataset
     if new_batch_sampler is None:
+        kwargs["drop_last"] = dataloader.drop_last
         kwargs["batch_size"] = dataloader.batch_size // num_processes if split_batches else dataloader.batch_size
 
     if dispatch_batches:


Test patch: