Repository: hpcaitech/ColossalAI

Issue:
[BUG]: GradAccumDataloader will ignore the last iter.
### 🐛 Describe the bug

I used the list as dataloader for testing.

```
class GradAccumDataloader:
    """A wrapper for dataloader to enable gradient accumulation by dropping the last incomplete steps.

    Note:
        The dataloader would drop the last incomplete steps for gradient accumulation.
        For example, if a dataloader has 10 batches of data and accumulate size is 4. The model parameters will
        be updated only twice at step 4 and step 8. The last two batches of data do not form a complete 4-step cycle.
        Thus, they will be automatically skipped by this class. If the dataloader is not standard PyTorch dataloader,
        (e.g. Dali dataloader), this class will automatically consume (load data for nothing) the remaining 2 batches.

    Args:
        dataloader (``Iterable``): Your dataloader object for gradient accumulation.
        accumulate_size (int): The number of steps to accumulate gradients.
    """

    def __init__(self, dataloader: Iterable, accumulate_size: int) -> None:
        self.dataloader = dataloader
        self.consume_remain_data = True
        self.steps_per_epoch = len(dataloader) - len(dataloader) % accumulate_size

    def __getattr__(self, __name: str) -> Any:
        return getattr(self.dataloader, __name)

    def __len__(self) -> int:
        return self.steps_per_epoch

    def __iter__(self) -> Iterable:
        self._cur_step = 0
        self._dataiter = iter(self.dataloader)
        return self

    def __next__(self) -> Union[Tensor, Tuple[Tensor]]:
        if self._cur_step < self.steps_per_epoch:
            self._cur_step += 1

            if self._cur_step == self.steps_per_epoch and self.consume_remain_data:
                # this is to handle non standard pytorch dataloader
                # such as dali dataloader
                while True:
                    try:
                        _ = next(self._dataiter)
                    except StopIteration:
                        break
            return next(self._dataiter)
        else:
            raise StopIteration

a = list(range(1,11))
print(len(a), a)
a_ = GradAccumDataloaderTorch(a, accumulate_size=4)
for i, val in enumerate(a_):
    print(i, val)
```

the output is:
```
10 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
0 1
1 2
2 3
3 4
4 5
5 6
6 7
```



### Environment

collosal : pip install colossalai==0.1.4+torch1.10cu11.1 -f https://release.colossalai.org
ubuntu 20.04


Discussion:

Resolving patch:
diff --git a/colossalai/engine/gradient_accumulation/_gradient_accumulation.py b/colossalai/engine/gradient_accumulation/_gradient_accumulation.py
--- a/colossalai/engine/gradient_accumulation/_gradient_accumulation.py
+++ b/colossalai/engine/gradient_accumulation/_gradient_accumulation.py
@@ -145,6 +145,7 @@ def __iter__(self) -> Iterable:
     def __next__(self) -> Union[Tensor, Tuple[Tensor]]:
         if self._cur_step < self.steps_per_epoch:
             self._cur_step += 1
+            data = next(self._dataiter)
 
             if self._cur_step == self.steps_per_epoch and self.consume_remain_data:
                 # this is to handle non standard pytorch dataloader
@@ -154,7 +155,7 @@ def __next__(self) -> Union[Tensor, Tuple[Tensor]]:
                         _ = next(self._dataiter)
                     except StopIteration:
                         break
-            return next(self._dataiter)
+            return data
         else:
             raise StopIteration
 


Test patch: